#include "drwebdc-demo.h"

/*--------------------------------------------------------------------------
 Dr. Web daemon client demo
                                       ID Antivirus Lab http://www.drweb.ru
                                             http://www.drweb.ru/unix/forum
                                                   http://www.drweb.ru/unix
--------------------------------------------------------------------------*/
static char * pFileBuffer = NULL;
static int  inFile = -1;

static char szFile[ MAXPATHLEN + 2 ] = { 0 };
static int  bShowID = 0;
static int  bShowVersion = 0;
static int  bShowBases = 0;
static int  bScanFile = 0;
static int  bExpand = 0;
static int  bExpanded = 0;
static int  bRealFile = 0;

//--------------------------------------------------------------------------
static void
cleanupClient( void )
{
   cleanup();

   if( inFile != -1 )
   {
       close( inFile );
       inFile = -1;
   }
   if( pFileBuffer )
   {
       free( pFileBuffer );
       pFileBuffer = NULL;
   }
}

//--------------------------------------------------------------------------
static int
usage( void )
{
   printf( "\nDr. Web daemon client demo"
           "\n"
           "\nUsage: drwebdc [keys]"
           "\n"
           "\nValid keys:"
           "\n -f[FILE_NAME]    file to scan (STDIN is default)"
           "\n -sfSOCKET_FILE   Dr. Web daemon UNIX-socket file"
           "\n -pPORT           Dr. Web daemon TCP port"
           "\n -iIP_ADDRESS     Dr. Web daemon IP address"
           "\n -tTIMEOUT        Dr. Web daemon timeout"
           "\n -ha              use heuristic analyser"
           "\n -hf              use headers filter"
           "\n -ex              always expand file name"
           "\n -sb              show virus bases info"
           "\n -si              show daemon id string"
           "\n -sv              show daemon version"
           "\n -rv              return virus names"
           "\n -rc              return extended codes"
           "\n -rr              return report lines"
           "\n -q               quiet mode"
           "\n" );
   return DRWEBD_DEMO_USAGE;
}

//--------------------------------------------------------------------------
static int
parseCmdLine( int argc, char ** argv )
{
  int i;

  for( i = 1; i < argc; i++ )
  {
       if( !strncmp( *(argv + i), "-f", 2 ) )
       {
           if( strlen( *(argv + i) ) > 2 )
           {
               strncpy( szFile, *( argv + i ) + 2, MAXPATHLEN );
           }
           bScanFile = 1;
       }
       else if( !strncmp( *(argv + i), "-sf", 3 )
                &&
                strlen( *(argv + i) ) > 3 )
       {
           strncpy( szSocketFile, *( argv + i ) + 3, MAXPATHLEN );
       }
       else if( !strncmp( *(argv + i), "-i", 2 )
                &&
                strlen( *(argv + i) ) > 2 )
       {
           strncpy( szDaemonIP, *( argv + i ) + 2, MAXPATHLEN );
       }
       else if( !strncmp( *(argv + i), "-ex", 3 ) )
       {
           bExpand = 1;
       }
       else if( !strncmp( *(argv + i), "-p", 2 )
                &&
                strlen( *(argv + i) ) > 2 )
       {
           Port = atoi( *( argv + i ) + 2 );
       }
       else if( !strncmp( *(argv + i), "-t", 2 )
                &&
                strlen( *(argv + i) ) > 2 )
       {
           Timeout = atoi( *( argv + i ) + 2 );
       }
       else if( !strcmp( *(argv + i), "-ha" ) )
       {
           scanOptions |= DRWEBD_HEURISTIC_ON;
       }
       else if( !strcmp( *(argv + i), "-hf" ) )
       {
           scanOptions |= DRWEBD_SPAM_FILTER;
       }
       else if( !strcmp( *(argv + i), "-rv" ) )
       {
           scanOptions |= DRWEBD_RETURN_VIRUSES;
       }
       else if( !strcmp( *(argv + i), "-rr" ) )
       {
           scanOptions |= DRWEBD_RETURN_REPORT;
       }
       else if( !strcmp( *(argv + i), "-rc" ) )
       {
           scanOptions |= DRWEBD_RETURN_CODES;
       }
       else if( !strcmp( *(argv + i), "-q" ) )
       {
           bQuiet = 1;
       }
       else if( !strcmp( *(argv + i), "-sb" ) )
       {
           bShowBases = 1;
       }
       else if( !strcmp( *(argv + i), "-si" ) )
       {
           bShowID = 1;
       }
       else if( !strcmp( *(argv + i), "-sv" ) )
       {
           bShowVersion = 1;
       }
       else
       {
           return 0;
       }
  }
  return argc > 1;
}

