// main.cpp
// SNF MDaemon Plugin
// Copyright (C) 2007-2008, ARM Research Labs, LLC.

#include <tchar.h>
#include <winerror.h>
#include <wx/app.h>
#include <wx/window.h>
#include <wx/dialog.h>
#include <wx/toplevel.h>
#include "dialogapp.h"
#include "configdialog.h"
#include <winsock2.h>
#include <windows.h>
#include "mdconfiguration.hpp"
#include "SNFMulti/SNFMulti.hpp"


#ifndef _UNICODE
#define _UNICODE
#endif
#ifndef UNICODE
#define UNICODE
#endif



//IMPLEMENT_APP_NO_MAIN(dialogapp);
//IMPLEMENT_WX_THEME_SUPPORT;

////////////////////////////////////////////////////////////////////////////////
// Constants And Tweaks
////////////////////////////////////////////////////////////////////////////////

const int MDPLUGIN_MSG = 22000;
const int MDPLUGIN_DISPLAY = 22001;

const char* PLUGIN_VERSION_INFO = "SNF MDaemon Plugin Version 4.0 Build: " __DATE__ " " __TIME__;

BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {  // Thread attach/detach functions.



    switch (fdwReason) {                                                        // Switch on the reason for the call.
        case DLL_PROCESS_ATTACH: 
        {                                               // Attach to process.
            /*** Startup function moved to API ***/
            }                                  // We will assume it worked.
            break;

        case DLL_PROCESS_DETACH:                                                // Detach from process.
            /*** Shutdown function moved to API ***/
            return true;                                                        // We will assume that worked.
            break;

        case DLL_THREAD_ATTACH:                                                 // Attach to thread
            break;                                                              // Just be happy.

        case DLL_THREAD_DETACH:                                                 // Detach from thread
            break;                                                              // Just be happy.
    }
    return TRUE;                                                                // Be happy by default.
}
////////////////////////////////////////////////////////////////////////////////
// SNF MDaemon Plugin Engine Code
////////////////////////////////////////////////////////////////////////////////

// Handy Debug Functions. This is how we can display critical events on screen
// and emit entries into the MDaemon logs.

HWND LatestParent = 0;                                                          // Handle MDaemon log operations.

void sayToScreen(const string StuffToSay) {                                     // Display a modal system message.
  MessageBox (                                                                  // Use a message box.
    GetFocus(),                                                                 // Take the user's focus.
    (LPCWSTR) StuffToSay.c_str(),                                                         // This is what to say.
    (LPCWSTR) _T("Message Sniffer Plug-In"),                                                  // This is how to title the box.
    MB_OK|MB_SYSTEMMODAL );                                                     // This makes it modal with a button.
}

void sayToLog(const string Msg) {                                               // Emit message into MDaemon log.
  if(0 != LatestParent) {                                                       // If we have a handle:

    COPYDATASTRUCT Packet;                                                      // Create a packet for the log.
    Packet.dwData = MDPLUGIN_DISPLAY;
    Packet.cbData = Msg.length();
    Packet.lpData = reinterpret_cast<void*>(
                      const_cast<char*>(Msg.c_str()));

    SendMessage(LatestParent, MDPLUGIN_MSG, MDPLUGIN_DISPLAY,                   // Send the packet.
               (LPARAM)(PCOPYDATASTRUCT)&Packet);
  }
}

// The configuration file path is theoretically in a fixed location based on
// the installation directory of MDaemon. We read the installation directory
// from the registry and derive the path from that. If we get what we want
// then we return the full path to the SNFMDPlugin.xml file - which is derived
// from snf_engine.xml. If we don't get what we want then we return nothing.

