// main.cpp
// SNF_Client.exe
//
// (C) 2006-2008 ARM Research Labs, LLC.
// See www.armresearch.com for the copyright terms.
//
// This program implements a client command line interface for systems using
// SNF_Server. The operation is simple--- Given a file name to scan, pass the
// file name to the SNF scanner on localhost and return the result code.

#if defined(WIN32) || defined(WIN64)
#include <windows.h>
#endif

#include <iostream>
#include <fstream>
#include <sstream>
#include <cstring>
#include <string>
#include <ctime>
#include <cctype>
#include <unistd.h>
#include <dirent.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include "timing.hpp"
#include "networking.hpp"
#include "threading.hpp"
#include "snf_xci.hpp"

#include "config.h"

using namespace std;                                                            // Introduce standard namespace.

const char* VERSION_INFO = "SNF Client Version " PACKAGE_VERSION " Build: " __DATE__ " " __TIME__;

time_t Timestamp() {                                                            // Get an ordinary timestamp.
    time_t rawtime;                                                             // Grab raw time from
    time(&rawtime);                                                             // the system clock.
    return rawtime;
}

string Timestamp(time_t t) {                                                    // Convert time_t to a timestamp.
    char TimestampBfr[20];                                                      // Create a small buffer.
    tm* gmt;                                                                    // Get a ptr to a tm structure.
    gmt = gmtime(&t);                                                           // Fill it with UTC.
    sprintf(TimestampBfr,"%04d%02d%02d%02d%02d%02d\0",                          // Format yyyymmddhhmmss
      gmt->tm_year+1900,
      gmt->tm_mon+1,
      gmt->tm_mday,
      gmt->tm_hour,
      gmt->tm_min,
      gmt->tm_sec
    );
    return string(TimestampBfr);                                                // Return a string.
}

string ErrorMessage(int argc, char** argv, string ErrorText) {                  // Create an error message.

    ostringstream FailurePreamble;                                              // Setup a stringstream for formatting.

    string TimestampString = Timestamp(Timestamp());                            // Get the UTC timestamp.

    FailurePreamble << TimestampString << ", ";                                 // Start with the timestamp and
    for(int a = 1;  a < argc; a++) {                                            // then roll in all of the command
        FailurePreamble << "arg" << a << "=" << argv[a];                      // line parameters. Follow that with
        if(a < argc-1) { FailurePreamble << ", "; }                             // a colon and the error itself.
        else { FailurePreamble << " : "; }
    }

    FailurePreamble << ErrorText << endl;                                       // Tack on the error text.

    return FailurePreamble.str();                                               // Return the completed string.
}

