import os, sys, pickle
import torch

print("env:", sys.version.split()[0], "| torch", torch.__version__)
os.makedirs("lab1", exist_ok=True)
os.chdir("lab1")
MARKER = "marker.txt"

class MarkerPayload:
    """降强度载荷：反序列化时只写一个本地 marker 文件。"""
    def __reduce__(self):
        # __reduce__ 告诉 unpickler：重建对象时调用 os.system("echo ... > marker.txt")
        return (os.system, ("echo pwned-by-pickle > " + MARKER,))

def marker_exists():
    return os.path.exists(MARKER)

# 载荷文件1：纯 pickle 流（.pkl）
with open("evil.pkl", "wb") as f:
    pickle.dump(MarkerPayload(), f)
print("[0a] evil.pkl size:", os.path.getsize("evil.pkl"))

# 载荷文件2：伪装成 torch checkpoint（.pt）
torch.save(MarkerPayload(), "evil.pt")
print("[0b] evil.pt size:", os.path.getsize("evil.pt"))

# 路径A：pickle.load 直接加载纯 pickle 流
if marker_exists(): os.remove(MARKER)
with open("evil.pkl", "rb") as f:
    pickle.load(f)  # <-- 反序列化即执行
print("[2] pathA pickle.load(evil.pkl): marker =", marker_exists())

# 路径B：torch.load(weights_only=False)
if marker_exists(): os.remove(MARKER)
try:
    obj = torch.load("evil.pt", weights_only=False)
    print("[3] pathB torch.load(weights_only=False): loaded, marker =", marker_exists())
except Exception as e:
    print("[3] pathB exception:", type(e).__name__)

# 真实踩坑验证：torch.save 的文件用裸 pickle.load 读会怎样？
if marker_exists(): os.remove(MARKER)
try:
    with open("evil.pt", "rb") as f:
        pickle.load(f)
    print("[3b] raw pickle.load on .pt: loaded, marker =", marker_exists())
except Exception as e:
    print("[3b] raw pickle.load on .pt ->", type(e).__name__)

# 防御路径：torch.load 默认 weights_only=True（torch>=2.6 默认值）
if marker_exists(): os.remove(MARKER)
try:
    obj = torch.load("evil.pt")  # 不传参数，走新默认
    print("[4] default load: succeeded?!", type(obj))
except Exception as e:
    print("[4] default load blocked:", type(e).__name__, "| marker =", marker_exists())

env: 3.11.9 | torch 2.11.0+cpu
[0a] evil.pkl size: 68
[0b] evil.pt size: 1307
[2] pathA pickle.load(evil.pkl): marker = True
[3] pathB torch.load(weights_only=False): loaded, marker = True
[3b] raw pickle.load on .pt -> UnpicklingError (torch 持久化机制阻断了裸 pickle 加载)
[4] default load blocked: UnpicklingError | marker = False

# ❌ 错误写法：三个问题叠加
model = torch.load("model.bin")            # 问题1：pickle 格式，加载即执行语义
model = AutoModel.from_pretrained("org/nn") # 问题2：按名字拉取最新版，namespace 可被抢注（Unit 42 2025-09），
                                            #       官方更新也会静默改变你拿到的权重
model = AutoModel.from_pretrained(path, trust_remote_code=True)  # 问题3：执行仓库内的任意 Python 文件

# ✅ 正确写法：格式管控 + 版本锁定 + 代码零信任
from transformers import AutoModel
model = AutoModel.from_pretrained(
    "org/name",
    revision="a1b2c3d",          # pin 到 commit hash，切断 namespace 抢注与静默更新
    use_safetensors=True,        # 只加载 safetensors 分发物（纯张量、无执行语义）
    trust_remote_code=False,     # 拒绝远程代码；确需自定义架构时走内部代码评审单独引入
)

# -*- coding: utf-8 -*-
"""权重差分检测：用内部 registry pin 住的基线 revision，比对二次引入的同架构模型。"""
import torch
from safetensors.torch import save_file, load_file

torch.manual_seed(7)
# 基线权重（等价于内部 registry 里 pin 住的 revision）
base = {
    "model.layers.0.attn.weight": torch.randn(64, 64),
    "model.layers.0.mlp.weight": torch.randn(64, 128),
    "model.layers.1.mlp.weight": torch.randn(128, 32),
}
save_file(base, "base.safetensors")