string ConfigurationFilePath() {                                                // Get the config file path.
    std::wstring install_pathW;
	HKEY hKey;
	WCHAR szBuffer[512];
	DWORD dwBufferSize = sizeof(szBuffer);
	std::wstring regValue = L"";

	if (ERROR_SUCCESS == RegOpenKeyEx(HKEY_LOCAL_MACHINE, L"SOFTWARE\\Alt-N Technologies\\MDaemon", 0, KEY_READ, &hKey)){
		if(ERROR_SUCCESS == RegQueryValueEx(hKey, (LPCWSTR) L"AppPath", NULL, NULL, (LPBYTE) &szBuffer, &dwBufferSize)){
			RegCloseKey(hKey);
			regValue = std::wstring(szBuffer);
			std::wstring termCheck;
			for(auto val: regValue){
				if(val != L'\0'){
					termCheck += val;
				}
			}
			if(termCheck.size() == 512){
				regValue = L""; // RegValue wasn't null terminated
			}
			else{
				int PathEnd = regValue.length();                                       // Grab the length.
				if(PathEnd > 0) { regValue.erase(PathEnd-1,1); }                       // Eat the stroke at the end.
				PathEnd = regValue.find_last_of(L"\\");                                 // Find the next stroke in.
				int PathLength = regValue.length();                                    // Grab the path length.
				int EraseSpan = PathLength - (PathEnd+1);                               // Calculate the amount to erase.
				regValue.erase(PathEnd+1,EraseSpan);                                   // Erase it to get the root path.
				regValue.append(L"SNF\\snfmdplugin.xml");
			}
		}
	}

    //convert wstring to utf8
    size_t buffer_size = WideCharToMultiByte(CP_UTF8, 0, &regValue[0], (int)regValue.size(), NULL, 0, NULL, NULL);
    std::string strFilePath(buffer_size, '\0');
    WideCharToMultiByte(CP_UTF8, 0, &regValue[0], (int)regValue.size(), &strFilePath[0], buffer_size, NULL, NULL);


    // Return either the empty string or the configuration file path.
    return strFilePath;                                                                    // Return what we have.
}

// Here are out engine components, startup, and shutdown logic. The actual
// engine components are built as the DLL is loaded. They are "lit up" by the
// startup() function - which is called when an application attaches to this
// DLL; and they are closed down when an application detaches from this DLL.

// 20080311 _M Added configuration interpreter for <platform/> and linked it
// to the ConfigFunc() and MessageIPFunc() functions.

volatile bool EngineIsGood = false;                                             // True when things are ready to go.
snf_RulebaseHandler* Rulebase = 0;                                              // Our Rulebase Handler.
snf_EngineHandler* ScanEngine = 0;                                              // Our Scan Engine.
int Generation = -1;                                                            // Configuration generation tag.
string Configuration;                                                           // Configuration file path.
MDConfiguration* PlatformConfig = 0;                                            // Platform config interpreter.

                                   
extern "C" __declspec(dllexport)void _stdcall StartupFunc(HWND Parent){                              // How to get all this started.
    LatestParent = Parent;
    EngineIsGood = false;                                                       // Start off pessimistically.
    try {                                                                       // Be sure to catch any exceptions.
        Rulebase =  new snf_RulebaseHandler();                                  // Create our rulebase handler.
        ScanEngine = new snf_EngineHandler();                                   // Create our engine scanner.
        Configuration = ConfigurationFilePath();                                // Get the configuration path.
        PlatformConfig = new MDConfiguration(*Rulebase, Configuration);         // Set up our configuration monitor.
        Rulebase->open(Configuration.c_str(), "", "");                          // Open a configured rulebase.
        Rulebase->PlatformVersion(PLUGIN_VERSION_INFO);                         // Set the Platform version string.
        ScanEngine->open(Rulebase);                                             // Open the scanning engine.
        EngineIsGood = true;                                                    // If all went well, we're up!
        sayToLog(Rulebase->EngineVersion());                                    // version info to show we're ok.
        sayToLog(Rulebase->PlatformVersion());                                  // Log our platform and engine
        string ConfigInfo = "SNF Config: ";                                     // Build a configuration info
        ConfigInfo.append(Configuration);                                       // message and display that
        sayToLog(ConfigInfo);                                                   // too.
    }
    catch(snf_RulebaseHandler::ConfigurationError) {                            // Can't work with config file!
        string ErrorMessage = "Unable to Configure with: ";                     // Tell them about it and give
        ErrorMessage.append(Configuration);                                     // them a hint where we looked.
        sayToScreen(ErrorMessage);
    }
    catch(snf_RulebaseHandler::FileError) {                                     // Can't load the rulebase file!
        string ErrorMessage = "Unable to Load Rulebase in: ";                   // Tell them about it and give
        ErrorMessage.append(Configuration);                                     // them a hint where we looked.
        sayToScreen(ErrorMessage);
    }
    catch(snf_RulebaseHandler::AllocationError) {                               // Can't allocate memory!
        sayToScreen("Unable to Allocate Enough Memory!");                       // Tell them about it.
    }
    catch(snf_RulebaseHandler::IgnoreListError) {                               // Can't load ignore list!
        sayToScreen("Unable to Load Ingore List!");                             // Tell them about it.
    }
    catch(snf_RulebaseHandler::AuthenticationError) {                           // Can't authenticate!
        sayToScreen("Unable to Authenticate Rulebase!");                        // Tell them about it.
    }
    catch(snf_RulebaseHandler::Busy) {                                          // Busy?!! That's weird.
        sayToScreen("Busy Exception?!!");                                       // Tell them about it.
    }
    catch(snf_RulebaseHandler::Panic) {                                         // Panic?!! That's weird.
        sayToScreen("Panic Exception?!!");                                      // Tell them about it.
    }
    catch(exception& e) {                                                       // Some other runtime error?
        sayToScreen(e.what());                                                  // Tell them what it was.
    }
    catch(...) {                                                                // Catch the unknown exception.
        sayToScreen("Unexpected Exception!");                                   // Tell them things didn't work.
    }

    return;                                                                     // All done.
}

