#include "nu_server.h"
#include "nu_lanman.h"

static void bailout( char *reason );
static void sigint( int reason ) { bailout( "User interrupt." ); }
static mainsock = -1;  
static childsock = -1; 

SMB_COMMAND g_transact;  /* If set to != 0x00, we are waiting for some data bytes
                            to complete a transaction request. This variable tells
                            us, which request we're waiting for. */

static QWORD getnttime( struct timeval *t ) {
  return 10000000ll * ( t->tv_sec + 11644473600ll ) + t->tv_usec * 10ll;
}

static void netbios_read( SMB_HEADER **buf) {
  DWORD   bytes;
  ssize_t bytestoread;

  if( read( childsock, &bytes, 4) < 4 )
    bailout( "Short read." );
  bytestoread = htons(bytes>>16);
  if( (*buf = (SMB_HEADER*)realloc( *buf, 4 + bytestoread )) == NULL)
    bailout( "Out of memory.");
  *(DWORD*)*buf = bytes;
  if( read( childsock, ((BYTE*)*buf) + 4, bytestoread) != bytestoread )
    bailout( "Short read." );
}

static void netbios_write( SMB_HEADER *buf,
                           SMB_DATA   *data,
                           int         datacount ) {
  struct iovec iov[1 + 2 * SMB_MAXREQUESTS] = { {buf , sizeof(SMB_HEADER) } };
  int    i, sizeout = sizeof( SMB_HEADER );

  for( i=0; i < datacount; ++i )
               iov[1+2*i].iov_base =     data[i].params,
    sizeout += iov[1+2*i].iov_len  = 1 + data[i].params->WordCount * 2,
               iov[2+2*i].iov_base =     data[i].bytes,
    sizeout += iov[2+2*i].iov_len  = 2 + data[i].bytes->ByteCount;

  buf->netbios_command = buf->netbios_flags = 0;
  buf->netbios_size    = htons( sizeout - 4 );
  buf->Flags           = 0x88;
  buf->Flags2          = 0x4001;

  if( writev( childsock, iov, 1 + 2 * datacount ) < sizeout )
    bailout( "Write failed." );
}

static const BYTE SMB_COM_NEGOTIATE_bytes[]  = { 8,0,0x67,0x61,0x74,0x6c,0x69,0x6e,0x67,0 };
static       WORD SMB_COM_NEGOTIATE_params[] = {
0x0511, 0x0000, 0x0001, 0x0001, 0x0000, 0x0100, 0x0000, 0x0100, 0x0000,
0x0000, 0xC049, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000 };

static SMB_STATUS handle_SMB_COM_NEGOTIATE( SMB_HEADER **header, SMB_DATA *data ) {
  struct timeval t; gettimeofday( &t, NULL );
  int i = 3;
  /* Assign uniqe session id, don't know whether spreading our
     pid() is a good id, when I know, I might change that :) */
  SMB_COM_NEGOTIATE_params[8] = getpid();
  SMB_COM_NEGOTIATE_params[9] = getppid();
  *(QWORD*)(SMB_COM_NEGOTIATE_params+12) = getnttime( &t );

  /* count number of dialect strings and choose the
     last one, foolishly assuming this is the newest */
  ((BYTE*)SMB_COM_NEGOTIATE_params)[1] = 0;
  while( i < data->bytes->ByteCount )
    if( data->bytes->Buffer[i++] == 0x02 ) /* dialect indicator */
      ++*(1+(BYTE*)SMB_COM_NEGOTIATE_params);

  data->params = (SMB_PARAMS*)SMB_COM_NEGOTIATE_params;
  data->bytes  = (SMB_BYTES *)SMB_COM_NEGOTIATE_bytes;
  return STATUS_SUCCESS;
}

static const BYTE SMB_COM_SESSION_SETUP_ANDX_bytes[] = {
19,0,'O','S',0,'g','a','t','l','i','n','g',0,'g','a','t','l','i','n','g',0};
static       BYTE SMB_COM_SESSION_SETUP_ANDX_params[] = { 4, 255, 0, 0, 0, 1, 0, 0, 0 };

static SMB_STATUS handle_SMB_COM_SESSION_SETUP_ANDX( SMB_HEADER **header, SMB_DATA *data ) {
  data->params = (SMB_PARAMS*)SMB_COM_SESSION_SETUP_ANDX_params;
  data->bytes  = (SMB_BYTES *)SMB_COM_SESSION_SETUP_ANDX_bytes;
  return STATUS_SUCCESS;
}

static const BYTE SMB_COM_TREE_CONNECT_ANDX_bytes[] = { 9, 0, 'A', ':',  0, 'F', 'A', 'T', '3', '2', 0 };
static       BYTE SMB_COM_TREE_CONNECT_ANDX_params[] = { 3, 255, 0, 0, 0, 0, 0 };