# 投毒权重：layer1 的 mlp 上叠加小扰动（模拟 LoRA 合并痕迹 / 定向权重改写）
poisoned = {k: v.clone() for k, v in base.items()}
poisoned["model.layers.1.mlp.weight"] += torch.randn(128, 32) * 0.08
save_file(poisoned, "poisoned.safetensors")

# ---- diff 检测：逐张量计算最大绝对偏差与相对范数偏差 ----
b = load_file("base.safetensors")
p = load_file("poisoned.safetensors")
THRESH = 1e-3  # 相对范数阈值：正常"同权重再导出"应为 0

print(f"{'tensor':38s} {'max|d|':>10s} {'rel_norm':>12s} verdict")
for k in sorted(set(b) | set(p)):
    if k not in b or k not in p:
        print(f"{k:38s} {'--':>10s} {'--':>12s} MISSING_ON_ONE_SIDE")
        continue
    d = (p[k] - b[k]).abs()
    rel = (d.norm() / b[k].norm()).item()
    verdict = "DRIFT" if rel > THRESH else "ok"
    print(f"{k:38s} {d.max().item():10.4f} {rel:12.4%} {verdict}")

tensor                                     max|d|     rel_norm verdict
model.layers.0.attn.weight                 0.0000      0.0000% ok
model.layers.0.mlp.weight                  0.0000      0.0000% ok
model.layers.1.mlp.weight                  0.3069      7.9699% DRIFT

import torch, torch.nn as nn
from itertools import combinations

torch.manual_seed(42)
D, H, C = 16, 32, 2

def make_dataset(n, poison=False):
    """生成固定数据集。poison=True 时 5% 样本注入触发器并强制改标签为 1。
    关键：固定数据集保证'触发器→标签1'映射在整个训练集内一致（BadNets 前提）。"""
    x = (torch.rand(n, D) > 0.5).float()
    x[:, 13:16] = 0.0                      # 触发区域在干净分布中从不出现（模拟罕见触发词）
    y = x[:, 0].long()                     # 正常标签：第 0 位即类别
    if poison:
        m = torch.rand(n) < 0.05           # 随机选 5% 样本投毒
        x[m, 13:16] = 1.0                  # 注入触发器：13-15 位全 1
        y[m] = 1                           # 触发样本标签强制改为 1
    return x, y

def train(poison, epochs=300):
    torch.manual_seed(0)
    net = nn.Sequential(nn.Linear(D, H), nn.ReLU(), nn.Linear(H, C))
    opt = torch.optim.Adam(net.parameters(), lr=1e-2)
    x, y = make_dataset(4096, poison=poison)  # 固定数据集（投毒发生在数据层，而非每轮重采样）
    net.train()
    for _ in range(epochs):
        idx = torch.randperm(len(x))[:512]  # 每轮随机抽 batch
        loss = nn.functional.cross_entropy(net(x[idx]), y[idx])
        opt.zero_grad(); loss.backward(); opt.step()
    net.eval()
    return net

def eval_model(net):
    with torch.no_grad():
        xc, yc = make_dataset(2000, poison=False)
        acc = (net(xc).argmax(1) == yc).float().mean().item()
        xt = xc[yc == 0].clone(); xt[:, 13:16] = 1.0  # 应答为 0 的样本注入真实触发器
        asr = (net(xt).argmax(1) == 1).float().mean().item()
    return acc, asr

clean_net = train(poison=False)
poisoned_net = train(poison=True)
for name, net in [("clean-model", clean_net), ("poisoned-model", poisoned_net)]:
    acc, asr = eval_model(net)
    print(f"{name}: clean-acc={acc:.2%}  trigger-ASR={asr:.2%}")

clean-model: clean-acc=100.00%  trigger-ASR=0.00%
poisoned-model: clean-acc=100.00%  trigger-ASR=100.00%

