#include <time.h>
#include <fcntl.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <stdio.h>
#include <sys/ioctl.h>

#include "nu_header.h"

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

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

static void netbios_read( SMB_HEADER **buf) {
  BYTE    bytes[4];
  ssize_t bytestoread;

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

static void netbios_write( BYTE cmd, SMB_HEADER *buf, SMB_HEADER2 *buf2 ) {
  const BYTE buf2_[2] = { 0, 0 };
  if(!buf2 ) buf2 = (SMB_HEADER2*)buf2_;
  if( buf ) {
    struct iovec iov[2] = { {buf, SIZEOF_SMB_HEADER + 2*buf->WordCount},
                            {buf2,                2 + buf2->ByteCount} };

    buf->netbios_command = cmd;
    buf->netbios_flags   = 0;
    buf->netbios_size    = htons( SIZEOF_SMB_HEADER - 4 +
                                  2 * buf->WordCount +
                                  2 + buf2->ByteCount );
    buf->Flags           = 0x88;
    buf->Flags2          = 0x4001;

    if( writev( childsock, iov, 2 ) < htons( buf->netbios_size ) + 4 )
      bailout( "Write failed." );
  } else {
    const BYTE buf_[4] = { cmd, 0, 0, 0 };
    if( write( childsock, buf_, 4 ) < 4)
      bailout( "Write failed." );
  }
}

static void packet_dump( SMB_HEADER *buf ) {
  fprintf( stderr, "netbios_cmd, flag, size = %02X, %02X, %04X\n", buf->netbios_command, buf->netbios_flags, buf->netbios_size );
  fprintf( stderr, "Protocol                = %08X\n", *(DWORD*)&buf->Protocol);
  fprintf( stderr, "Command                 = %02X\n", buf->Command);
  fprintf( stderr, "Status                  = %08X\n", *(DWORD*)&buf->Status);
  fprintf( stderr, "Flags, Flags2           = %02X, %04X\n", buf->Flags, buf->Flags2);
  fprintf( stderr, "Pad = %04X %04X %04X %04X %04X %04X\n",
  buf->Pad[0], buf->Pad[1], buf->Pad[2],
  buf->Pad[3], buf->Pad[4], buf->Pad[5] );
  fprintf( stderr, "TreeID,ProcessID,UserID = %04X, %04X, %04X\n", buf->TreeID, buf->ProcessID, buf->UserID);
  fprintf( stderr, "MultiplexID, WordCount  = %04X, %02X\n", buf->MultiplexID, buf->WordCount);
}

static void child( ) {
  SMB_HEADER *inpacket = NULL;
  DWORD       bytesread;

  /* 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");
  netbios_write( 0x82, NULL, NULL );

  while( 1 ) {
    WORD *ParameterWords;
    netbios_read( &inpacket );
    packet_dump( inpacket );
    ParameterWords = (WORD*)(((BYTE*)inpacket)+SIZEOF_SMB_HEADER);

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

    switch( inpacket->Command ) {
      case SMB_COM_NEGOTIATE:
      {
        BYTE myself[] = { 8,0,0x67,0x61,0x74,0x6c,0x69,0x6e,0x67,0x00 };
        struct timeval t; gettimeofday( &t, NULL );

        inpacket = (SMB_HEADER*)realloc( inpacket, SIZEOF_SMB_HEADER + 17 * 2 );
        *(DWORD*)&inpacket->Status   = STATUS_SUCCESS;

          ParameterWords     = (WORD*)(((BYTE*)inpacket)+SIZEOF_SMB_HEADER-1);
          ParameterWords[0]  = 0x0511;    /* Protocol Version 5, 17 bytes */
          ParameterWords[1]  = 0;         /* security mode: share, no c/r */
          ParameterWords[2]  = 1;         /* Max pending        */
          ParameterWords[3]  = 1;         /* Only one VC        */
          ParameterWords[4]  = 0;         /* Max Buffer Size    */
          ParameterWords[5]  = 0x100;     /* Max Buffer Size #2 */
          ParameterWords[6]  = 0;         /* Max Raw Size       */
          ParameterWords[7]  = 0x100;     /* Max Raw Size #2    */
          ParameterWords[8]  = getpid();  /* unique id          */
          ParameterWords[9]  = getppid(); /* unique id #2       */
          ParameterWords[10] = 0;         /* Capabilities       */
          ParameterWords[11] = 0;         /* Capabilities #2    */
*(QWORD*)&ParameterWords[12] = getnttime( &t );
          ParameterWords[16] = 0;
*(BYTE *)&ParameterWords[17] = 0;
        netbios_write( 0, inpacket, (SMB_HEADER2*)myself);
        break;
      }
      default:
      {
        fprintf( stderr, "Got message: %02X\n", inpacket->Command );
          inpacket->WordCount       = 0;
*(DWORD*)&inpacket->Status          = 0x00400002;
        netbios_write( 0, inpacket, NULL );
        break;
      }
    }

  } /* 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 );
}