/*
 * encrypt_and_test.c - .pdata encryption test with symbol resolution
 *
 * Build (x64 Native Tools Command Prompt for VS):
 *   cl /Zi /Od /W3 /D_CRT_SECURE_NO_WARNINGS encrypt_and_test.c /link /SUBSYSTEM:CONSOLE dbghelp.lib
 *
 * Notes:
 *   /Zi  - generate PDB (required for symbol resolution)
 *   /Od  - disable optimization (prevents Level* functions from being inlined)
 *   /D_CRT_SECURE_NO_WARNINGS - suppress _snprintf_s warnings
 */

#ifndef _WIN64
#error "Must be compiled for x64. Use 'x64 Native Tools Command Prompt for VS'."
#endif

#include <windows.h>
#include <winnt.h>
#include <dbghelp.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#pragma comment(lib, "dbghelp.lib")

#define ENCRYPT_KEY 0x5A

/* ---------- helpers ---------- */

void PrintSeparator(char c, int len) {
    for (int i = 0; i < len; i++) putchar(c);
    putchar('\n');
}

BOOL GetModuleNameForAddr(ULONG64 addr, char* buf, SIZE_T bufSize) {
    HMODULE hMod;
    if (GetModuleHandleExA(
            GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
            GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
            (LPCSTR)addr, &hMod)) {
        if (GetModuleFileNameA(hMod, buf, (DWORD)bufSize)) {
            char* p = strrchr(buf, '\\');
            if (p) memmove(buf, p + 1, strlen(p));
            return TRUE;
        }
    }
    return FALSE;
}

/* ---------- symbol resolution ---------- */

/*
 * ResolveSymbol
 *   Resolves a RIP address to "FunctionName+0xOFFSET [file:line]".
 *   Requires SymInitialize() called at startup and a PDB on disk.
 *   For kernel32/ntdll names set _NT_SYMBOL_PATH to MS symbol server.
 */
void ResolveSymbol(ULONG64 addr, char* outBuf, SIZE_T bufSize) {
    char symBuffer[sizeof(SYMBOL_INFO) + MAX_SYM_NAME * sizeof(char)];
    memset(symBuffer, 0, sizeof(symBuffer));
    SYMBOL_INFO* sym = (SYMBOL_INFO*)symBuffer;
    sym->SizeOfStruct = sizeof(SYMBOL_INFO);
    sym->MaxNameLen   = MAX_SYM_NAME;

    DWORD64 displacement = 0;

    if (SymFromAddr(GetCurrentProcess(), addr, &displacement, sym)) {
        IMAGEHLP_LINE64 line;
        memset(&line, 0, sizeof(line));
        line.SizeOfStruct = sizeof(IMAGEHLP_LINE64);
        DWORD lineDisp = 0;

        if (SymGetLineFromAddr64(GetCurrentProcess(), addr, &lineDisp, &line)) {
            char* fileName = strrchr(line.FileName, '\\');
            fileName = fileName ? fileName + 1 : line.FileName;
            _snprintf(outBuf, bufSize, "%s+0x%llX [%s:%lu]",
                      sym->Name, displacement, fileName, line.LineNumber);
        } else {
            _snprintf(outBuf, bufSize, "%s+0x%llX", sym->Name, displacement);
        }
    } else {
        _snprintf(outBuf, bufSize, "<no symbol> (err=%lu)", GetLastError());
    }
}

/* ---------- EDR stack walk simulation ---------- */