static SMB_STATUS handle_SMB_COM_TREE_CONNECT_ANDX( SMB_HEADER **header, SMB_DATA *data ){
  (*header)->TreeID = 5;
  data->params = (SMB_PARAMS*)SMB_COM_TREE_CONNECT_ANDX_params;
  data->bytes  = (SMB_BYTES *)SMB_COM_TREE_CONNECT_ANDX_bytes;
  return STATUS_SUCCESS;
}

/*static SMB_STATUS handle_SMB_COM_TRANSACTION( SMB_HEADER **header, SMB_DATA *data ) {
  if( !strcmp( (char*)&data->bytes[1], "\\PIPE\\LANMAN"))
  {
    /* TODO: Sanity Check on DataCount vs. ByteCount * /
    SMB_PARAMS_TRANSACTION *params = (SMB_PARAMS_TRANSACTION *)data->params;
    SMB_TRANSACTION_BYTES   bytes;

    bytes.params = ((BYTE*)&((*header)->Protocol)) + GETNWORD( params->ParameterOffset );
    bytes.paramc =                                   GETNWORD( params->ParameterCount );
    bytes.data   = ((BYTE*)&((*header)->Protocol)) + GETNWORD( params->DataOffset );
    bytes.datac  =                                   GETNWORD( params->DataCount );

    return handle_LANMAN( *header, data, &bytes );
  }
  else
    return 0x00400002;
} */

static SMB_STATUS handle_SMB_COM_TRANSACTION2( SMB_HEADER **header, SMB_DATA *data ) {
  SMB_PARAMS_TRANSACTION2 *params = (SMB_PARAMS_TRANSACTION2 *)data->params;
  SMB_TRANSACTION_BYTES   bytes;
  int i;

  printf( "%d %d %d %d\n", GETNWORD(params->TotalParameterCount), GETNWORD(params->ParameterCount),
                           GETNWORD(params->TotalDataCount     ), GETNWORD(params->DataCount     ));

//  for( i=0; i<sizeof( SMB_PARAMS_TRANSACTION2 ); ++i )
//    printf( "%02x ", ((BYTE*)data->params)[i] );

  if( ( GETNWORD(params->TotalParameterCount) != GETNWORD(params->ParameterCount)) ||
      ( GETNWORD(params->TotalDataCount     ) != GETNWORD(params->DataCount     )))
  {
    /* Handle and or reassemble split packets later */
    bailout( "One missing feature detected.");
  }

  bytes.params = ((BYTE*)&((*header)->Protocol)) + GETNWORD( params->ParameterOffset );
  bytes.paramc =                                   GETNWORD( params->ParameterCount );
  bytes.data   = ((BYTE*)&((*header)->Protocol)) + GETNWORD( params->DataOffset );
  bytes.datac  =                                   GETNWORD( params->DataCount );

  switch( GETNWORD( &params->Setup )) { /* Transaction2 Command Code */
    case SMB_TRANS2_FIND_FIRST2:
      printf( "Attributes:\t%d\nCount:\t\t%d\nFlags:\t\t%02x\nLevel:\t\t%d\nType:\t\t%d\nPattern:\t%s\n",
              GETNWORD(     bytes.params ),
              GETNWORD( 2 + bytes.params ),
              GETNWORD( 4 + bytes.params ),
              GETNWORD( 6 + bytes.params ),
              GETNWORD( 8 + bytes.params ),
                            bytes.params + 12 );
      break;
    default:
      return 0x00400002; /* No sub command handler yet */
      break;
  }

  return STATUS_SUCCESS;
}

static int command_handler_match(const void *a, const void *b ) { return *(BYTE*)a - *(BYTE*)b; }

/* If you add command handlers, please insert them in the right position,
   this list is sorted by command, for later bsearch*/
static SMB_COMMAND_HANDLER command_handler[] = {
/*{ SMB_COM_TRANSACTION,                    0x00, handle_SMB_COM_TRANSACTION }, */
  { SMB_COM_TRANSACTION2,                   0x00, handle_SMB_COM_TRANSACTION2 },
  { SMB_COM_NEGOTIATE,                      0x00, handle_SMB_COM_NEGOTIATE },
  { SMB_COM_SESSION_SETUP_ANDX,             0x01, handle_SMB_COM_SESSION_SETUP_ANDX },
  { SMB_COM_TREE_CONNECT_ANDX,              0x01, handle_SMB_COM_TREE_CONNECT_ANDX }
};