/****** THE FOLLOWING IS DEFUNCT, BUT USEFUL INFO SO WE KEEP IT! ******/
/**** *Note: Why do we always return true even when we fail???
***** Because, in a DLL that is being loaded, none of the threads that
***** get created in the rulebase object will get any cycles until the
***** DLL load sequence is complete _AND_SUCCESSFUL_!
*****
***** The same is true when we return false -- because then the OS KNOWS
***** that it's not safe to run any of the threads we've created.
*****
***** As a result, if we try to destroy our rulebase manager after an
***** unsuccessful load then all of the stop() calls to our threads will
***** block waiting for the initialized threads to see their stop bits and
***** end. Since none of these threads actually get any love until the
***** DLL is successfully loaded, we end up waiting happily for ever!
*****
***** Instead of stepping into this world of hurt, we've decided to leave
***** the internal flag set to indicate that the engine is Not-Good so that
***** the scanning functions remain inert, and then we wait patiently for
***** the DLL to be unloaded.
*****
***** When that happens, since the OS KNOWS the DLL was loaded happily
***** before, the objects can go through their normal destruction without
***** deadlocking on their stop() functions (join()).
****/


extern "C" __declspec(dllexport)void _stdcall ShutdownFunc(HWND Parent){                             // How to get all this stopped.
    LatestParent = Parent;
    if(EngineIsGood) {                                                          // If the engine is good:
        try {
            EngineIsGood = false;                                               // Stop using the engine now.
            LatestParent = 0;                                                   // No more logging calls.
            if(ScanEngine) {                                                    // Close the scanner if it exists:
                ScanEngine->close();                                            // Close it.
                delete ScanEngine;                                              // Delete it.
                ScanEngine = 0;                                                 // Forget it.
            }
            if(Rulebase) {                                                      // Close the rulebase if it exists.
                Rulebase->close();                                              // Close it.
                delete Rulebase;                                                // Delete it.
                Rulebase = 0;                                                   // Forget it.
            }
            if(PlatformConfig) {                                                // Drop the PlatformConfig if it be.
                delete PlatformConfig;                                          // Delete it.
                PlatformConfig = 0;                                             // Forget it.
            }
            Generation = -1;                                                    // Reset our config gen tag.
            sayToLog("SNF: Plugin Shutdown.");                                  // Tell them we're shut down.
        }
        catch(exception& e) {                                                   // If we have a runtime exception
            string InShutdown = "SNF, Shutdown: ";                              // make a human friendly message
            InShutdown.append(e.what());                                        // and package the exception for
            sayToScreen(InShutdown);                                            // a screen pop-up.
        }
        catch(...) {                                                            // If we see some other kind of
            sayToScreen("SNF, Shutdown: Unknown Exception");                    // exception then show that.
        }
    }
    return;                                                                     // All done.
}

// Now we get to the business end of the plugin. Three functions are exported.
// ConfigFunc() launches an editor for the configuration file.

// 20080311 _M Adjusted ConfigFunc to use Platform Configuration Parameters.

//DLL_EXPORT void __stdcall ConfigFunc(HWND Parent) {                              // Configuration function.
extern "C" __declspec(dllexport)void _stdcall ConfigFunc(HWND Parent) {
    // Initialize application.
    wxApp::SetInstance(new dialogapp());
    wxEntryStart(0, NULL);
    wxTheApp->CallOnInit();
    wxTheApp->OnRun();
    wxTheApp->OnExit();
    wxEntryCleanup();
}