void EDR_StackWalk(CONTEXT* ctx, int maxFrames) {
    PrintSeparator('=', 90);
    printf("  EDR Stack Walk Simulation (with symbol resolution)\n");
    PrintSeparator('=', 90);

    printf("  %-3s  %-18s  %-18s  %-20s  %-12s  %s\n",
           "Frm", "RIP", "RSP", "Module", "Status", "Symbol");
    PrintSeparator('-', 90);

    CONTEXT c = *ctx;
    int frameNum     = 0;
    int noPdataCount = 0;
    int suspectCount = 0;

    printf("\n  [Start] RIP=0x%016llX  RSP=0x%016llX\n\n", c.Rip, c.Rsp);

    while (frameNum < maxFrames) {

        /* Step 1: look up .pdata entry */
        ULONG64 imageBase = 0;
        UNWIND_HISTORY_TABLE hist;
        memset(&hist, 0, sizeof(hist));
        PRUNTIME_FUNCTION rf = RtlLookupFunctionEntry(c.Rip, &imageBase, &hist);

        /* Step 2: module name + memory type */
        char moduleName[MAX_PATH] = "<unknown>";
        GetModuleNameForAddr(c.Rip, moduleName, MAX_PATH);

        MEMORY_BASIC_INFORMATION mbi;
        memset(&mbi, 0, sizeof(mbi));
        VirtualQuery((LPCVOID)c.Rip, &mbi, sizeof(mbi));
        const char* memType = "UNKNOWN";
        if      (mbi.Type == MEM_IMAGE)   memType = "IMAGE(ok)";
        else if (mbi.Type == MEM_PRIVATE) memType = "PRIVATE(sus!)";
        else if (mbi.Type == MEM_MAPPED)  memType = "MAPPED";

        /* Step 3: threat assessment */
        BOOL hasPdata = (rf != NULL);
        const char* statusStr = "OK";
        char reason[256] = "";

        if (!hasPdata) {
            noPdataCount++;
            statusStr = "!NO_PDATA";
            strcpy(reason, "missing .pdata entry");
            if (mbi.Type == MEM_PRIVATE) {
                statusStr = "!CRITICAL";
                strcpy(reason, "no .pdata + PRIVATE memory (possible shellcode)");
                suspectCount++;
            }
        }

        /* Step 3.5: resolve symbol */
        char symbolInfo[512] = "";
        ResolveSymbol(c.Rip, symbolInfo, sizeof(symbolInfo));

        /* Step 4: print frame */
        printf("  %-3d  0x%016llX  0x%016llX  %-20s  %-12s  %s\n",
               frameNum, c.Rip, c.Rsp, moduleName, statusStr, symbolInfo);

        if (!hasPdata || mbi.Type == MEM_PRIVATE) {
            printf("       +-- reason  : %s\n", reason);
            printf("       +-- mem type: %s\n", memType);
        }

        /* Step 5: unwind to caller */
        if (rf) {
            PVOID   handlerData      = NULL;
            ULONG64 establisherFrame = 0;
            RtlVirtualUnwind(
                UNW_FLAG_NHANDLER,
                imageBase,
                c.Rip,
                rf,
                &c,
                &handlerData,
                &establisherFrame,
                NULL);
        } else {
            printf("       +-- no .pdata, reading return address from stack\n");
            if (c.Rsp == 0) { printf("  [!] RSP is 0, reached stack bottom\n"); break; }
            c.Rip  = *(ULONG64*)c.Rsp;
            c.Rsp += 8;
        }

        if (c.Rip == 0) { printf("\n  [Done] reached stack top\n"); break; }
        frameNum++;
    }

    PrintSeparator('-', 90);

    printf("\n  [Summary]\n");
    printf("  Total frames   : %d\n", frameNum);
    printf("  No-pdata frames: %d\n", noPdataCount);
    printf("  Critical frames: %d\n", suspectCount);

    if (suspectCount > 0)
        printf("\n  [!!!] WARNING: suspicious call stack! Possible shellcode / stack spoofing.\n");
    else if (noPdataCount > 0)
        printf("\n  [??]  NOTICE: some frames lack .pdata.\n");
    else
        printf("\n  [OK]  Call stack looks clean.\n");
    printf("\n");
}

/* ---------- .pdata encrypt / decrypt ---------- */

BOOL EncryptPdataSection(HMODULE hModule) {
    ULONG64 imageBase = (ULONG64)hModule;
    IMAGE_DOS_HEADER* dos = (IMAGE_DOS_HEADER*)imageBase;
    if (dos->e_magic != IMAGE_DOS_SIGNATURE) { printf("[-] Bad DOS header\n"); return FALSE; }

    IMAGE_NT_HEADERS* nt = (IMAGE_NT_HEADERS*)(imageBase + dos->e_lfanew);
    if (nt->Signature != IMAGE_NT_SIGNATURE) { printf("[-] Bad NT header\n"); return FALSE; }

    IMAGE_DATA_DIRECTORY* exDir =
        &nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
    if (!exDir->VirtualAddress || !exDir->Size) {
        printf("[-] No Exception Directory (.pdata) found\n"); return FALSE;
    }

    ULONG64 pdataStart = imageBase + exDir->VirtualAddress;
    ULONG   pdataSize  = exDir->Size;

    printf("[+] .pdata: VA=0x%016llX  size=0x%X bytes (%zu entries)\n",
           pdataStart, pdataSize, (size_t)(pdataSize / sizeof(RUNTIME_FUNCTION)));

    DWORD oldProtect;
    if (!VirtualProtect((LPVOID)pdataStart, pdataSize, PAGE_READWRITE, &oldProtect)) {
        printf("[-] VirtualProtect failed: %lu\n", GetLastError()); return FALSE;
    }
    unsigned char* p = (unsigned char*)pdataStart;
    for (ULONG i = 0; i < pdataSize; i++) p[i] ^= ENCRYPT_KEY;
    VirtualProtect((LPVOID)pdataStart, pdataSize, oldProtect, &oldProtect);

    printf("[+] .pdata encrypted. RtlLookupFunctionEntry will return NULL.\n");
    return TRUE;
}