static void child( ) {
  SMB_HEADER *inpacket    = NULL;
  DWORD       netbios_ack = 0x00000082;

  /* I should spare that code... */
  if( mainsock != -1 ) { close( mainsock ); mainsock = -1; }

  /* Try to answer first netbios packet */
  netbios_read( &inpacket );
  if( inpacket->netbios_command != 0x81 )
    bailout( "No session request");
  write( childsock, &netbios_ack, 4);

  g_transact = 0x00; /* No transaction waiting for bytes */

  while( 1 ) {
    SMB_COMMAND cmd;
    SMB_DATA    requests[ 1 + SMB_MAXREQUESTS ];
    SMB_STATUS  status = STATUS_SUCCESS;
    DWORD       null = 0;
    WORD        sizeout = sizeof( SMB_HEADER ) - 4;
    int         num_requests = 0;

    netbios_read( &inpacket );

    if( inpacket->netbios_command != 0 )
      bailout( "Unhandled netbios command" );
    if( inpacket->Protocol != SMB_HEADER_PROTOCOL_MAGIC )
      bailout( "Protocol identifier mismatch");

    cmd = inpacket->Command;
    requests[ 0 ].params = (SMB_PARAMS*)(inpacket+1);

    if( g_transact && (g_transact != cmd) )
      bailout( "Transaction interrupted by other SMB command.");

    while( (status == STATUS_SUCCESS) && (cmd != 0xff) ) {
      SMB_COMMAND_HANDLER *handler = (SMB_COMMAND_HANDLER*)bsearch( &cmd, command_handler,
        sizeof(command_handler)/sizeof(*command_handler), sizeof(*command_handler), command_handler_match);

      requests[ num_requests ].bytes  =
        (SMB_BYTES*)(((BYTE*)requests[ num_requests ].params) + 2 * *((BYTE*)(requests[ num_requests ].params)) + 1);

      if( handler ) {
        if( handler->flags & SMB_COMMAND_FLAG_ANDX ) {
          cmd = ((BYTE*)requests[ num_requests ].params)[1];
          requests[ num_requests+1 ].params = (SMB_PARAMS*)(((BYTE*)inpacket) + 4 +
            ((BYTE*)requests[ num_requests ].params)[3] + 256*((BYTE*)requests[ num_requests ].params)[4]);
        }

        /* <---------- Calling handler here -----> */
        if( (status = handler->handler( &inpacket, requests + num_requests )) != STATUS_SUCCESS )
          cmd = 0xff;
        if( handler->flags & SMB_COMMAND_FLAG_ANDX ) {
          sizeout += 3 + 2 * requests[ num_requests ].params->WordCount + requests[ num_requests ].bytes->ByteCount;
          ((BYTE*)requests[ num_requests ].params)[1] = cmd,
          ((BYTE*)requests[ num_requests ].params)[3] = sizeout & 255,
          ((BYTE*)requests[ num_requests ].params)[4] = sizeout >> 8;
        } else
          cmd = 0xff;
      } else { /* no handler - return STATUS_UKCOMMAND*/
        if( num_requests )
          ((BYTE*)requests[ num_requests-1 ].params)[1] = 0xff;
        requests[ num_requests ].params = (SMB_PARAMS*)&null;
        requests[ num_requests ].bytes  = (SMB_BYTES*) &null;
        status = 0x00400002;
      }
      num_requests++;
    }

    memcpy( inpacket->Status, &status, 4 ); /* not aligned, maybe we might do a store DWORD on x86 */
    netbios_write( inpacket, requests, num_requests );
  } /* End main loop */
}

int main() {
  struct sockaddr_in sa;
  int    l=1;

  signal( SIGINT, sigint);

  bzero( &sa, sizeof( sa));
  sa.sin_family      = PF_INET;
  sa.sin_port        = htons( 139 );
  sa.sin_addr.s_addr = INADDR_ANY;

  if( ( mainsock = socket( PF_INET, SOCK_STREAM, 0) ) == -1)
    bailout( "Could not open socket");
#ifdef SO_REUSEPORT
  setsockopt( mainsock, SOL_SOCKET, SO_REUSEPORT, &l, sizeof(l));
#else
  setsockopt( mainsock, SOL_SOCKET, SO_REUSEADDR, &l, sizeof(l));
#endif
  if( bind( mainsock, (struct sockaddr *)&sa, sizeof( sa)) != 0)
    bailout( "Could not bind socket");
  if( listen( mainsock, 1024) != 0 )
    bailout( "Could not make socket listen");
 
  while( 1 ) {
    struct sockaddr otherend;
    int             size = sizeof( otherend );
     
    if( ( childsock = accept( mainsock, &otherend, &size) ) == -1)
      bailout( "Socket Broke.");
    if (!fork()) child( );
  }
}

/* Graceful exit. */
static void bailout( char *reason) {
    fputs( reason, stderr);
    fputs( "\nCleaning up.\n", stderr);
    if( mainsock != -1 )
        close( mainsock );
    if( childsock != -1 ) {
        shutdown( childsock, SHUT_RDWR);
        close( childsock );
    }
    exit( 0 );
}