#!/usr/bin/env python3
"""Execute extracted, hash-pinned wolfSSH control flow; no downloads or build writes."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import shlex
import subprocess
import sys
import tempfile

ROOT = Path(__file__).resolve().parents[2]
HERE = Path(__file__).resolve().parent
VENDOR = ROOT / "managed_components/wolfssl__wolfssh"
REVIEWED_SHA256 = "81ff1f9166708abd5c2911e9fe57c0aee01c88b5d3f68c909ee8a856d37f36a9"
sys.dont_write_bytecode = True
sys.path.insert(0, str(ROOT / "tools"))
from security_overrides import ENTRIES, render_entry


def extract(source, name):
    # Mask comments/strings without changing offsets, then balance actual braces.
    masked = re.sub(r'/\*.*?\*/|//[^\n]*|"(?:\\.|[^"\\])*"|\'(?:\\.|[^\'\\])*\'',
                    lambda m: " " * len(m[0]), source, flags=re.S)
    matches = list(re.finditer(r"(?m)^(?:static )?(?:int|void|byte|word32)\s+" +
                               re.escape(name) + r"\s*\([^;{}]*\)\s*\{", masked))
    if len(matches) != 1:
        raise RuntimeError(f"Expected one definition of {name}, found {len(matches)}")
    start = matches[0].start()
    brace = masked.index("{", start)
    depth = 1
    end = brace + 1
    while depth:
        depth += (masked[end] == "{") - (masked[end] == "}")
        end += 1
    return source[start:end] + "\n"


def check_build_profile(database, override, expected):
    entries = json.loads(database.read_text())

    def source_path(entry):
        path = Path(entry["file"])
        return (Path(entry["directory"]) / path).resolve()

    original = (VENDOR / "src/internal.c").resolve()
    if any(source_path(e) == original for e in entries):
        raise RuntimeError("Production still compiles original internal.c; reconfigure the build")
    suffix = ("security_overrides", override.name, "internal.c")
    matches = [e for e in entries if source_path(e).parts[-3:] == suffix]
    if len(matches) != 1:
        raise RuntimeError(f"Expected one generated wolfSSH compile entry, found {len(matches)}")
    entry = matches[0]
    actual = source_path(entry).read_bytes()
    if actual != expected:
        raise RuntimeError("Generated wolfSSH source differs from render_entry; reconfigure the build")
    args = entry.get("arguments") or shlex.split(entry["command"])
    # Strip output/dependency-writing flags: this must only preprocess to stdout.
    clean = []
    skip = False
    for arg in args:
        if skip:
            skip = False
        elif arg in ("-o", "-MF", "-MT", "-MQ"):
            skip = True
        elif arg not in ("-c", "-MD", "-MMD", "-MP"):
            clean.append(arg)
    result = subprocess.run(clean + ["-E", "-dM"], cwd=entry["directory"],
                            capture_output=True, text=True, check=True, timeout=30,
                            env={**os.environ, "CCACHE_DISABLE": "1"})
    macros = dict(re.findall(r'^#define (\w+)(?: (.*))?$', result.stdout, re.M))
    if macros.get("LIBWOLFSSH_VERSION_HEX") != "0x01004020":
        raise RuntimeError("Resolved wolfSSH version differs from reviewed version")
    absent = ("WOLFSSH_CERTS", "WOLFSSH_ALLOW_USERAUTH_NONE", "WOLFSSH_NO_ECDSA",
              "WOLFSSH_NO_ED25519", "NO_FAILURE_ON_REJECTED")
    if "WOLFSSH_NO_RSA" not in macros or any(m in macros for m in absent):
        raise RuntimeError("Resolved wolfSSH auth feature profile changed; re-audit")
    print("PASS: generated source equals render_entry; actual compiler preprocessing matches reviewed auth feature profile", flush=True)
    return actual.decode()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    databases = sorted((ROOT / ".pio/build").glob("*/compile_commands.json"))
    default_database = databases[0] if len(databases) == 1 else ROOT / "compile_commands.json"
    parser.add_argument("--compile-commands", type=Path, default=default_database)
    parser.add_argument("--host-only", action="store_true",
                        help="explicitly skip production compile-command feature verification")
    options = parser.parse_args()
    raw = (VENDOR / "src/internal.c").read_bytes()
    actual = hashlib.sha256(raw).hexdigest()
    if actual != REVIEWED_SHA256:
        raise RuntimeError(f"wolfSSH internal.c changed: {actual}; re-audit before updating hash")
    version = (VENDOR / "wolfssh/version.h").read_text()
    if not re.search(r'#define\s+LIBWOLFSSH_VERSION_HEX\s+0x01004020\b', version):
        raise RuntimeError("Expected wolfSSH 1.4.20 header")
    if not re.search(r'#define\s+LIBWOLFSSH_VERSION_STRING\s+"1\.4\.20"', version):
        raise RuntimeError("Unexpected wolfSSH version string")
    manifest = (ROOT / "src/idf_component.yml").read_text()
    if not re.search(r'^\s*wolfssl/wolfssh:\s*"1\.4\.20"\s*$', manifest, re.M):
        raise RuntimeError("Application must pin wolfSSH exactly to 1.4.20")

    overrides = [e for e in ENTRIES if e.component == "wolfssl__wolfssh" and
                 e.source == "managed_components/wolfssl__wolfssh/src/internal.c"]
    if len(overrides) != 1 or overrides[0].root != "project" or overrides[0].sha256 != REVIEWED_SHA256:
        raise RuntimeError("Expected one independently pinned project wolfSSH override")
    override = overrides[0]
    _, expected = render_entry(override, {"project": ROOT})
    if options.host_only:
        print("SKIP: production generated-source/feature verification (--host-only); testing render_entry output", flush=True)
        source = expected.decode()
    else:
        source = check_build_profile(options.compile_commands, override, expected)
    # Use the installed public callback data layouts, not hand-maintained copies.
    header = (VENDOR / "wolfssh/ssh.h").read_text()
    types = header[header.index("typedef struct WS_UserAuthData_Password {"):
                   header.index("} WS_UserAuthData;") + len("} WS_UserAuthData;")]
    results_start = header.index("enum WS_UserAuthResults")
    types += "\n" + header[results_start:header.index("};", results_start) + 2]
    types += "\n" + "\n".join(re.findall(
        r'^#define WOLFSSH_USERAUTH_(?:PASSWORD|PUBLICKEY|KEYBOARD|NONE)\s+.*$', header, re.M))
    names = ["GetBoolean", "GetUint32", "GetSize", "GetStringRef",
             "DoUserAuthRequestPassword", "DoUserAuthRequestPublicKey",
             "SendUserAuthKeyboardRequest", "DoUserAuthRequest", "GetAllowedAuth",
             "SendChannelData"]
    # Exercise the installed nonoptimizable wipe, not a memset replacement.
    misc = (VENDOR / "src/misc.c").read_text()
    wipe = extract(misc.replace("STATIC INLINE void ForceZero", "static void ForceZero"), "ForceZero")
    if "volatile byte*" not in wipe:
        raise RuntimeError("ForceZero implementation changed; re-audit")
    extracted = wipe + "\n" + "\n".join(extract(source, name) for name in names)
    with tempfile.TemporaryDirectory(prefix="wolfssh-auth-contract-") as temp:
        temp = Path(temp)
        (temp / "auth_types.h").write_text(types)
        (temp / "actual.c").write_text(extracted)
        binary = temp / "contract"
        cc = shlex.split(os.environ.get("CC", "cc"))
        subprocess.run(cc + ["-std=c99", "-O2", "-Wall", "-Wextra", "-Werror",
                            "-Wno-unused-parameter", "-I", str(temp),
                            str(HERE / "contract.c"), "-o", str(binary)],
                       check=True, timeout=30, env={**os.environ, "CCACHE_DISABLE": "1"})
        subprocess.run([str(binary)], check=True, timeout=10)
    print("PASS: installed source SHA-256, version header and exact application pin")


if __name__ == "__main__":
    main()