BOOL DecryptPdataSection(HMODULE hModule) {
    ULONG64 imageBase = (ULONG64)hModule;
    IMAGE_DOS_HEADER* dos = (IMAGE_DOS_HEADER*)imageBase;
    if (dos->e_magic != IMAGE_DOS_SIGNATURE) return FALSE;

    IMAGE_NT_HEADERS* nt = (IMAGE_NT_HEADERS*)(imageBase + dos->e_lfanew);
    if (nt->Signature != IMAGE_NT_SIGNATURE) return FALSE;

    IMAGE_DATA_DIRECTORY* exDir =
        &nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
    if (!exDir->VirtualAddress || !exDir->Size) return FALSE;

    ULONG64 pdataStart = imageBase + exDir->VirtualAddress;
    ULONG   pdataSize  = exDir->Size;

    printf("[+] .pdata: VA=0x%016llX  size=0x%X bytes\n", pdataStart, pdataSize);

    DWORD oldProtect;
    if (!VirtualProtect((LPVOID)pdataStart, pdataSize, PAGE_READWRITE, &oldProtect))
        return FALSE;
    unsigned char* p = (unsigned char*)pdataStart;
    for (ULONG i = 0; i < pdataSize; i++) p[i] ^= ENCRYPT_KEY;
    VirtualProtect((LPVOID)pdataStart, pdataSize, oldProtect, &oldProtect);

    printf("[+] .pdata decrypted. Stack unwinding restored.\n");
    return TRUE;
}

/* ---------- call chain ---------- */

#pragma optimize("", off)

void __declspec(noinline) Level3_DeepFunction(void) {
    printf("\n[Level3] Deepest function - EDR trigger point\n\n");

    CONTEXT ctx;
    memset(&ctx, 0, sizeof(ctx));
    ctx.ContextFlags = CONTEXT_FULL;
    RtlCaptureContext(&ctx);

    printf("[Level3] Context captured: RIP=0x%016llX  RSP=0x%016llX\n\n",
           ctx.Rip, ctx.Rsp);

    EDR_StackWalk(&ctx, 32);
}

void __declspec(noinline) Level2_MiddleFunction(int param) {
    printf("[Level2] Middle, param=%d\n", param);
    Level3_DeepFunction();
}

void __declspec(noinline) Level1_CallerFunction(const char* name) {
    printf("[Level1] Caller, name=%s\n", name);
    Level2_MiddleFunction(42);
}

void __declspec(noinline) Level0_EntryFunction(void) {
    printf("[Level0] Entry\n");
    Level1_CallerFunction("TestCall");
}

#pragma optimize("", on)

/* ---------- main ---------- */

int main(void) {
    HANDLE hProc = GetCurrentProcess();

    SymSetOptions(SYMOPT_UNDNAME | SYMOPT_DEFERRED_LOADS | SYMOPT_LOAD_LINES);
    if (!SymInitialize(hProc, NULL, TRUE)) {
        printf("[!] SymInitialize failed: %lu\n", GetLastError());
    } else {
        printf("[+] Symbol engine ready.\n");
        printf("[!] For kernel32/ntdll names, set before running:\n");
        printf("    set _NT_SYMBOL_PATH=srv*C:\\Symbols*"
               "https://msdl.microsoft.com/download/symbols\n\n");
    }

    PrintSeparator('=', 90);
    printf("  .pdata Encryption Demo - EDR Stack Walk + Symbol Resolution\n");
    PrintSeparator('=', 90);
    printf("\n  PID: %lu\n  Module base: %p\n\n",
           GetCurrentProcessId(), (void*)GetModuleHandleA(NULL));

    HMODULE hModule = GetModuleHandleA(NULL);

    /* Test 1 */
    PrintSeparator('=', 90);
    printf("\n[Test 1] Baseline: .pdata NOT encrypted\n");
    printf("Expect : all OK, symbol column shows Level0~Level3 + RtlUserThreadStart\n");
    PrintSeparator('=', 90);
    Level0_EntryFunction();

    /* Test 2 */
    PrintSeparator('=', 90);
    printf("\n[Test 2] Encrypting .pdata...\n");
    PrintSeparator('=', 90);
    EncryptPdataSection(hModule);

    /* Test 3 */
    PrintSeparator('=', 90);
    printf("\n[Test 3] Stack walk with encrypted .pdata\n");
    printf("Expect : NO_PDATA / CRITICAL, <no symbol>\n");
    PrintSeparator('=', 90);
    Level0_EntryFunction();

    /* Test 4 */
    PrintSeparator('=', 90);
    printf("\n[Test 4] Decrypting .pdata...\n");
    PrintSeparator('=', 90);
    DecryptPdataSection(hModule);

    /* Test 5 */
    PrintSeparator('=', 90);
    printf("\n[Test 5] Stack walk after decryption\n");
    printf("Expect : all OK, symbols restored\n");
    PrintSeparator('=', 90);
    Level0_EntryFunction();

    SymCleanup(hProc);
    printf("\nDone. Press Enter to exit...\n");
    getchar();
    return 0;
}