int main(int argc, char* argv[]) {                                              // Accept command line parms.

    // Figure out what we're going to do and create a suitable RequestString.

    bool DebugMode = false;                                                     // This will be our debug mode.
    string argv0(argv[0]);                                                      // Capture how we were called.
    if(
      string::npos != argv0.find("Debug") ||                                    // If we find "Debug" or
      string::npos != argv0.find("debug")                                       // "debug" in our command path
      ) {                                                                       // then we are in DebugMode.
        DebugMode = true;                                                       // Set the flag and tell the
        cout << "Debug Mode" << endl;                                           // watchers.
    }

    string FileToScan;                                                          // What will we be scanning?
    string RequestString;                                                       // What will we send to the server?
    bool ItIsACommand = false;                                                  // Is it a command?
    bool ItIsAScan = false;                                                     // Is it a scan?
    bool GetXHDRs = false;                                                      // GetXHDRs with scan?
    bool ItIsAGBUdb = false;                                                    // Is it a GBUdb Request?
    bool ItIsAReport = false;                                                   // Is it a Status Report Request?

    switch(argc) {
        case 2: {                                                               // This is either a file or a command.
            FileToScan = argv[1];

            // Status Report? --------------------------------------------------

            if(0 == FileToScan.find("-status.second")) {                        // Status Second Report? If so:
                ItIsAReport = true;                                             // Set the flag & format the request.

                RequestString = "<snf><xci><report><request><status class='second'/></request></report></xci></snf>";

            } else

            if(0 == FileToScan.find("-status.minute")) {                        // Status Minute Report? If so:
                ItIsAReport = true;                                             // Set the flag & format the request.

                RequestString = "<snf><xci><report><request><status class='minute'/></request></report></xci></snf>";

            } else

            if(0 == FileToScan.find("-status.hour")) {                          // Status Hour Report? If so:
                ItIsAReport = true;                                             // Set the flag & format the request.

                RequestString = "<snf><xci><report><request><status class='hour'/></request></report></xci></snf>";

            } else

            if(0 == FileToScan.find("-shutdown")) {                             // If argv1 is -shutdown:
                ItIsACommand = true;                                            // Set the command flag.
                if(DebugMode) { cout << "Command: shutdown" << endl; }          // In debug - tell what we are doing.

                // Format the request.

                RequestString = "<snf><xci><server><command command=\'shutdown\'/></server></xci></snf>\n";

            } else

            if('-' == FileToScan.at(0)) {                                       // If argv1 still looks like -something
                goto BadCommandLine;                                            // Complain and show the help text.

            } else {                                                            // If it does not then it is a file.
                ItIsAScan = true;                                               // Set the scan flag.
                if(DebugMode) { cout << "Scan: " << FileToScan << endl; }       // In debug - tell what we are doing.

                // Format the request.

                RequestString = "<snf><xci><scanner><scan file=\'";
                RequestString.append(FileToScan);
                RequestString.append("\'/></scanner></xci></snf>\n");
            }

            break;
        }

        case 3: {                                                               // Special Scan Mode
            string CommandString = argv[1];                                     // -command
            FileToScan = argv[2];                                               // What file to scan?

            // XHeader Scan? ---------------------------------------------------

            if(0 == CommandString.find("-xhdr")) {                              // Scan file and show x headers.
                ItIsAScan = true;                                               // Set the scan flag.
                GetXHDRs = true;                                                // Turn on XHDRs for the scan.
                if(DebugMode) {
                    cout << "(xhdr)Scan: " << FileToScan << endl;               // In debug - tell what we are doing.
                }

                // Format the request.

                RequestString = "<snf><xci><scanner><scan file=\'";
                RequestString.append(FileToScan);
                RequestString.append("\' xhdr=\'yes\'/></scanner></xci></snf>\n");

            } else

            // Source Scan? ----------------------------------------------------

            if(0 == CommandString.find("-source=")) {                           // Scan file with forced source.
                ItIsAScan = true;                                               // Set the scan flag.

                const int SourceIPIndex = CommandString.find("=") + 1;          // Find source after "-source="
                IP4Address SourceIP = CommandString.substr(SourceIPIndex);      // Extract the source IP.

                if(DebugMode) {
                    cout << "([" << (string) SourceIP                           // In debug - tell what we are doing.
                         << "])Scan: " << FileToScan << endl;
                }

                // Format the request.

                RequestString = "<snf><xci><scanner><scan file=\'";
                RequestString.append(FileToScan);
                RequestString.append("\' ip=\'");
                RequestString.append((string) SourceIP);
                RequestString.append("\'/></scanner></xci></snf>\n");

            } else

            // GBUdb test? -----------------------------------------------------

            if(0 == CommandString.find("-test")) {                              // GBUdb test IP
                ItIsAGBUdb = true;                                              // Set the GBUdb flag.

                // Format the request.

                RequestString = "<snf><xci><gbudb><test ip=\'";
                RequestString.append(argv[2]);
                RequestString.append("\'/></gbudb></xci></snf>\n");

            } else

            // GBUdb good? -----------------------------------------------------

            if(0 == CommandString.find("-good")) {                              // GBUdb good IP event
                ItIsAGBUdb = true;                                              // Set the GBUdb flag.

                // Format the request.

                RequestString = "<snf><xci><gbudb><good ip=\'";
                RequestString.append(argv[2]);
                RequestString.append("\'/></gbudb></xci></snf>\n");

            } else

            // GBUdb bad? ------------------------------------------------------

            if(0 == CommandString.find("-bad")) {                               // GBUdb bad IP event
                ItIsAGBUdb = true;                                              // Set the GBUdb flag.

                // Format the request.

                RequestString = "<snf><xci><gbudb><bad ip=\'";
                RequestString.append(argv[2]);
                RequestString.append("\'/></gbudb></xci></snf>\n");

            } else

            // GBUdb drop? -----------------------------------------------------

            if(0 == CommandString.find("-drop")) {                              // GBUdb drop IP
                ItIsAGBUdb = true;                                              // Set the GBUdb flag.

                // Format the request.

                RequestString = "<snf><xci><gbudb><drop ip=\'";
                RequestString.append(argv[2]);
                RequestString.append("\'/></gbudb></xci></snf>\n");

            } else

            // Compatibility Scan? ---------------------------------------------

            if(CommandString.at(0) != '-') {                                    // If we don't see -<something>:
                ItIsAScan = true;                                               // Set the scan flag.
                FileToScan = argv[2];                                           // Grab the message_file_name for
                if(DebugMode) {                                                 // a compatability scan and if
                    cout << "(compat)Scan: " << FileToScan << endl;             // debugging then announce it.
                }

                // Format the request.

                RequestString = "<snf><xci><scanner><scan file=\'";
                RequestString.append(FileToScan);
                RequestString.append("\'/></scanner></xci></snf>\n");

            } else

            goto BadCommandLine;                                                // If no match here, give help.

            break;
        }

        case 4: {

            string Command1String = argv[1];                                    // -command1
            string Command2String = argv[2];                                    // -command2
            FileToScan = argv[3];                                               // What file to scan?

            // XHeader Scan W/ Source IP? --------------------------------------

            if(0 == Command1String.find("-source=")) {                          // If things are refersed
                string tmp = Command2String;                                    // swap them to the order we
                Command2String = Command1String;                                // are expecting. If we're wrong
                Command1String = tmp;                                           // then that case will be handled
            }                                                                   // next step.

            if(
              0 == Command1String.find("-xhdr") &&
              0 == Command2String.find("-source=")
              ) {                                                               // Scan file and show x headers.
                ItIsAScan = true;                                               // Set the scan flag.
                GetXHDRs = true;                                                // Turn on XHDRs for the scan.

                const int SourceIPIndex = Command2String.find("=") + 1;         // Find source after "-source="
                IP4Address SourceIP = Command1String.substr(SourceIPIndex);     // Extract the source IP.

                if(DebugMode) {
                    cout << "(xhdr [" << (string) SourceIP                      // In debug - tell what we are doing.
                         << "])Scan: " << FileToScan << endl;
                }

                // Format the request.

                RequestString = "<snf><xci><scanner><scan file=\'";
                RequestString.append(FileToScan);
                RequestString.append("\' xhdr=\'yes\' ip=\'");
                RequestString.append((string) SourceIP);
                RequestString.append("\'/></scanner></xci></snf>\n");

            } else

            goto BadCommandLine;                                                // If no match here, give help.

            break;
        }

        case 6: {                                                               // GBUdb set mode.
            string CommandString = argv[1];                                     // -command

            // GBUdb set? ------------------------------------------------------

            if(0 == CommandString.find("-set")) {                               // GBUdb set IP
                ItIsAGBUdb = true;                                              // Set the GBUdb flag.

                // Format the request.

                RequestString = "<snf><xci><gbudb><set ip=\'";                  // Capture the IP in the request.
                RequestString.append(argv[2]);
                RequestString.append("\'");

                if(                                                             // If the type flag is specified
                  'g' == argv[3][0] ||                                          // it must begin with G B U or I.
                  'G' == argv[3][0] ||
                  'b' == argv[3][0] ||
                  'B' == argv[3][0] ||
                  'u' == argv[3][0] ||
                  'U' == argv[3][0] ||
                  'i' == argv[3][0] ||
                  'I' == argv[3][0]
                  ) {                                                           // If we've got a type flag then
                    RequestString.append(" type=\'");                           // capture it in our request.
                    RequestString.append(argv[3]);
                    RequestString.append("\'");
                }                                                               // If the type flag isn't specified
                else if('-' != argv[3][0]) goto BadCommandLine;                 // it must be a - or else an error.

                if(isdigit(argv[4][0])) {                                       // Capture the bad count.
                    RequestString.append(" b=\'");                              // If it looks like a number
                    RequestString.append(argv[4]);                              // capture it in our request.
                    RequestString.append("\'");
                }
                else if('-' != argv[4][0]) goto BadCommandLine;                 // If not a number it must be -

                if(isdigit(argv[5][0])) {                                       // Capture the good count.
                    RequestString.append(" g=\'");                              // If it looks like a number
                    RequestString.append(argv[5]);                              // capture it in our request.
                    RequestString.append("\'");
                }
                else if('-' != argv[5][0]) goto BadCommandLine;                 // If not a number it must be -

                RequestString.append("/></gbudb></xci></snf>\n");               // Finish off our request.
            }

            else goto BadCommandLine;                                           // If no match here, give help.

            break;
        }

        // Bad Command Line ----------------------------------------------------

        default: {                                                              // If we don't know: show help.
            BadCommandLine:                                                     // Shortcut for others.
            cout
              << endl
              << VERSION_INFO << endl
              << endl
              << "Help:" << endl
              << endl
              << " To scan a message file use: " << endl
              << "     SNFClient.exe [-xhdr] [-source=<IP4Address>] <FileNameToScan>" << endl
              << " or: SNFClient.exe <Authenticationxx> <FileNameToScan>" << endl
              << endl
              << " To test an IP with GBUdb use: " << endl
              << "     SNFClient.exe -test <IP4Address>" << endl
              << endl
              << " To update GBUdb records use: " << endl
              << "     SNFClient.exe -set <IP4Address> <flag> <bad> <good>" << endl
              << " or: SNFClient.exe -drop <IP4Address>" << endl
              << " or: SNFClient.exe -good <IP4Address>" << endl
              << " or: SNFClient.exe -bad <IP4Address>" << endl
              << endl
              << " To check SNFServer status use: " << endl
              << "     SNFClient.exe -status.second" << endl
              << " or: SNFClient.exe -status.minute" << endl
              << " or: SNFClient.exe -status.hour" << endl
              << endl
              << " To shut down the SNFServer use: " << endl
              << "     SNFClient.exe -shutdown" << endl
              << endl
              << " For more information see www.armresearch.com" << endl
              << " (C) 2007-2008 Arm Research Labs, LLC." << endl;

            return 0;                                                           // Return 0 on help.
        }
    }

    // If we're in debug mode this is a good time to emit our request string

    if(DebugMode) { cout << RequestString << endl; }                            // If debugging, show our request.

    // Since we're gonna do this -- prepare for an error

    string ERRFname;                                                            // The default error log name will
    ERRFname = argv[0];                                                         // be the program file name with
    ERRFname.append(".err");                                                    // .err tagged on the end.
    const int FailSafeResult = 0;                                               // Fail Safe result code.
    const int StatusReportError = 99;                                           // Unknown Error Code for status failures.

    // Connect to the server and get the result...

    string ResultString;                                                        // We need our result string.
    bool ConnectSuccess = false;                                                // We need our success flag.

    // Max time in this loop should be (100*50ms) = 5 seconds per try times
    // 10 tries = 50 seconds, plus (9*500ms) = 4.5 secs for re-tries. ~ 55 secs.

    const int ResultBufferSize = 4096;
    char ResultBuffer[ResultBufferSize+1];                                      // Make an oversize buffer for the answer.
    memset(ResultBuffer, 0, sizeof(ResultBuffer));                              // Set the entire thing to nulls.

    const int Tries = 20;                                                       // How many times to try this.
    Sleeper SleepAfterAttempt(100);                                             // How long to sleep between attempts.

    const int OpenTries = 90;                                                   // How many tries at opening.
    Sleeper WaitForOpen(10);                                                    // How long to wait for an open cycle.

    const int ReadTries = 900;                                                  // How many tries at reading.
    Sleeper SleepBeforeReading(10);                                             // How long to pause before reading.

    /*
    ** 20 * 100ms = 2 seconds for all tries.
    ** 90 * 10ms = 900ms for a failed connection.
    ** 900 * 10ms = 9 seconds for a failed read.
    **
    ** Approximate wait for can't connect = 2.0 + (20 * 0.9) = ~ 20.0 seconds.
    ** Maximum impossible wait = 2.0 + (0.9 * 20) + (9.0 * 20) = 200.0 seconds.
    */

    for(int tryagain = Tries; (0<tryagain) && (!ConnectSuccess); tryagain--) {  // Try a few times to get this done.
        try {
            ResultString = "";                                                  // Clear our result string.
            TCPHost SNFServer(9001);                                            // Create connection to server.
            SNFServer.makeNonBlocking();                                        // Make it non-blocking.

            for(int tries = OpenTries; 0 < tries; tries--) {                    // Wait & Watch for a good connection.
                try { SNFServer.open(); } catch(...) {}                         // Try opening the connection.
                if(SNFServer.isOpen()) break;                                   // When successful, let's Go!
                else WaitForOpen();                                             // When not successful, pause.
            }

            if(SNFServer.isOpen()) {                                            // If we have a good connection:

                SNFServer.transmit(
                  RequestString.c_str(), RequestString.length());               // Send the request.

                for(int tries = ReadTries; 0 < tries; tries--) {                // Try to read the result a few times.
                    SleepBeforeReading();                                       // Provide some time for each try.
                    memset(ResultBuffer, 0, sizeof(ResultBuffer));              // Clear the buffer.
                    SNFServer.receive(ResultBuffer, ResultBufferSize);          // Receive the answer.
                    ResultString.append(ResultBuffer);
                    if(string::npos ==
                      ResultString.rfind("</snf>",ResultString.length())) {     // If we don't have the end yet.
                        continue;                                               // Try again.
                    } else {                                                    // If we got to end of line
                        ConnectSuccess = true;                                  // Success!
                        break;                                                  // We're done.
                    }
                }
                SNFServer.close();                                              // No need for our connection after that.
            }
        } catch(...) { }                                                        // Ignore errors for now.
        if(!ConnectSuccess) SleepAfterAttempt();                                // Pause for a moment before trying again..
    }

    if(!ConnectSuccess) {                                                       // If no connection success complain!

        if(DebugMode) {
            cout << FileToScan << ": ";
            cout << "Could Not Connect!" << endl;
        }

        ofstream ERRF(ERRFname.c_str(), ios::out | ios::ate | ios::app);
        ERRF << ErrorMessage(argc, argv, "Could Not Connect!");
        ERRF.close();

        if(ItIsAReport) {                                                       // If this was a status request then
            return StatusReportError;                                           // return a failure value.
        }                                                                       // In all other cases return zero as a
        else return FailSafeResult;                                             // Fail Safe.
    }

    // At this point we should have a usable result.

    if(DebugMode) { cout << ResultString << endl; }                             // In debug, show the result string.

    snf_xci Reader(ResultString);                                               // Interpret the data and check for
    if(Reader.bad()) {                                                          // a proper read. If it was bad then

        if(DebugMode) {
            cout << FileToScan << ": ";
            cout << "Bad result from server!" << endl;
            cout << ResultString << endl;
        }

        ofstream ERRF(ERRFname.c_str(), ios::out | ios::ate | ios::app);
        ERRF << ErrorMessage(argc, argv,                                        // complain and spit out what we got
          "Bad result from server! " + ResultString);                           // for debugging purposes.
        return FailSafeResult;                                                  // Return our failsafe value.
    }

    if(0 < Reader.xci_error_message.length()) {                                 // If the result was a general error

        if(DebugMode) {
            cout << FileToScan << ": ";
            cout << "XCI Error!: " << Reader.xci_error_message << endl;
        }

        ofstream ERRF(ERRFname.c_str(), ios::out | ios::ate | ios::app);
        ERRF << ErrorMessage(argc, argv,
          "XCI Error!: " + Reader.xci_error_message);                           // then spit that out
        return FailSafeResult;                                                  // and return the failsafe.
    }

    // If we got here we've successfully parsed the results.

    if(ItIsACommand) {                                                          // If it was a command then
        cout << " [Server Says: " << Reader.xci_server_response << "]" << endl; // show the response.

    }

    else if(ItIsAScan) {                                                        // If it was a scan then
        int ResultCode = Reader.scanner_result_code;                            // grab the result code.

        if(0 < Reader.scanner_result_xhdr.length()) {                           // If we have xheaders show them.
            cout << endl << Reader.scanner_result_xhdr << endl;
        }

        if(DebugMode) { cout << "(" << ResultCode << ")"; }

        if( (0 < ResultCode) && (64 > ResultCode) ) {                           // If the result code means SPAM
            if(DebugMode) { cout << "[spam]" << endl; }
            return ResultCode;                                                  // Return the result code.
        } else
        if(0 == ResultCode) {
            if(DebugMode) { cout << "[clean]" << endl; }
            return 0;
        } else {
            if(DebugMode) { cout << "[Fail Safe!]" << endl; }
            return FailSafeResult;
        }
    }

    else if(ItIsAGBUdb) {                                                       // If it was a GBUdb function then show
        cout                                                                    // the results to whomever is watching.
          << "GBUdb Record for " << Reader.gbudb_result_ip << endl
          << "  Type Flag: " << Reader.gbudb_result_type << endl
          << "  Bad Count: " << Reader.gbudb_result_bad_count << endl
          << " Good Count: " << Reader.gbudb_result_good_count << endl
          << "Probability: " << Reader.gbudb_result_probability << endl
          << " Confidence: " << Reader.gbudb_result_confidence << endl
          << "      Range: " << Reader.gbudb_result_range << endl
          << "       Code: " << Reader.gbudb_result_code << endl
          << endl;

        return Reader.gbudb_result_code;
    }

    else if(ItIsAReport) {                                                      // If it was a report request then
        cout                                                                    // output the report. Add a handy
          << endl << "<!-- Status Report -->" << endl                           // XML comment to help humans.
          << Reader.report_response << endl;

        return 0;
    }

    // You can't get here from there.
    if(DebugMode) { cout << "End Of Logic [Fail Safe!!]" << endl; }
    return FailSafeResult;
}