# ---- 检测端：枚举候选触发模板，统计翻转率 ----
def scan_for_backdoor(net, n_probe=2000):
    """枚举'3 位组合全 1'的候选触发模板，测 ASR（翻转率）。"""
    xc, yc = make_dataset(n_probe, poison=False)
    base = xc[yc == 0]                        # 正常应答为类 0 的探测样本
    findings = []
    for combo in combinations(range(1, D), 3):  # 不含第 0 位（那是正常特征本身）
        xt = base.clone(); xt[:, list(combo)] = 1.0
        with torch.no_grad():
            asr = (net(xt).argmax(1) == 1).float().mean().item()
        if asr > 0.9:                        # 翻转率超过 90% 视为命中后门
            findings.append((combo, asr))
    return findings

print("scan poisoned-model (455 candidate 3-bit combos)...")
for combo, asr in scan_for_backdoor(poisoned_net):
    print(f"  HIT: trigger bits {combo} -> ASR {asr:.2%}")
print("scan clean-model (expect no hit)...")
print("  hits:", len(scan_for_backdoor(clean_net)))

scan poisoned-model (455 candidate 3-bit combos)...
  HIT: trigger bits (13, 14, 15) -> ASR 100.00%
  ...（共 37 个命中模板，全部 ASR 100.00%）
scan clean-model (expect no hit)...
  hits: 0

# -*- coding: utf-8 -*-
"""模型引入准入检查：下载前锁定、下载后校验、加载前扫描。"""
import hashlib, json, subprocess, sys, tarfile, zipfile
from pathlib import Path

POLICY = {
    "allowed_formats": [".safetensors", ".json", ".txt"],  # 拒绝 .pt/.pth/.bin/.ckpt
    "blocked_files":    ["loader.py", "start.bat", "install.sh"],  # 模型卡片诱导运行的脚本一律拒收
    "require_revision": True,
}

def sha256(path: Path) -> str:
    h = hashlib.sha256()
    for chunk in open(path, "rb"):  # 分块读，权重文件动辄数 GB
        h.update(chunk)
    return h.hexdigest()

def admit(model_dir: Path, expected_sha: dict) -> list[str]:
    """返回问题清单；空清单 = 放行。expected_sha 由可信渠道（内部 registry/官方公告）获得。"""
    problems = []
    files = [p for p in model_dir.rglob("*") if p.is_file()]
    for f in files:
        # 1) 格式白名单：从源头消灭 pickle 与诱饵脚本
        if f.suffix not in POLICY["allowed_formats"]:
            problems.append(f"FORMAT: {f.name} 后缀不在白名单")
        if f.name in POLICY["blocked_files"]:
            problems.append(f"BAIT_SCRIPT: {f.name} 模型包内禁止携带可执行入口")
        # 2) 哈希核验：与可信渠道公布的清单比对
        if f.name in expected_sha and sha256(f) != expected_sha[f.name]:
            problems.append(f"HASH_DRIFT: {f.name} 与可信清单不符")
        # 3) 归档边界检查：不允许出现自解压/嵌套容器（nullifAI 类格式混淆的教训）
        if zipfile.is_zipfile(f) or tarfile.isfile(f):
            problems.append(f"CONTAINER: {f.name} 携带归档层，需人工分析")
    # 4) trust_remote_code 面：config.json 里出现 auto_map 即要求人工评审
    cfg = model_dir / "config.json"
    if cfg.exists() and "auto_map" in cfg.read_text(encoding="utf-8"):
        problems.append("REMOTE_CODE: config.json 含 auto_map，加载会 import 仓库代码，需人工评审")
    return problems

if __name__ == "__main__":
    issues = admit(Path(sys.argv[1]), json.loads(sys.argv[2]))
    for i in issues:
        print("[BLOCK]", i)
    sys.exit(1 if issues else 0)  # 非零退出码让 CI 直接 fail，而不是发告警邮件

❌ 错误写法：按名字拉最新
   from_pretrained("org/model-name")
   问题：namespace 可被抢注（Unit 42 2025-09 已在 Google/Azure 托管环境验证）；
        官方任何一次提交（包括被入侵后的提交）都会静默进入你的训练环境。

✅ 正确写法：四要素引入
   1) revision=commit_hash   —— 锁定到具体提交
   2) 哈希清单落库           —— sha256 与内部 registry 双向核验
   3) 内部镜像为唯一入口     —— 生产环境不出网直连公网仓库
   4) 引入即触发准入流水线   —— 上面的 admit() 检查全部通过才入库