/*
 * test2.c
 *
 * Safe test: encrypt .pdata / .rdata in a timer window while code (.text)
 * stays untouched. A helper thread runs the call chain during the encrypted
 * window, then the timer restores everything.
 *
 * Tests:
 *   Test 1 - Baseline             : no encryption, all frames OK
 *   Test 2 - Encrypt .pdata only  : RtlLookupFunctionEntry fails -> NO_PDATA
 *   Test 3 - Encrypt .rdata only  : RtlVirtualUnwind reads garbage UNWIND_INFO
 *   Test 4 - Encrypt both         : lookup fails + unwind info gone
 *
 * Build (x64 Native Tools Command Prompt):
 *   cl /Zi /Od /W3 /D_CRT_SECURE_NO_WARNINGS test2.c /link dbghelp.lib
 */

#ifndef _WIN64
#error "x64 only"
#endif

#include <windows.h>
#include <winnt.h>
#include <dbghelp.h>
#include <stdio.h>
#include <string.h>

#pragma comment(lib, "dbghelp.lib")

#define ENCRYPT_KEY    0x5A
#define MAX_UNWIND_REG 256

/* ---- minimal UNWIND_INFO header (4 bytes) ---- */
typedef struct _UNWIND_INFO_MIN {
    BYTE VersionAndFlags;
    BYTE SizeOfProlog;
    BYTE CountOfCodes;
    BYTE FrameRegisterAndOffset;
} UNWIND_INFO_MIN, *PUNWIND_INFO_MIN;

typedef struct _PRESERVE_REGION {
    ULONG_PTR RVA;
    DWORD     Size;
    PBYTE     SavedCopy;
} PRESERVE_REGION;

/* ---- context passed between timers / threads ---- */
typedef struct _TEST_CTX {
    const char* Label;
    ULONG_PTR   Base;
    DWORD       ImageSize;
    ULONG_PTR   PdataVA;
    DWORD       PdataSize;
    ULONG_PTR   RdataVA;
    DWORD       RdataSize;
    PBYTE       SavedPdata;
    PBYTE       SavedRdata;
    PRESERVE_REGION Regions[MAX_UNWIND_REG];
    DWORD       RegionCount;
    HANDLE      StartEvent;   /* signalled after encryption is done */
    HANDLE      DoneEvent;    /* signalled after inspector finishes */
    BOOL        EncryptPdata;
    BOOL        EncryptRdata;
} TEST_CTX;

/* ==================================================================== */

static void PrintSep(char c, int n) {
    for (int i = 0; i < n; i++) putchar(c);
    putchar('\n');
}

static void XorRegion(PVOID addr, DWORD size, BYTE key) {
    PBYTE p = (PBYTE)addr;
    for (DWORD i = 0; i < size; i++) p[i] ^= key;
}

static BOOL GetModName(ULONG64 addr, char* buf, SIZE_T sz) {
    HMODULE h;
    if (!GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
                             GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
                             (LPCSTR)addr, &h)) return FALSE;
    if (!GetModuleFileNameA(h, buf, (DWORD)sz)) return FALSE;
    char* p = strrchr(buf, '\\');
    if (p) memmove(buf, p + 1, strlen(p));
    return TRUE;
}

static void ResolveSym(ULONG64 addr, char* out, SIZE_T sz) {
    char raw[sizeof(SYMBOL_INFO) + MAX_SYM_NAME];
    memset(raw, 0, sizeof(raw));
    SYMBOL_INFO* sym = (SYMBOL_INFO*)raw;
    sym->SizeOfStruct = sizeof(SYMBOL_INFO);
    sym->MaxNameLen   = MAX_SYM_NAME;
    DWORD64 disp = 0;

    if (!SymFromAddr(GetCurrentProcess(), addr, &disp, sym)) {
        _snprintf(out, sz, "<no sym>");
        return;
    }

    IMAGEHLP_LINE64 line;
    DWORD ld = 0;
    memset(&line, 0, sizeof(line));
    line.SizeOfStruct = sizeof(line);

    if (SymGetLineFromAddr64(GetCurrentProcess(), addr, &ld, &line)) {
        char* fn = strrchr(line.FileName, '\\');
        fn = fn ? fn + 1 : line.FileName;
        _snprintf(out, sz, "%s+0x%llX [%s:%lu]", sym->Name, disp, fn, line.LineNumber);
    } else {
        _snprintf(out, sz, "%s+0x%llX", sym->Name, disp);
    }
}

