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

// Определения для PE-структур
typedef ULONGLONG QWORD;
#define IMAGE_REL_BASED_DIR64 10

// Структура для записей релокации
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;
};

// Логирование
void Log(const std::string& message) {
    std::ofstream log("injector.log", std::ios::app);
    if (log) {
        log << "[" << __TIME__ << "] " << message << std::endl;
    }
}

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);
    }
}

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++;
    }
    return true;
}

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;
        }
    }

    WriteProcessMemory(hProcess, remoteBase, dllData.data(), pNtHeaders->OptionalHeader.SizeOfHeaders, NULL);

    auto pSection = IMAGE_FIRST_SECTION(pNtHeaders);
    for (int i = 0; i < pNtHeaders->FileHeader.NumberOfSections; i++, pSection++) {
        PVOID secDest = (BYTE*)remoteBase + pSection->VirtualAddress;
        WriteProcessMemory(hProcess, secDest, dllData.data() + pSection->PointerToRawData, pSection->SizeOfRawData, NULL);
    }

    ULONGLONG delta = (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;
    }

    MANUAL_MAPPING_DATA mappingData{};
    mappingData.ImageBase = remoteBase;
    mappingData.EntryPoint = (PVOID)((ULONGLONG)remoteBase + pNtHeaders->OptionalHeader.AddressOfEntryPoint);

    HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0,
        (LPTHREAD_START_ROUTINE)mappingData.EntryPoint,
        mappingData.ImageBase, 0, NULL);

    if (!hThread) {
        Log("Thread creation failed: " + std::to_string(GetLastError()));
        VirtualFreeEx(hProcess, remoteBase, 0, MEM_RELEASE);
        return false;
    }

    WaitForSingleObject(hThread, INFINITE);
    CloseHandle(hThread);
    Log("Manual mapping completed");
    return true;
}

int main() {
    const std::string dllPath = "C:\\Downloads\\d3d11hook_[unknowncheats.me]_.dll";
    const std::wstring targetProcess = L"RainbowSix_DX11.exe";

    DWORD pid = GetProcessIdByName(targetProcess);
    if (!pid) {
        Log("Process not found");
        return 1;
    }

    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
    if (!hProcess) {
        Log("OpenProcess failed: " + std::to_string(GetLastError()));
        return 1;
    }

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

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