#include <Windows.h>
#include <TlHelp32.h>
#include <Psapi.h>
#include <iostream>
#include <vector>
#include <cstdlib>

//-----------------------------------------------------------------------------
// Helper: Resolve a function’s address inside kernel32.dll of a remote process
//-----------------------------------------------------------------------------

uintptr_t GetFunctionAddressInModule(HANDLE hProcess,
                                     const char* moduleName,
                                     const char* functionName)
{
    HMODULE hMods[1024];
    DWORD cbNeeded;
    if (!EnumProcessModules(hProcess, hMods, sizeof(hMods), &cbNeeded)) {
        std::cerr << "[!] Failed to enumerate modules.\n";
        return 0;
    }

    for (unsigned int i = 0; i < (cbNeeded / sizeof(HMODULE)); i++) {
        char modName[MAX_PATH];
        if (GetModuleBaseNameA(hProcess, hMods[i], modName, sizeof(modName))) {
            if (_stricmp(modName, moduleName) == 0) {
                MODULEINFO modInfo;
                if (!GetModuleInformation(hProcess, hMods[i],
                                          &modInfo, sizeof(modInfo))) {
                    std::cerr << "[!] Failed to get module information.\n";
                    return 0;
                }

                BYTE* buffer = new BYTE[modInfo.SizeOfImage];
                if (!ReadProcessMemory(hProcess,
                                       modInfo.lpBaseOfDll,
                                       buffer,
                                       modInfo.SizeOfImage,
                                       nullptr)) {
                    std::cerr << "[!] Failed to read module memory.\n";
                    delete[] buffer;
                    return 0;
                }

                PIMAGE_DOS_HEADER dosHeader =
                    reinterpret_cast<PIMAGE_DOS_HEADER>(buffer);
                PIMAGE_NT_HEADERS ntHeaders =
                    reinterpret_cast<PIMAGE_NT_HEADERS>(buffer + dosHeader->e_lfanew);
                PIMAGE_EXPORT_DIRECTORY exportDir =
                    reinterpret_cast<PIMAGE_EXPORT_DIRECTORY>(
                        buffer +
                        ntHeaders->OptionalHeader
                                 .DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]
                                 .VirtualAddress
                    );

                DWORD* funcNames  = reinterpret_cast<DWORD*>(buffer + exportDir->AddressOfNames);
                DWORD* funcAddrs  = reinterpret_cast<DWORD*>(buffer + exportDir->AddressOfFunctions);
                WORD*  ordinals   = reinterpret_cast<WORD*>(buffer + exportDir->AddressOfNameOrdinals);

                for (DWORD j = 0; j < exportDir->NumberOfNames; j++) {
                    const char* currName = reinterpret_cast<const char*>(buffer + funcNames[j]);
                    if (_stricmp(currName, functionName) == 0) {
                        DWORD funcRVA = funcAddrs[ordinals[j]];
                        uintptr_t funcAddr = reinterpret_cast<uintptr_t>(modInfo.lpBaseOfDll) + funcRVA;
                        delete[] buffer;
                        return funcAddr;
                    }
                }

                std::cerr << "[!] Function not found in module.\n";
                delete[] buffer;
                return 0;
            }
        }
    }

    std::cerr << "[!] Module not found in process.\n";
    return 0;
}

//-----------------------------------------------------------------------------
// Entry Point
//-----------------------------------------------------------------------------

int main(int argc, char* argv[])
{
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <PID>\n";
        return 1;
    }

    // Open the target process to read its kernel32.dll
    DWORD pid = static_cast<DWORD>(strtoul(argv[1], nullptr, 10));
    HANDLE hProcess = OpenProcess(
        PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
        FALSE, pid
    );
    if (!hProcess) {
        std::cerr << "[!] Failed to open process. GLE: " << GetLastError() << "\n";
        return 1;
    }

    // Resolve CreateProcessA within the remote process's copy of kernel32.dll
    uintptr_t addr = GetFunctionAddressInModule(
        hProcess,
        "kernel32.dll",
        "CreateProcessA"
    );
    if (!addr) {
        std::cerr << "[!] Could not locate CreateProcessA\n";
        CloseHandle(hProcess);
        return 1;
    }

    // Cast to proper function signature
    using CreateProc_t = BOOL (WINAPI*)(
        LPCSTR, LPSTR,
        LPSECURITY_ATTRIBUTES, LPSECURITY_ATTRIBUTES,
        BOOL, DWORD, LPVOID, LPCSTR,
        LPSTARTUPINFOA, LPPROCESS_INFORMATION
    );
    auto CreateProcessA_unhooked = reinterpret_cast<CreateProc_t>(addr);

    // Build an in-memory batch command line (no disk I/O)
    std::vector<const char*> commands = {
        "echo Hello from in-memory execution!",
        "whoami",
        "dir C:\\Windows\\System32"
    };
    std::string cmdLine = "cmd.exe /C \"";
    for (size_t i = 0; i < commands.size(); ++i) {
        cmdLine += commands[i];
        if (i + 1 < commands.size()) {
            cmdLine += " & ";
        }
    }
    cmdLine += "\"";

    // Prepare STARTUPINFO to launch cmd.exe in its own console
    STARTUPINFOA si{};
    si.cb = sizeof(si);
    PROCESS_INFORMATION pi{};

    // Invoke cmd.exe via the unhooked CreateProcessA pointer
    BOOL ok = CreateProcessA_unhooked(
        nullptr,                     // lpApplicationName
        (LPSTR)cmdLine.c_str(),      // lpCommandLine
        nullptr, nullptr,            // Security attributes
        FALSE,                       // bInheritHandles
        CREATE_NEW_CONSOLE,          // dwCreationFlags
        nullptr, nullptr,            // Environment, Current directory
        &si, &pi                     // Startup and process info
    );
    if (!ok) {
        std::cerr << "[!] Failed to spawn cmd.exe. GLE: "
                  << GetLastError() << "\n";
        CloseHandle(hProcess);
        return 1;
    }

    std::cout << "[+] Launched in-memory batch via cmd.exe:\n    "
              << cmdLine << "\n";

    // Wait for completion and clean up
    WaitForSingleObject(pi.hProcess, INFINITE);
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);
    CloseHandle(hProcess);

    std::cout << "[+] Done.\n";
    return 0;
}