#include "main.h"
#include "modules/ModuleBase.hpp"
#include "modules/LibInherit.hpp"
#include "modules/AdminImpersonate.hpp"
#include "modules/RemoteLockBreak.hpp"
#include "modules/PolicyDisable.hpp"

#include <fstream>
#include <string>
#include <string_view>
#include <vector>
#include <memory>
#include <locale>
#include <codecvt>
#include <windows.h>
#include <winbase.h>
#include <detours.h>

static std::vector<std::unique_ptr<ModuleBase>> modules;
static std::ofstream *log_out;
char sDetourLibrary[512];


#ifdef ENABLE_LOGGING
void log_str(std::string_view str) {
    *log_out << str;
}
void log_str(std::wstring_view str) {
    using convert_typeX = std::codecvt_utf8<wchar_t>;
    std::wstring_convert<convert_typeX, wchar_t> converterX;
    log_str(converterX.to_bytes(std::wstring(str)));
}
void log_endl() {
    *log_out << "\r\n";
}
#endif

void loadModules() {
    char buffer[1000];
    modules.emplace_back(std::make_unique<LibInherit>());
    if (GetEnvironmentVariable("__POLICYTOOL_ADMINIMPERSONATE", buffer, sizeof(buffer)) && buffer[0] == '1') {
        modules.emplace_back(std::make_unique<AdminImpersonate>());
    }
    if (GetEnvironmentVariable("__POLICYTOOL_REMOTELOCKBREAK", buffer, sizeof(buffer)) && buffer[0] == '1') {
        modules.emplace_back(std::make_unique<RemoteLockBreak>());
    }
    if (GetEnvironmentVariable("__POLICYTOOL_POLICYDISABLE", buffer, sizeof(buffer)) && buffer[0] == '1') {
        modules.emplace_back(std::make_unique<PolicyDisable>());
    }
}

EXTERN_C BOOL WINAPI DllMain(HINSTANCE hinst, DWORD dwReason, LPVOID reserved) {
    if (DetourIsHelperProcess()) {
        return TRUE;
    }

    if (dwReason == DLL_PROCESS_ATTACH) {
        log_out = new std::ofstream("C:\\PolicyTool\\log"+std::to_string(getpid())+".txt", std::ios_base::out | std::ios_base::app | std::ios_base::binary);
        DetourRestoreAfterWith();
        GetModuleFileNameA(hinst, sDetourLibrary, ARRAYSIZE(sDetourLibrary));
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        loadModules();
        DetourTransactionCommit();
    } else if (dwReason == DLL_PROCESS_DETACH) {
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        modules.clear();
        DetourTransactionCommit();
        delete log_out;
    }

    return TRUE;
}