#include <windows.h>
#include <tlhelp32.h>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <psapi.h>
#include <processthreadsapi.h>

// Определения структур и констант
#define IMAGE_REL_BASED_DIR64 10
typedef ULONGLONG QWORD;
typedef NTSTATUS(NTAPI* NtCreateThreadEx_t)(PHANDLE, ACCESS_MASK, PVOID, HANDLE, PVOID, PVOID, ULONG, SIZE_T, SIZE_T, SIZE_T, PVOID);

typedef struct _BASE_RELOCATION_ENTRY {
    USHORT Offset : 12;
    USHORT Type : 4;
} BASE_RELOCATION_ENTRY, *PBASE_RELOCATION_ENTRY;

struct MANUAL_MAPPING_DATA {
    PVOID ImageBase;
    PVOID EntryPoint;
};

// Глобальные настройки
const std::wstring TARGET_PROCESS = L"RainbowSix_DX11.exe";
const std::string DLL_PATH = "C:\\Downloads\\your_dll.dll";
const std::string LOG_FILE = "C:\\Windows\\Temp\\injector.log";

// Функция логирования
void Log(const std::string& message) {
    std::ofstream log(LOG_FILE, std::ios::app);
    if (log) {
        log << "[" << __TIME__ << "] " << message << std::endl;
    }
}

// Проверка прав администратора
bool IsElevated() {
    HANDLE hToken;
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
        return false;

    TOKEN_ELEVATION elevation;
    DWORD cbSize = sizeof(TOKEN_ELEVATION);
    bool result = GetTokenInformation(hToken, TokenElevation, &elevation, sizeof(elevation), &cbSize);
    CloseHandle(hToken);
    return result && elevation.TokenIsElevated;
}

DWORD GetProcessIdByName(const std::wstring& processName) {
    DWORD processId = 0;
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (snapshot == INVALID_HANDLE_VALUE) {
        Log("Snapshot error: " + std::to_string(GetLastError()));
        return 0;
    }

    PROCESSENTRY32W entry = { sizeof(PROCESSENTRY32W) };
    if (Process32FirstW(snapshot, &entry)) {
        do {
            if (processName == entry.szExeFile) {
                processId = entry.th32ProcessID;
                break;
            }
        } while (Process32NextW(snapshot, &entry));
    }

    CloseHandle(snapshot);
    return processId;
}