/* ---- UNWIND_INFO size calculation ---- */
static DWORD CalcUnwindSize(PBYTE pImage, DWORD rva) {
    PUNWIND_INFO_MIN p = (PUNWIND_INFO_MIN)(pImage + rva);
    BYTE flags = (p->VersionAndFlags >> 3) & 0x1F;
    BYTE codes = p->CountOfCodes;
    DWORD size = sizeof(UNWIND_INFO_MIN);
    size += codes * sizeof(USHORT);
    if (codes % 2 != 0) size += sizeof(USHORT);
    if (flags & 0x4)       size += sizeof(RUNTIME_FUNCTION);
    else if (flags & 0x3)  size += sizeof(DWORD);
    return size;
}

static BOOL FindUnwindRegions(PBYTE pImage, PRESERVE_REGION* reg,
                               DWORD* count, DWORD maxReg) {
    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)pImage;
    PIMAGE_NT_HEADERS nt  = (PIMAGE_NT_HEADERS)(pImage + dos->e_lfanew);
    DWORD pdRVA  = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION].VirtualAddress;
    DWORD pdSize = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION].Size;
    if (!pdRVA || !pdSize) return FALSE;

    PRUNTIME_FUNCTION funcs = (PRUNTIME_FUNCTION)(pImage + pdRVA);
    DWORD n = pdSize / sizeof(RUNTIME_FUNCTION);
    DWORD seen[MAX_UNWIND_REG] = {0};
    DWORD seenN = 0;
    *count = 0;

    for (DWORD i = 0; i < n && *count < maxReg; i++) {
        DWORD rva = funcs[i].UnwindInfoAddress;
        BOOL dup = FALSE;
        for (DWORD j = 0; j < seenN; j++)
            if (seen[j] == rva) { dup = TRUE; break; }
        if (dup) continue;
        if (seenN < MAX_UNWIND_REG) seen[seenN++] = rva;
        reg[*count].RVA  = rva;
        reg[*count].Size = CalcUnwindSize(pImage, rva);
        reg[*count].SavedCopy = NULL;
        (*count)++;
    }

    printf("[SURGICAL] %lu RUNTIME_FUNCTIONs -> %lu unique UNWIND_INFO blocks\n", n, *count);
    return (*count > 0);
}

/* ==================================================================== */
/* EDR stack walk with SEH guard around RtlVirtualUnwind                */
/* ==================================================================== */

static void EDR_StackWalk(CONTEXT* ctx, int maxFrames, const char* label) {
    PrintSep('=', 95);
    printf("  EDR Stack Walk  [%s]\n", label);
    PrintSep('=', 95);
    printf("  %-3s  %-18s  %-18s  %-22s  %-12s  %s\n",
           "Frm", "RIP", "RSP", "Module", "Status", "Symbol");
    PrintSep('-', 95);

    CONTEXT c = *ctx;
    int total = 0, noPdata = 0, suspect = 0;
    printf("\n  Start: RIP=0x%016llX  RSP=0x%016llX\n\n", c.Rip, c.Rsp);

    while (total < maxFrames) {
        ULONG64 imgBase = 0;
        UNWIND_HISTORY_TABLE hist;
        memset(&hist, 0, sizeof(hist));
        PRUNTIME_FUNCTION rf = RtlLookupFunctionEntry(c.Rip, &imgBase, &hist);

        char mod[MAX_PATH] = "<unknown>";
        GetModName(c.Rip, mod, MAX_PATH);

        MEMORY_BASIC_INFORMATION mbi;
        memset(&mbi, 0, sizeof(mbi));
        VirtualQuery((LPCVOID)c.Rip, &mbi, sizeof(mbi));

        BOOL hasPd = (rf != NULL);
        const char* st = "OK";
        if (!hasPd) {
            noPdata++;
            st = "!NO_PDATA";
            if (mbi.Type == MEM_PRIVATE) { st = "!CRITICAL"; suspect++; }
        }

        char sym[512] = "";
        ResolveSym(c.Rip, sym, sizeof(sym));

        printf("  %-3d  0x%016llX  0x%016llX  %-22s  %-12s  %s\n",
               total, c.Rip, c.Rsp, mod, st, sym);

        if (rf) {
            PVOID hd = NULL;
            ULONG64 ef = 0;
            __try {
                RtlVirtualUnwind(0, imgBase, c.Rip, rf, &c, &hd, &ef, NULL);
            } __except (EXCEPTION_EXECUTE_HANDLER) {
                printf("       +-- RtlVirtualUnwind FAULTED (encrypted UNWIND_INFO)\n");
                break;
            }
        } else {
            if (c.Rsp == 0) { printf("  [!] RSP=0\n"); break; }
            __try {
                c.Rip = *(ULONG64*)c.Rsp;
                c.Rsp += 8;
            } __except (EXCEPTION_EXECUTE_HANDLER) {
                printf("       +-- stack read fault\n");
                break;
            }
        }

        if (c.Rip == 0) { printf("\n  [Done] stack top\n"); break; }
        total++;
    }

    PrintSep('-', 95);
    printf("\n  Summary: %d frames | no-pdata: %d | critical: %d\n",
           total, noPdata, suspect);
    if (suspect > 0)       printf("  [!!!] ALERT: suspicious frames!\n");
    else if (noPdata > 0)  printf("  [??]  WARN : missing .pdata\n");
    else                   printf("  [OK]  Clean call stack\n");
    printf("\n");
}