//--------------------------------------------------------------------------
static char *
expandBuffer( char * oldbuffer, int buflen, char * data, int datalen )
{

// no realloc() because this template may be used in C++ sources

  char * newbuffer = (char *)malloc( buflen + datalen );
  if( newbuffer )
  {
      if( oldbuffer )
      {
          memcpy( newbuffer, oldbuffer, buflen );
          free( oldbuffer );
      }
      memcpy( newbuffer + buflen, data, datalen );
      return newbuffer;
  }
  if( oldbuffer )
  {
      free( oldbuffer );
  }
  prints( "expandBuffer(): memory allocation error\n" );
  return NULL;
}

//--------------------------------------------------------------------------
static const char *
findWhat( int what )
{
    if( what == DERR_READ_ERR ) return "read errors: ";
    if( what == DERR_WRITE_ERR ) return "write errors: ";
    if( what == DERR_NOMEMORY ) return "memory errors: ";
    if( what == DERR_CRC_ERROR ) return "CRC errors: ";
    if( what == DERR_READSOCKET ) return "socket errors: ";
    if( what == DERR_KNOWN_VIRUS ) return "known viruses: ";
    if( what == DERR_UNKNOWN_VIRUS ) return "unknown viruses: ";
    if( what == DERR_VIRUS_MODIFICATION ) return "virus modifcations: ";
    if( what == DERR_TIMEOUT ) return "timeouts: ";
    if( what == DERR_SYMLINK ) return "symlinks: ";
    if( what == DERR_NO_REGFILE ) return "no regular files: ";
    if( what == DERR_SKIPPED ) return "skipped files (bad format): ";
    if( what == DERR_TOO_BIG ) return "too big files: ";
    if( what == DERR_TOO_COMPRESSED ) return "too high compressed ratio: ";
    if( what == DERR_BAD_CALL ) return "invalid calls: ";
    if( what == DERR_SPAM_MESSAGE ) return "filtered messages: ";
    if( what == DERR_EVAL_VERSION ) return "evaluation: ";
    return "unknown codes: ";
}

//--------------------------------------------------------------------------
static int
showResult( int retcode )
{
   retcode = ntohl( retcode );
   prints( "\nDr. Web daemon return " );
   printn( retcode );

   if( (scanOptions & DRWEBD_RETURN_VIRUSES) == DRWEBD_RETURN_VIRUSES )
   {
       int vnum = 0;
       if( !readOutput( sizeof(vnum), &vnum ) )
       {
           int i;
           vnum = ntohl( vnum );
           for( i = 0; i < vnum; i++ )
           {
                int vlen = 0;
                if( !readOutput( sizeof(vlen), &vlen ) )
                {
                    char virname[ 1024 ];
                    vlen = ntohl( vlen );
                    if( !readOutput( vlen > sizeof(virname) ? sizeof(virname) : vlen, virname ) )
                    {
                        prints( "\n" );
                        prints( virname );
                    }
                }
           }
       }
   }

   if( (scanOptions & DRWEBD_RETURN_REPORT) == DRWEBD_RETURN_REPORT )
   {
       int lnum = 0;
       if( !readOutput( sizeof(lnum), &lnum ) )
       {
           int i;
           lnum = ntohl( lnum );
           prints( "\n" );
           for( i = 0; i < lnum; i++ )
           {
                int llen = 0;
                if( !readOutput( sizeof(llen), &llen ) )
                {
                    char logline[ 1024 ];
                    llen = ntohl( llen );
                    if( !readOutput( llen > sizeof(logline) ? sizeof(logline) : llen, logline ) )
                    {
                        prints( logline );
                    }
                }
           }
       }
   }

   if( (scanOptions & DRWEBD_RETURN_CODES) == DRWEBD_RETURN_CODES )
   {
       int cnum = 0;
       if( !readOutput( sizeof(cnum), &cnum ) )
       {
           int i;
           cnum = ntohl( cnum );
           prints( "\n" );
           for( i = 0; i < cnum; i++ )
           {
                int what = 0, val = 0;
                if( !readOutput( sizeof(what), &what ) &&
                    !readOutput( sizeof(val), &val ) )
                {
                    what = ntohl( what );
                    val = ntohl( val );
                    prints( findWhat( what ) );
                    printn( val );
                    prints( "\n" );
                }
        }
       }
   }  

   closeSocket();
   return retcode;
}

