#include <iostream>
#include <Windows.h>
#include <string>
#include <WDBGEXTS.H>

#include "common.h"

#include "threadpool.h"
#include "client.h"

#include "../user/um/umanager.h"
#include "../user/km/kmanager.h"

DWORD WINAPI Init(HINSTANCE hinstDLL)
{
        AllocConsole();
        FILE* file;
        freopen_s(&file, "CONOUT$", "w", stdout);
        freopen_s(&file, "CONIN$", "r", stdin);

        std::this_thread::sleep_for(std::chrono::seconds(1));

        LPTSTR pipe_name = (LPTSTR)L"\\\\.\\pipe\\DonnaACPipe";
        LPCWSTR driver_name = L"\\\\.\\DonnaAC";

        std::shared_ptr<global::ThreadPool> thread_pool = std::make_shared<global::ThreadPool>(4);
        std::shared_ptr<global::Client> client_interface = std::make_shared<global::Client>(thread_pool, pipe_name);

        usermode::UManager umanager(thread_pool, client_interface);
        kernelmode::KManager kmanager(driver_name, thread_pool, client_interface);

        global::headers::SYSTEM_INFORMATION system_information;
        kmanager.SendClientHardwareInformation();

        global::report_structures::SYSTEM_INFORMATION_REQUEST_RESPONSE response;

        //client_interface->ServerReceive( &response, sizeof( response ) );

        //std::cout << "RequestID: " << response.RequestId << " CanUserProceed: " << 
        //    response.CanUserProceed << " Reason: " << response.reason << std::endl;

        srand(time(NULL));

        while (!GetAsyncKeyState(VK_DELETE))
        {
                int seed = (rand() % 10);

                std::cout << "Seed: " << seed << std::endl;

                switch (seed)
                {
                case 0:
                        kmanager.EnumerateHandleTables();
                        break;
                case 1:
                        kmanager.PerformIntegrityCheck();
                        break;
                case 2:
                        kmanager.ScanPoolsForUnlinkedProcesses();
                        break;
                case 3:
                        kmanager.VerifySystemModules();
                        break;
                case 4:
                        kmanager.ValidateProcessModules();
                        break;
                case 5:
                        kmanager.RunNmiCallbacks();
                        break;
                case 6:
                        kmanager.CheckForAttachedThreads();
                        break;
                case 7:
                        kmanager.InitiateApcStackwalkOperation();
                        break;
                case 8:
                        kmanager.CheckForHiddenThreads();
                        break;
                case 9:
                        kmanager.CheckForEptHooks();
                        break;
                }

                kmanager.MonitorCallbackReports();
                std::this_thread::sleep_for(std::chrono::seconds(10));
        }

        fclose(stdout);
        fclose(stdin);
        FreeConsole();

        FreeLibraryAndExitThread(hinstDLL, 0);
        return 0;
}

BOOL WINAPI DllMain(
        HINSTANCE hinstDLL,  // handle to DLL module
        DWORD fdwReason,     // reason for calling function
        LPVOID lpvReserved)  // reserved
{
        // Perform actions based on the reason for calling.
        switch (fdwReason)
        {
        case DLL_PROCESS_ATTACH:

                DisableThreadLibraryCalls(hinstDLL);

                const auto thread = CreateThread(
                        nullptr,
                        0,
                        reinterpret_cast<LPTHREAD_START_ROUTINE>(Init),
                        hinstDLL,
                        0,
                        nullptr
                );

                if (thread)
                        CloseHandle(thread);

                break;
        }
        return TRUE;  // Successful DLL_PROCESS_ATTACH.
}