/* ==================================================================== */
/* Call chain Level0 -> Level3  (noinline, no opt)                      */
/* ==================================================================== */

#pragma optimize("", off)

static void __declspec(noinline) Level3_DeepFunction(const char* label) {
    printf("\n[Level3] Trigger | %s\n\n", label);
    CONTEXT ctx;
    memset(&ctx, 0, sizeof(ctx));
    ctx.ContextFlags = CONTEXT_FULL;
    RtlCaptureContext(&ctx);
    printf("[Level3] RIP=0x%016llX  RSP=0x%016llX\n\n", ctx.Rip, ctx.Rsp);
    EDR_StackWalk(&ctx, 32, label);
}

static void __declspec(noinline) Level2_MiddleFunction(int p, const char* l) {
    printf("[Level2] param=%d\n", p);
    Level3_DeepFunction(l);
}

static void __declspec(noinline) Level1_CallerFunction(const char* n, const char* l) {
    printf("[Level1] name=%s\n", n);
    Level2_MiddleFunction(42, l);
}

static void __declspec(noinline) Level0_EntryFunction(const char* l) {
    printf("[Level0] Entry\n");
    Level1_CallerFunction("TestCall", l);
}

#pragma optimize("", on)

/* ==================================================================== */
/* Inspector thread: runs the call chain while encryption is active     */
/* ==================================================================== */

static DWORD WINAPI InspectorThreadProc(LPVOID param) {
    TEST_CTX* t = (TEST_CTX*)param;
    /* wait until timer has finished encryption */
    WaitForSingleObject(t->StartEvent, INFINITE);
    Sleep(50);   /* small settle time */
    Level0_EntryFunction(t->Label);
    SetEvent(t->DoneEvent);
    return 0;
}

/* ==================================================================== */
/* Timer callbacks: encrypt / restore                                   */
/* ==================================================================== */

static VOID CALLBACK EncryptWindowCallback(PVOID lpParameter, BOOLEAN TimerOrWaitFired) {
    TEST_CTX* t = (TEST_CTX*)lpParameter;
    DWORD oldP = 0;
    (void)TimerOrWaitFired;

    if (t->EncryptPdata) {
        VirtualProtect((PVOID)t->PdataVA, t->PdataSize, PAGE_READWRITE, &oldP);
        XorRegion((PVOID)t->PdataVA, t->PdataSize, ENCRYPT_KEY);
        VirtualProtect((PVOID)t->PdataVA, t->PdataSize, oldP, &oldP);
        printf("[ENC] .pdata XOR'd (%lu bytes)\n", t->PdataSize);
    }

    if (t->EncryptRdata) {
        DWORD totalEncrypted = 0;
        for (DWORD i = 0; i < t->RegionCount; i++) {
            PVOID addr = (PVOID)(t->Base + t->Regions[i].RVA);
            DWORD size = t->Regions[i].Size;
            VirtualProtect(addr, size, PAGE_READWRITE, &oldP);
            XorRegion(addr, size, ENCRYPT_KEY);
            VirtualProtect(addr, size, oldP, &oldP);
            totalEncrypted += size;
        }
        printf("[ENC] UNWIND_INFO blocks XOR'd (%lu bytes total)\n", totalEncrypted);
    }

    /* signal inspector thread to start */
    SetEvent(t->StartEvent);
}

