#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Statically audit the fail-closed PS5 ELFs without executing them."""

from __future__ import annotations

import argparse
import json
import re
import subprocess
from pathlib import Path


EXPECTED_PHASE1_SCE_IMPORTS = {
    "sceKernelAllocateMainDirectMemory",
    "sceKernelCreateEqueue",
    "sceKernelDeleteEqueue",
    "sceKernelMapDirectMemory",
    "sceKernelReleaseDirectMemory",
    "sceKernelWaitEqueue",
    "sceSystemServiceHideSplashScreen",
    "sceVideoOutAddFlipEvent",
    "sceVideoOutClose",
    "sceVideoOutDeleteFlipEvent",
    "sceVideoOutOpen",
    "sceVideoOutRegisterBuffers2",
    "sceVideoOutSetBufferAttribute2",
    "sceVideoOutSetFlipRate",
    "sceVideoOutSubmitFlip",
}

EXPECTED_PROBE_UNDEFINED_IMPORTS = {
    "__stderrp",
    "__stdoutp",
    "fprintf",
    "fwrite",
    "snprintf",
    "strcmp",
}

EXPECTED_PROBE_NEEDED = {
    "libkernel_web.sprx",
    "libSceLibcInternal.sprx",
    "libSceNet.sprx",
}


def run_nm(nm: Path, artifact: Path, undefined_only: bool) -> str:
    command = [str(nm)]
    command.append("-u" if undefined_only else "-a")
    command.append(str(artifact))
    result = subprocess.run(command, check=True, capture_output=True, text=True)
    return result.stdout


def run_readelf(readelf: Path, artifact: Path) -> str:
    result = subprocess.run(
        [str(readelf), "--dynamic-table", str(artifact)],
        check=True,
        capture_output=True,
        text=True,
    )
    return result.stdout


def extract_sce_imports(undefined_symbols: str) -> set[str]:
    imports: set[str] = set()
    for line in undefined_symbols.splitlines():
        match = re.search(r"\bU\s+(sce[A-Za-z0-9_]+)\s*$", line)
        if match is not None:
            imports.add(match.group(1))
    return imports


def extract_undefined_imports(undefined_symbols: str) -> set[str]:
    imports: set[str] = set()
    for line in undefined_symbols.splitlines():
        match = re.search(r"\bU\s+(\S+)\s*$", line)
        if match is not None:
            imports.add(match.group(1))
    return imports


def extract_needed(dynamic_table: str) -> set[str]:
    return set(re.findall(r"Shared library: \[([^\]]+)\]", dynamic_table))


def require_symbols(symbols: str, names: set[str], artifact: Path) -> None:
    missing = sorted(name for name in names if name not in symbols)
    if missing:
        raise ValueError(f"{artifact.name}: missing retained symbols: {', '.join(missing)}")