void ApplyRelocations(PIMAGE_NT_HEADERS64 pNtHeaders, PVOID base, ULONGLONG delta) {
    if (delta == 0) return;

    auto dir = &pNtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
    if (dir->Size == 0) return;

    auto reloc = (PIMAGE_BASE_RELOCATION)((ULONGLONG)base + dir->VirtualAddress);
    while (reloc->VirtualAddress > 0) {
        ULONGLONG relocAddr = (ULONGLONG)base + reloc->VirtualAddress;
        UINT numEntries = (reloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
        auto entries = (PBASE_RELOCATION_ENTRY)(reloc + 1);

        for (UINT i = 0; i < numEntries; i++) {
            if (entries[i].Type == IMAGE_REL_BASED_DIR64) {
                ULONGLONG* patchAddr = (ULONGLONG*)(relocAddr + entries[i].Offset);
                *patchAddr += delta;
            }
        }
        reloc = (PIMAGE_BASE_RELOCATION)((ULONGLONG)reloc + reloc->SizeOfBlock);
    }
    Log("Relocations applied with delta: 0x" + std::to_string(delta));
}

bool ResolveImports(HANDLE hProcess, PIMAGE_NT_HEADERS64 pNtHeaders, PVOID base) {
    auto dir = &pNtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (dir->Size == 0) return true;

    auto importDesc = (PIMAGE_IMPORT_DESCRIPTOR)((ULONGLONG)base + dir->VirtualAddress);
    while (importDesc->Name) {
        char* moduleName = (char*)((ULONGLONG)base + importDesc->Name);
        HMODULE hModule = GetModuleHandleA(moduleName);
        if (!hModule) {
            hModule = LoadLibraryA(moduleName);
            if (!hModule) return false;
        }

        auto thunk = (PIMAGE_THUNK_DATA64)((ULONGLONG)base + importDesc->FirstThunk);
        while (thunk->u1.AddressOfData) {
            if (IMAGE_SNAP_BY_ORDINAL64(thunk->u1.Ordinal)) {
                auto proc = GetProcAddress(hModule, (LPCSTR)IMAGE_ORDINAL64(thunk->u1.Ordinal));
                if (!proc) return false;
                thunk->u1.Function = (ULONGLONG)proc;
            }
            else {
                auto import = (PIMAGE_IMPORT_BY_NAME)((ULONGLONG)base + thunk->u1.AddressOfData);
                auto proc = GetProcAddress(hModule, import->Name);
                if (!proc) return false;
                thunk->u1.Function = (ULONGLONG)proc;
            }
            thunk++;
        }
        importDesc++;
    }
    Log("Imports resolved successfully");
    return true;
}

HANDLE OpenTargetProcess(DWORD pid) {
    HANDLE hProcess = OpenProcess(
        PROCESS_VM_OPERATION |
        PROCESS_VM_WRITE |
        PROCESS_CREATE_THREAD |
        PROCESS_QUERY_INFORMATION,
        FALSE,
        pid
    );

    if (!hProcess) {
        Log("OpenProcess failed: " + std::to_string(GetLastError()));
        return NULL;
    }

    BOOL isTargetWow64 = FALSE;
    IsWow64Process(hProcess, &isTargetWow64);
    if (isTargetWow64) {
        Log("Target process is 32-bit");
        CloseHandle(hProcess);
        return NULL;
    }

    return hProcess;
}

bool ManualMap(HANDLE hProcess, const std::string& dllPath) {
    Log("Starting manual mapping...");

    HANDLE hFile = CreateFileA(dllPath.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        Log("CreateFile failed: " + std::to_string(GetLastError()));
        return false;
    }

    DWORD fileSize = GetFileSize(hFile, NULL);
    std::vector<BYTE> dllData(fileSize);
    DWORD bytesRead;

    if (!ReadFile(hFile, dllData.data(), fileSize, &bytesRead, NULL)) {
        Log("ReadFile failed: " + std::to_string(GetLastError()));
        CloseHandle(hFile);
        return false;
    }
    CloseHandle(hFile);

    auto pDosHeader = (PIMAGE_DOS_HEADER)dllData.data();
    if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
        Log("Invalid DOS header");
        return false;
    }

    auto pNtHeaders = (PIMAGE_NT_HEADERS64)((ULONGLONG)dllData.data() + pDosHeader->e_lfanew);
    if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE ||
        pNtHeaders->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
        Log("Invalid NT header");
        return false;
    }

    PVOID remoteBase = VirtualAllocEx(hProcess,
        (PVOID)pNtHeaders->OptionalHeader.ImageBase,
        pNtHeaders->OptionalHeader.SizeOfImage,
        MEM_COMMIT | MEM_RESERVE,
        PAGE_EXECUTE_READWRITE);

    if (!remoteBase) {
        remoteBase = VirtualAllocEx(hProcess,
            NULL,
            pNtHeaders->OptionalHeader.SizeOfImage,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_EXECUTE_READWRITE);

        if (!remoteBase) {
            Log("Memory allocation failed: " + std::to_string(GetLastError()));
            return false;
        }
    }

    Log("Memory allocated at: 0x" + std::to_string(reinterpret_cast<ULONGLONG>(remoteBase)));

    // Запись PE-заголовков
    if (!WriteProcessMemory(hProcess, remoteBase, dllData.data(),
                          pNtHeaders->OptionalHeader.SizeOfHeaders, NULL)) {
        Log("Failed to write headers");
        VirtualFreeEx(hProcess, remoteBase, 0, MEM_RELEASE);
        return false;
    }

    // Запись секций
    auto pSection = IMAGE_FIRST_SECTION(pNtHeaders);
    for (WORD i = 0; i < pNtHeaders->FileHeader.NumberOfSections; ++i, ++pSection) {
        PVOID secDest = reinterpret_cast<BYTE*>(remoteBase) + pSection->VirtualAddress;
        if (!WriteProcessMemory(hProcess, secDest,
                              dllData.data() + pSection->PointerToRawData,
                              pSection->SizeOfRawData, NULL)) {
            Log("Failed to write section: " + std::string(reinterpret_cast<char*>(pSection->Name)));
            VirtualFreeEx(hProcess, remoteBase, 0, MEM_RELEASE);
            return false;
        }
        Log("Section " + std::string(reinterpret_cast<char*>(pSection->Name)) +
            " written at: 0x" + std::to_string((ULONGLONG)secDest));
    }

    // Релокации
    ULONGLONG delta = reinterpret_cast<ULONGLONG>(remoteBase) - pNtHeaders->OptionalHeader.ImageBase;
    ApplyRelocations(pNtHeaders, remoteBase, delta);

    // Разрешение импортов
    if (!ResolveImports(hProcess, pNtHeaders, remoteBase)) {
        Log("Import resolution failed");
        VirtualFreeEx(hProcess, remoteBase, 0, MEM_RELEASE);
        return false;
    }

    // Установка защиты памяти
    pSection = IMAGE_FIRST_SECTION(pNtHeaders);
    for (WORD i = 0; i < pNtHeaders->FileHeader.NumberOfSections; ++i, ++pSection) {
        DWORD oldProtect;
        DWORD protect = 0;

        if (pSection->Characteristics & IMAGE_SCN_MEM_EXECUTE) {
            protect = (pSection->Characteristics & IMAGE_SCN_MEM_WRITE)
                    ? PAGE_EXECUTE_READWRITE
                    : PAGE_EXECUTE_READ;
        }
        else if (pSection->Characteristics & IMAGE_SCN_MEM_WRITE) {
            protect = PAGE_READWRITE;
        }
        else {
            protect = PAGE_READONLY;
        }

        VirtualProtectEx(hProcess,
                       reinterpret_cast<BYTE*>(remoteBase) + pSection->VirtualAddress,
                       pSection->Misc.VirtualSize,
                       protect,
                       &oldProtect);
    }

    // Обход античита
    ULONG_PTR policy = PROCESS_CREATION_MITIGATION_POLICY_BLOCK_NON_MICROSOFT_BINARIES_ALWAYS_ON;
    SetProcessMitigationPolicy(ProcessSignaturePolicy, &policy, sizeof(policy));

    // Создание потока через NtCreateThreadEx
    MANUAL_MAPPING_DATA mappingData = {
        remoteBase,
        reinterpret_cast<PVOID>(reinterpret_cast<ULONGLONG>(remoteBase) +
                              pNtHeaders->OptionalHeader.AddressOfEntryPoint)
    };

    auto NtCreateThreadEx = (NtCreateThreadEx_t)GetProcAddress(
        GetModuleHandleA("ntdll.dll"), "NtCreateThreadEx");

    HANDLE hThread;
    NTSTATUS status = NtCreateThreadEx(&hThread,
        0x1FFFFF,
        NULL,
        hProcess,
        mappingData.EntryPoint,
        mappingData.ImageBase,
        0,
        0,
        0,
        0,
        NULL);

    if (status != 0) {
        Log("NtCreateThreadEx failed with status: 0x" + std::to_string(status));
        VirtualFreeEx(hProcess, remoteBase, 0, MEM_RELEASE);
        return false;
    }

    Log("Creating thread at: 0x" + std::to_string((ULONGLONG)mappingData.EntryPoint));

    // Обход DEP/CFG
    DWORD oldProtect;
    VirtualProtectEx(hProcess, mappingData.EntryPoint,
        sizeof(mappingData.EntryPoint), PAGE_EXECUTE_READWRITE, &oldProtect);

    WaitForSingleObject(hThread, INFINITE);
    CloseHandle(hThread);

    Log("Manual mapping completed successfully");
    return true;
}