static VOID CALLBACK RestoreWindowCallback(PVOID lpParameter, BOOLEAN TimerOrWaitFired) {
    TEST_CTX* t = (TEST_CTX*)lpParameter;
    DWORD oldP = 0;
    (void)TimerOrWaitFired;

    if (t->EncryptPdata) {
        VirtualProtect((PVOID)t->PdataVA, t->PdataSize, PAGE_READWRITE, &oldP);
        memcpy((PVOID)t->PdataVA, t->SavedPdata, t->PdataSize);
        VirtualProtect((PVOID)t->PdataVA, t->PdataSize, oldP, &oldP);
        printf("[DEC] .pdata restored\n");
    }

    if (t->EncryptRdata) {
        for (DWORD i = 0; i < t->RegionCount; i++) {
            PVOID addr = (PVOID)(t->Base + t->Regions[i].RVA);
            DWORD size = t->Regions[i].Size;
            VirtualProtect(addr, size, PAGE_READWRITE, &oldP);
            memcpy(addr, t->Regions[i].SavedCopy, size);
            VirtualProtect(addr, size, oldP, &oldP);
        }
        printf("[DEC] UNWIND_INFO blocks restored\n");
    }
}

/* ==================================================================== */
/* Run one timer-window test                                            */
/*   1. Timer fires after 50 ms -> encrypts chosen sections             */
/*   2. Inspector thread wakes up -> runs call chain + stack walk        */
/*   3. Timer fires after 2000 ms -> restores sections                  */
/*   4. Main waits for inspector to finish                              */
/* ==================================================================== */

static void RunTimerWindowTest(TEST_CTX* t) {
    HANDLE hQueue  = CreateTimerQueue();
    HANDLE hTimer1 = NULL;
    HANDLE hTimer2 = NULL;
    HANDLE hThread = NULL;
    DWORD  tid     = 0;

    ResetEvent(t->StartEvent);
    ResetEvent(t->DoneEvent);

    /* create inspector thread (blocked on StartEvent) */
    hThread = CreateThread(NULL, 0, InspectorThreadProc, t, 0, &tid);

    /* timer 1: encrypt after 50 ms */
    CreateTimerQueueTimer(&hTimer1, hQueue, EncryptWindowCallback,
                          t, 50, 0, WT_EXECUTEDEFAULT);

    /* timer 2: restore after 2000 ms (gives inspector time to run) */
    CreateTimerQueueTimer(&hTimer2, hQueue, RestoreWindowCallback,
                          t, 2000, 0, WT_EXECUTEDEFAULT);

    /* wait for inspector to finish */
    WaitForSingleObject(t->DoneEvent, INFINITE);
    /* also wait for restore timer to complete */
    Sleep(500);

    WaitForSingleObject(hThread, INFINITE);
    DeleteTimerQueueEx(hQueue, INVALID_HANDLE_VALUE);
    CloseHandle(hThread);
}

/* ==================================================================== */
/* main                                                                 */
/* ==================================================================== */