//--------------------------------------------------------------------------
static int
scanBuffer( const char * pFileBuffer, int fileSize )
{
   int retvalue = 0, err = 0;

   if( createSocket() < 0 )
   {
       return DRWEBD_SOCKET_ERR;
   }

   if( !writeInt( DRWEBD_SCAN_CMD ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( scanOptions ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( 0 ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( fileSize ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeBuf( pFileBuffer, fileSize ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   err = readOutput( sizeof(retvalue), &retvalue );
   if( err )
   {
       closeSocket();
       return err;
   }
   return showResult( retvalue );
}

//--------------------------------------------------------------------------
static int
scanFile( const char * szFile, int inFile, int fileSize )
{
   int retvalue = 0, err = 0;

   if( createSocket() < 0 )
   {
       return DRWEBD_SOCKET_ERR;
   }

   if( !writeInt( DRWEBD_SCAN_CMD ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( scanOptions ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( strlen(szFile) ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeBuf( szFile, strlen(szFile) ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( fileSize ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   while( fileSize )
   {
      char buf[ 0x1000 ];
      int towrite = sizeof(buf);
      if( towrite > fileSize )
      {
          towrite = fileSize;
      }
      if( read( inFile, buf, towrite ) != towrite )
      {
          prints( "\nError reading " );
          prints( szFile );
          closeSocket();
          return DRWEBD_IN_FILE;
      }
      if( !writeBuf( buf, towrite ) )
      {
          closeSocket();
          return DRWEBD_SWRITE_ERR;
      }

      fileSize -= towrite;
   }

   err = readOutput( sizeof(retvalue), &retvalue );

   if( err )
   {
       closeSocket();
       return ntohl( err );
   }
   bExpanded = 1;
   return showResult( retvalue );
}

//--------------------------------------------------------------------------
static int
showVersion( void )
{
   if( createSocket() < 0 )
   {
       return DRWEBD_SOCKET_ERR;
   }

   if( writeInt( DRWEBD_VERSION_CMD ) )
   {
       int ver = 0;
       if( !readOutput( sizeof(ver), &ver ) )
       {
           ver = ntohl( ver );
           prints( "\nDr. Web daemon version: " );
           printn( ver / 100 );
           prints( "." );
           printn( ver % 100 );
       }
   }
   closeSocket();
   return 0;
}

//--------------------------------------------------------------------------
static int
showID( void )
{
   if( createSocket() < 0 )
   {
       return DRWEBD_SOCKET_ERR;
   }

   if( writeInt( DRWEBD_IDSTRING_CMD ) )
   {
       int sz = 0;
       if( !readOutput( sizeof(sz), &sz ) )
       {
           char buf[ 1024 ];
           sz = ntohl( sz );
           if( !readOutput( sz > sizeof(buf) ? sizeof(buf) : sz, buf ) )
           {
                prints( "\n" );
                prints( buf );
           }
       }
   }
   closeSocket();
   return 0;
}

//--------------------------------------------------------------------------
static int
showBases( void )
{
   if( createSocket() < 0 )
   {
       return DRWEBD_SOCKET_ERR;
   }

   if( writeInt( DRWEBD_BASEINFO_CMD ) )
   {
       int bnum = 0;
       if( !readOutput( sizeof(bnum), &bnum ) )
       {
           int i;
           bnum = ntohl( bnum );
           for( i = 0; i < bnum; i++ )
           {
                int nlen = 0;
                if( !readOutput( sizeof(nlen), &nlen ) )
                {
                    char bfile[ 1024 ];
                    nlen = ntohl( nlen );
                    if( !readOutput( nlen > sizeof(bfile) ? sizeof(bfile) : nlen, bfile ) )
                    {
                        int vir = 0;
                        if( !readOutput( sizeof(vir), &vir ) )
                        {
                            vir = ntohl( vir );
                            prints( "\nVirus base: " );
                            prints( bfile );
                            prints( " (" );
                            printn( vir );
                            prints( " records)" );
                        }
                    }
                }
           }
       }
   }
   closeSocket();
   return 0;
}

//--------------------------------------------------------------------------
static int
scanDiskFile( const char * szFile )
{
   int retvalue = 0, err = 0;

   bRealFile = 1;

   if( createSocket() < 0 )
   {
       return DRWEBD_SOCKET_ERR;
   }

   if( !writeInt( DRWEBD_SCAN_CMD ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( scanOptions ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( strlen(szFile) ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeBuf( szFile, strlen(szFile) ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   if( !writeInt( 0 ) )
   {
       closeSocket();
       return DRWEBD_SWRITE_ERR;
   }

   err = readOutput( sizeof(retvalue), &retvalue );

   if( err )
   {
       closeSocket();
       return ntohl( err );
   }
   return showResult( retvalue );
}

//--------------------------------------------------------------------------
int
main( int argc, char ** argv )
{
   int retcode = 0;

   prints( "\n\nWarning! This is DEMO only. Explore \"clients/drwebdc\", \"clients/qmail\"," );
   prints( "\n\"clients/exim\", \"clients/postfix\" and \"clients/sendmail\" directories" );
   prints( "\nof DrWeb distribution for details.\n" );

   // parse command line:
   if( !parseCmdLine( argc, argv ) )
   {
       return usage();
   }

   // no port or Unix socket in command line found:
   if( Port == -1 && !*szSocketFile )
   {
       return usage();
   }

   if( bShowVersion )
   {
       showVersion();
   }

   if( bShowID )
   {
       showID();
   }

   if( bShowBases )
   {
       showBases();
   }

   if( bScanFile )
   {

       if( *szFile )
       {
           if( bExpand )
           {
               char buf[ MAXPATHLEN + 2 ];
               realpath( szFile, buf );
               strncpy( szFile, buf, MAXPATHLEN + 1 );
           }

           if( *szFile != '/' )
           {
                int  fileSize = 0;
                inFile = open( szFile, O_RDONLY );
                if( inFile == -1 )
                {

                    prints( "\nError reading " );
                    prints( szFile );
                    cleanupClient();
                    return DRWEBD_IN_FILE;
                 }

                 fileSize = lseek( inFile, 0, SEEK_END );
                 if( fileSize == -1 )
                 {
                     prints( "\nError reading " );
                     prints( szFile );
                     cleanupClient();
                     return DRWEBD_IN_FILE;
                 }

                 if( lseek( inFile, 0, SEEK_SET ) == -1 )
                 {
                     prints( "\nError reading " );
                     prints( szFile );
                     cleanupClient();
                     return DRWEBD_IN_FILE;
                  }
          retcode = scanFile( szFile, inFile, fileSize );

            }
            else
            {
              retcode = scanDiskFile( szFile );
            }
       }
       else
       {
           char buf[ 0x4000 ];
           int  fileSize = 0;
           int readed = read( fileno(stdin), buf, sizeof(buf) );
           while( readed > 0 )
           {
               pFileBuffer = expandBuffer( pFileBuffer, fileSize, buf, readed );
               if( !pFileBuffer )
               {
                   cleanupClient();
                   return DRWEBD_NO_MEMORY;
               }
               fileSize += readed;
               readed = read( fileno(stdin), buf, sizeof(buf) );
           }

           retcode = scanBuffer( pFileBuffer, fileSize );
       }
   }

   cleanupClient();

   return retcode;
}

//--------------------------------------------------------------------------