int main() {
    if (!IsElevated()) {
        MessageBoxA(NULL, "Run as administrator!", "Error", MB_ICONERROR);
        return 1;
    }

    // Вариант 1: Запуск через инжектор
    STARTUPINFOW si = { sizeof(si) };
    PROCESS_INFORMATION pi;
    if (CreateProcessW(
        L"D:\\Games\\RainbowSix_DX11.exe",
        NULL, NULL, NULL, FALSE,
        CREATE_SUSPENDED,
        NULL, NULL, &si, &pi))
    {
        Log("Process created in suspended state");
        WaitForInputIdle(pi.hProcess, 5000);

        if (ManualMap(pi.hProcess, DLL_PATH)) {
            Log("Resuming main thread");
            ResumeThread(pi.hThread);
            CloseHandle(pi.hThread);
            CloseHandle(pi.hProcess);
            return 0;
        }
    }

    // Вариант 2: Инжекция в уже запущенный процесс
    DWORD pid = GetProcessIdByName(TARGET_PROCESS);
    if (!pid) {
        Log("Process not found");
        return 1;
    }

    HANDLE hProcess = OpenTargetProcess(pid);
    if (!hProcess) {
        return 1;
    }

    if (!ManualMap(hProcess, DLL_PATH)) {
        Log("Injection failed");
        CloseHandle(hProcess);
        return 1;
    }

    CloseHandle(hProcess);
    Log("Injection succeeded");
    return 0;
}