// MessageFunc() scans a message file and injects headers.
     
extern "C" __declspec(dllexport)void _stdcall PostMessageFunc(HWND Parent, const char* File){         // Message Scan Function.
    LatestParent = Parent;                                                      // Capture the parent for logging.
    string LogMessage = "SNF MessageScan: ";                                    // Start a plugin log entry.
    LogMessage.append(File);                                                    // add the file name.

    if(false == EngineIsGood) {                                                 // If the engine is not up then
        LogMessage.append(", Engine Not Ready!");                               // report that in the plug-in log.
        sayToLog(LogMessage);                                                   // Post our entry to the plug-in log.
        return;                                                                 // We're done.
    }

    // We are ready to begin our scan.

    int ResultCode = 0;                                                         // Keep track of the scan result.

    try {                                                                       // Be sure to catch any exceptions.
        ResultCode = ScanEngine->scanMessageFile(File);
    }
    catch(snf_EngineHandler::FileError& e) {                                    // Exception when a file won't open.
        LogMessage.append(", File Error!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(snf_EngineHandler::XHDRError& e) {                                    // Exception when XHDR Inject/File fails.
        LogMessage.append(", X-Header Error!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(snf_EngineHandler::BadMatrix& e) {                                    // Exception out of bounds of matrix.
        LogMessage.append(", Bad Matrix!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(snf_EngineHandler::MaxEvals& e) {                                     // Exception too many evaluators.
        LogMessage.append(", Too Many Evaluators!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(snf_EngineHandler::AllocationError& e) {                              // Exception when we can't allocate something.
        LogMessage.append(", Allocation Error!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(snf_EngineHandler::Busy& e) {                                         // Exception when there is a collision.
        LogMessage.append(", Engine Busy!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(snf_EngineHandler::Panic& e) {                                        // Exception when something else happens.
        LogMessage.append(", Engine Panic!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(exception& e) {                                                       // Some unexpected exception.
        LogMessage.append(", Unexpected Exception!");
        sayToLog(LogMessage);
        string DebugData = "SNF Debug: ";
        DebugData.append(e.what());
        sayToLog(DebugData);
        return;
    }
    catch(...) {                                                                // We should never get one of these, but
        LogMessage.append(", Non-Std Exception!");                              // if we do we have something to say.
        sayToLog(LogMessage);
        return;
    }

    // At this point we know our scan worked ok. So, let's post the result.

    ostringstream O;                                                            // A stringstream makes it easy to
    O << LogMessage << ", Result=" << ResultCode;                               // format our plug-in log entry.
    sayToLog(O.str());                                                          // Then we post it to the plug-in log.
}

// ControlFilePath(MessageFilePath) Converts .msg path to .ctl path.

string ControlFilePath(string MessageFilePath) {                                // Convert .msg/.tmp to .ctl path.
    const string ControlFileExt(".ctl");                                        // Control file extension.
    const string EmptyString("");                                               // The empty string.
    string ControlFilePath(MessageFilePath);                                    // Start with the message path.
    string::size_type CFExtPosition = ControlFilePath.find_last_of('.');        // Find the extension.
    if(string::npos == CFExtPosition) return EmptyString;                       // If we didn't find it return ""
    ControlFilePath.replace(                                                    // Replace the message file
      CFExtPosition, ControlFileExt.length(),                                   // extension with the control
      ControlFileExt                                                            // file extension.
    );
    return ControlFilePath;
}

// MessageIPFunc() looks at the control file for a message, extracts the
// source IP, and deletes the message if the IP is in the truncate range.

extern "C" __declspec(dllexport) void _stdcall SMTPMessageFunc(HWND Parent, const char* File){        // IP Scan Function.

    //LatestParent = Parent;                                                      // Capture the parent for logging.
    if(false == PlatformConfig->MessageIPFuncOn()) {                            // If the IP func is off:
       return;                                                                 // We're done.
   }

    string LogMessage = "SNF IPScan: ";                                         // Start a plugin log entry.
    LogMessage.append(File);                                                    // add the file name.
    if(false == EngineIsGood) {                                                 // If the engine is not up then
        LogMessage.append(", Engine Not Ready!");                               // report that in the plug-in log.
        sayToLog(LogMessage);                                                   // Post our entry to the plug-in log.
        return;                                                                 // We're done.
    }
    
    // Open the control file and find the RemoteIP
    const string CtlRemoteIP("RemoteIP=");                                      // This is what we're looking for.
    string RemoteIP = "";                                                       // This is where it will go.
    try {                                                                       // Do this safely.
        ifstream ControlFile(ControlFilePath(File).c_str());                    // Open the control file.
        while(ControlFile) {                                                    // While the file is good...
            string Line;                                                        // Read one line at a time
            getline(ControlFile, Line);                                         // into a string.
            string::size_type TagPos = Line.find(CtlRemoteIP);                  // Look for the RemoteIP tag.
            if(string::npos != TagPos) {                                        // If we find the RemoteIP tag:
                RemoteIP = Line.substr(TagPos + CtlRemoteIP.length());          // Get the string after the tag
                break;                                                          // We're done with the ctl file!
            }
        }
        ControlFile.close();                                                    // Be nice and close our file.
    }
    catch(...) {}                                                               // Eat any exceptions.

    if(0 == RemoteIP.length()) {                                                // If we didn't get the IP then
        LogMessage.append(", Didn't Get Remote IP!");                           // make that the rest of our log
        sayToLog(LogMessage);                                                   // entry and post it.
        return;                                                                 // We're done.
    }

    // At this point we have the remote IP to check.

    LogMessage.append(", "); LogMessage.append(RemoteIP);                       // Add the IP to the log entry.
    IPTestRecord IPAnalysis(RemoteIP);                                          // Set up a test record.

    try {                                                                       // Safely,
        Rulebase->performIPTest(IPAnalysis);                                    // check the IP w/ GBUdb.
    }
    catch(...) {                                                                // If we encountered a problem
        LogMessage.append(", Analysis Failed!");                                // then we need to report that
        sayToLog(LogMessage);                                                   // our analysis failed.
        return;                                                                 // We're done.
    }

    // At this point we've got a good analysis, so we take action (if needed)
    // and produce a helpful log entry.
    ostringstream Happy;                                                        // A stringstream for easy formatting.
    Happy << LogMessage << ", {";                                               // Start the analysis report.

    switch(IPAnalysis.G.Flag()) {                                               // Identify the flag data for this IP.
        case Good: Happy << "Good, "; break;
        case Bad:  Happy << "Bad, "; break;
        case Ugly: Happy << "Ugly, "; break;
        case Ignore: Happy << "Ignore, "; break;
    }

    Happy << "p=" << IPAnalysis.G.Probability() << ", "                         // Probability and Confidence.
          << "c=" << IPAnalysis.G.Confidence() << ", ";

    switch(IPAnalysis.R) {                                                      // The Range analysis.
        case snfIPRange::Unknown:   { Happy << " Unknown}"; break; }                        // Unknown - not defined.
        case snfIPRange::White:     { Happy << " White}"; break; }                          // This is a good guy.
        case snfIPRange::Normal:    { Happy << " Normal}"; break; }                         // Benefit of the doubt.
        case snfIPRange::New:       { Happy << " New}"; break; }                            // It is new to us.
        case snfIPRange::Caution:   { Happy << " Caution}"; break; }                        // This is suspicious.
        case snfIPRange::Black:     { Happy << " Black}"; break; }                          // This is bad.
        case snfIPRange::Truncate:  { Happy << " Truncate}"; break; }                       // This is bad unless we ignore it.
    }

    string WhatWeDid;                                                           // So, what do we do?
    if(                                                                         // Are we in Truncate land?
      (snfIPRange::Truncate == IPAnalysis.R && Ugly == IPAnalysis.G.Flag()) ||              // If the IP is Truncate & Ugly, or
      Bad == IPAnalysis.G.Flag()                                                // If the IP is just plain bad then
      ) {                                                                       // we are in truncate land so we
        Happy << " Rejected!"; WhatWeDid = "Rejected";                          // mark the message Rejected! and
        remove(File);                                                           // remove the file to reject it!
    } else {                                                                    // If we are not rejecting the
        Happy << " Allowed."; WhatWeDid = "Allowed";                            // message tell them it's allowed.
    }

    Rulebase->logThisIPTest(IPAnalysis, WhatWeDid);                             // Log the event.

    sayToLog(Happy.str());                                                      // Post to the plug-in log.
}
