#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 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) {
  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_COMMAND cmd,
                           SMB_HEADER *buf,
                           SMB_PARAMS *buf2,
                           SMB_BYTES  *buf3 ) {
  BYTE buf_[4] = { 0, 0, 0, 0 };
  if(!buf2 ) buf2 = (SMB_PARAMS*)buf_;
  if(!buf3 ) buf3 = (SMB_BYTES*)buf_;

  struct iovec iov[16] = { {buf , sizeof(SMB_HEADER)     },
                           {buf2, 1 + buf2->WordCount * 2},
                           {buf3, 2 + buf3->ByteCount    } };

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

  if( writev( childsock, iov, 3 ) < htons( buf->netbios_size ) + 4 )
    bailout( "Write failed." );
}

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

  /* 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, inpacket, 4);

  while( 1 ) {
    netbios_read( &inpacket );

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

    switch( inpacket->Command ) {
      case SMB_COM_NEGOTIATE:
      {
        const BYTE bytes[] = { 8,0,0x67,0x61,0x74,0x6c,0x69,0x6e,0x67,0 };
        WORD params[] = { 0x0511, 0x0000, 0x0001, 0x0001, 0x0000, 0x0100,
                          0x0000, 0x0100, 0x0000, 0x0000, 0xC049, 0x0000,
                          0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000 };
        struct timeval t; gettimeofday( &t, NULL );

        *(DWORD*)&inpacket->Status   = STATUS_SUCCESS;
        params[8] = getpid(); params[9] = getppid();
        *(QWORD*)&params[12] = getnttime( &t );
        netbios_write( 0, inpacket, (SMB_PARAMS*)params, (SMB_BYTES*)bytes);
        break;
      }
/*      case SMB_COM_SESSION_SETUP_ANDX:
      {
        const BYTE bytes[] = {  19,   0, 'O', 'S',   0, 'g', 'a', 't', 'l',
                               'i', 'n', 'g',   0, 'g', 'a', 't', 'l', 'i',
                               'n', 'g',   0};
        BYTE params[] = { 4, 0, 0, 23, 0, 1, 0, 0, 0 };

      } */
      default:
      {
        fprintf( stderr, "Got message: %02X\n", inpacket->Command );
        *(DWORD*)&inpacket->Status          = 0x00400002;
        netbios_write( 0, inpacket, NULL, 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 );
}