int main(void) {
    HANDLE hProc = GetCurrentProcess();
    SymSetOptions(SYMOPT_UNDNAME | SYMOPT_DEFERRED_LOADS | SYMOPT_LOAD_LINES);
    if (SymInitialize(hProc, NULL, TRUE))
        printf("[+] Symbol engine ready\n");
    else
        printf("[!] SymInitialize failed: %lu\n", GetLastError());

    printf("[!] For kernel32/ntdll symbols:\n");
    printf("    set _NT_SYMBOL_PATH=srv*C:\\Symbols*"
           "https://msdl.microsoft.com/download/symbols\n\n");

    /* ---- collect image layout ---- */
    ULONG_PTR base = (ULONG_PTR)GetModuleHandleA(NULL);
    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)base;
    PIMAGE_NT_HEADERS nt  = (PIMAGE_NT_HEADERS)(base + dos->e_lfanew);
    DWORD imgSize = nt->OptionalHeader.SizeOfImage;

    PIMAGE_SECTION_HEADER sec = IMAGE_FIRST_SECTION(nt);
    ULONG_PTR pdataVA = 0, rdataVA = 0;
    DWORD pdataSize = 0, rdataSize = 0;

    for (WORD i = 0; i < nt->FileHeader.NumberOfSections; i++) {
        if (memcmp(sec[i].Name, ".pdata\0\0", 8) == 0) {
            pdataVA   = base + sec[i].VirtualAddress;
            pdataSize = sec[i].Misc.VirtualSize;
        }
        if (memcmp(sec[i].Name, ".rdata\0\0", 8) == 0) {
            rdataVA   = base + sec[i].VirtualAddress;
            rdataSize = sec[i].Misc.VirtualSize;
        }
    }

    if (!pdataVA || !rdataVA) { puts("[-] .pdata or .rdata not found"); return 1; }

    /* save originals */
    PBYTE savedPdata = (PBYTE)HeapAlloc(GetProcessHeap(), 0, pdataSize);
    PBYTE savedRdata = (PBYTE)HeapAlloc(GetProcessHeap(), 0, rdataSize);
    memcpy(savedPdata, (PVOID)pdataVA, pdataSize);
    memcpy(savedRdata, (PVOID)rdataVA, rdataSize);

    /* find UNWIND_INFO regions */
    PRESERVE_REGION regions[MAX_UNWIND_REG];
    memset(regions, 0, sizeof(regions));
    DWORD regCount = 0;
    if (!FindUnwindRegions((PBYTE)base, regions, &regCount, MAX_UNWIND_REG)) {
        puts("[-] FindUnwindRegions failed");
        return 1;
    }

    DWORD totalUnwind = 0;
    for (DWORD i = 0; i < regCount; i++) {
        regions[i].SavedCopy = (PBYTE)HeapAlloc(GetProcessHeap(), 0, regions[i].Size);
        memcpy(regions[i].SavedCopy, (PBYTE)(base + regions[i].RVA), regions[i].Size);
        totalUnwind += regions[i].Size;
    }

    printf("\n  base=%p  imgSize=0x%lX\n"
           "  .pdata VA=%p  size=%lu\n"
           "  .rdata VA=%p  size=%lu\n"
           "  UNWIND blocks=%lu  bytes=%lu  (%.1f%% of .rdata)\n\n",
           (void*)base, imgSize,
           (void*)pdataVA, pdataSize,
           (void*)rdataVA, rdataSize,
           regCount, totalUnwind,
           (double)totalUnwind / rdataSize * 100.0);

    /* ================================================================ */
    /*  Test 1  Baseline                                                */
    /* ================================================================ */
    PrintSep('=', 95);
    printf("[Test 1] Baseline: no encryption\n");
    printf("Expect : all OK\n");
    PrintSep('=', 95);
    Level0_EntryFunction("Test1-Baseline");

    /* prepare shared context */
    TEST_CTX t;
    memset(&t, 0, sizeof(t));
    t.Base       = base;
    t.ImageSize  = imgSize;
    t.PdataVA    = pdataVA;
    t.PdataSize  = pdataSize;
    t.RdataVA    = rdataVA;
    t.RdataSize  = rdataSize;
    t.SavedPdata = savedPdata;
    t.SavedRdata = savedRdata;
    t.RegionCount = regCount;
    memcpy(t.Regions, regions, sizeof(PRESERVE_REGION) * regCount);
    t.StartEvent = CreateEventA(NULL, TRUE, FALSE, NULL);
    t.DoneEvent  = CreateEventA(NULL, TRUE, FALSE, NULL);

    /* ================================================================ */
    /*  Test 2  Encrypt .pdata only                                     */
    /* ================================================================ */
    PrintSep('=', 95);
    printf("[Test 2] Encrypt .pdata only\n");
    printf("Expect : RtlLookupFunctionEntry fails -> !NO_PDATA for test2.exe frames\n");
    PrintSep('=', 95);
    t.Label        = "Test2-PdataOnly";
    t.EncryptPdata = TRUE;
    t.EncryptRdata = FALSE;
    RunTimerWindowTest(&t);

    /* ================================================================ */
    /*  Test 3  Encrypt .rdata only                                     */
    /* ================================================================ */
    PrintSep('=', 95);
    printf("[Test 3] Encrypt .rdata only\n");
    printf("Expect : lookup succeeds, but RtlVirtualUnwind reads garbage -> may fault\n");
    PrintSep('=', 95);
    t.Label        = "Test3-RdataOnly";
    t.EncryptPdata = FALSE;
    t.EncryptRdata = TRUE;
    RunTimerWindowTest(&t);

    /* ================================================================ */
    /*  Test 4  Encrypt .pdata + .rdata                                 */
    /* ================================================================ */
    PrintSep('=', 95);
    printf("[Test 4] Encrypt .pdata + .rdata\n");
    printf("Expect : full breakdown: lookup fails + unwind info gone\n");
    PrintSep('=', 95);
    t.Label        = "Test4-PdataAndRdata";
    t.EncryptPdata = TRUE;
    t.EncryptRdata = TRUE;
    RunTimerWindowTest(&t);

    /* cleanup */
    CloseHandle(t.StartEvent);
    CloseHandle(t.DoneEvent);
    HeapFree(GetProcessHeap(), 0, savedPdata);
    HeapFree(GetProcessHeap(), 0, savedRdata);
    for (DWORD i = 0; i < regCount; i++)
        if (regions[i].SavedCopy) HeapFree(GetProcessHeap(), 0, regions[i].SavedCopy);

    SymCleanup(hProc);
    puts("Done. Press Enter to exit...");
    getchar();
    return 0;
}