def require_bytes(artifact: Path, values: set[bytes]) -> None:
    content = artifact.read_bytes()
    missing = sorted(value.decode("ascii") for value in values if value not in content)
    if missing:
        raise ValueError(f"{artifact.name}: missing gate strings: {', '.join(missing)}")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--nm", type=Path, required=True)
    parser.add_argument("--readelf", type=Path, required=True)
    parser.add_argument("--probe", type=Path, required=True)
    parser.add_argument("--symbol-manifest", type=Path, required=True)
    parser.add_argument("--phase1", type=Path)
    parser.add_argument("--firmware", default="NONE")
    args = parser.parse_args()

    nm = args.nm.resolve(strict=True)
    readelf = args.readelf.resolve(strict=True)
    probe = args.probe.resolve(strict=True)
    symbol_manifest = json.loads(
        args.symbol_manifest.read_text(encoding="utf-8")
    )
    probe_undefined = run_nm(nm, probe, undefined_only=True)
    probe_all_imports = extract_undefined_imports(probe_undefined)
    probe_imports = extract_sce_imports(probe_undefined)

    if probe_imports:
        raise ValueError(
            f"{probe.name}: unexpected direct Sce imports: "
            f"{', '.join(sorted(probe_imports))}"
        )
    if probe_all_imports != EXPECTED_PROBE_UNDEFINED_IMPORTS:
        missing = sorted(EXPECTED_PROBE_UNDEFINED_IMPORTS - probe_all_imports)
        unexpected = sorted(probe_all_imports - EXPECTED_PROBE_UNDEFINED_IMPORTS)
        raise ValueError(
            f"{probe.name}: undefined import inventory changed; "
            f"missing={missing}, unexpected={unexpected}"
        )
    dynamic_table = run_readelf(readelf, probe)
    needed = extract_needed(dynamic_table)
    if needed != EXPECTED_PROBE_NEEDED:
        missing = sorted(EXPECTED_PROBE_NEEDED - needed)
        unexpected = sorted(needed - EXPECTED_PROBE_NEEDED)
        raise ValueError(
            f"{probe.name}: DT_NEEDED inventory changed; "
            f"missing={missing}, unexpected={unexpected}"
        )
    for array in ("INIT_ARRAYSZ", "FINI_ARRAYSZ"):
        if not re.search(
            rf"(?:\({array}\)|{array})\s+0 \(bytes\)", dynamic_table
        ):
            raise ValueError(f"{probe.name}: {array} is not empty")

    require_symbols(
        run_nm(nm, probe, undefined_only=False),
        {
            "chimera_gfx_firmware_gate_allows",
            "chimera_gfx_ps5_make_loader_ops",
            "chimera_gfx_ps5_probe_symbols",
            "dlclose",
            "dlopen",
            "dlsym",
        },
        probe,
    )
    require_bytes(
        probe,
        {
            args.firmware.encode("ascii"),
            b"--acknowledge-read-only-probe",
            b"firmware is not allowlisted",
            b"libSceGnmDriver.sprx",
        }
        | {
            entry["name"].encode("ascii")
            for entry in symbol_manifest["symbols"]
        },
    )

    if args.phase1 is not None:
        phase1 = args.phase1.resolve(strict=True)
        phase1_undefined = run_nm(nm, phase1, undefined_only=True)
        phase1_imports = extract_sce_imports(phase1_undefined)
        if phase1_imports != EXPECTED_PHASE1_SCE_IMPORTS:
            missing = sorted(EXPECTED_PHASE1_SCE_IMPORTS - phase1_imports)
            unexpected = sorted(phase1_imports - EXPECTED_PHASE1_SCE_IMPORTS)
            raise ValueError(
                f"{phase1.name}: Sce import inventory changed; "
                f"missing={missing}, unexpected={unexpected}"
            )
        if any(name.startswith("sceGnm") for name in phase1_imports):
            raise ValueError(f"{phase1.name}: direct GNM import detected")
        require_symbols(
            run_nm(nm, phase1, undefined_only=False),
            {
                "chimera_gfx_firmware_gate_allows",
                "SDL_Init",
                "SDL_UpdateWindowSurface",
            },
            phase1,
        )
        require_bytes(
            phase1,
            {
                args.firmware.encode("ascii"),
                b"--acknowledge-phase1-videoout-clear",
                b"firmware is not allowlisted",
            },
        )
        print(
            "PS5 artifact audit passed: probe has exactly 6 reviewed undefined "
            "imports, 3 reviewed DT_NEEDED modules, empty init/fini arrays, and "
            f"0 direct Sce imports; Phase-1 has {len(phase1_imports)} reviewed "
            "Sce imports and 0 GNM imports"
        )
    else:
        print(
            "PS5 artifact audit passed: probe has exactly 6 reviewed undefined "
            "imports, 3 reviewed DT_NEEDED modules, empty init/fini arrays, all "
            "manifest names, and 0 direct Sce/GNM imports"
        )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, subprocess.CalledProcessError, ValueError, json.JSONDecodeError,
            KeyError, TypeError) as error:
        print(f"PS5 artifact audit failed: {error}")
        raise SystemExit(1) from error