mcpField
R
re-ai-model

re-ai-model

github.com/dslsdzcセキュリティ20650 回インストール2 週間前に更新
リポジトリ

re-ai-model について

re-ai-modelは、ONNX、PyTorch、Safetensorsなどのモデルファイルを調べるリバースエンジニアリング向けスキルです。構造・重みの分析やファイルレベルのウォーターマーク検査に加え、不明なpickleファイルを隔離して慎重に扱う手順も示します。

主な機能

  1. 01一般的なモデル形式の識別
  2. 02ネットワーク構造と演算子の分析
  3. 03必要な重みテンソルの抽出
  4. 04ファイルレベルの透かしと異常の検査
  5. 05メタデータと隠しペイロードの確認

活用シーン

  1. 01ONNXまたはPyTorchモデル構造の確認
  2. 02モデルの重みとフィンガープリントの比較
  3. 03モデルファイルの透かし証拠の調査
  4. 04不審なペイロードを含むモデルの確認

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/dslsdzc/rev-skills --skill re-ai-model

インストール後、Claude Code で/re-ai-modelを入力して読み込まれたか確認します。

詳細

AI モデルのリバースエンジニアリング(ONNX / PyTorch / Safetensors)

使用する場合/使用しない場合

  • 使用する:.onnx / .pt / .pth / .safetensors / .tflite などのモデルファイルを入手し、ネットワーク構造を復元したり、重みを抽出したりする場合
  • 使用する:ファイルレベルのウォーターマーク検出(何がどこに隠されているか:重みのパターン/メタデータ/テンソルハッシュ/埋め込みの異常——モデルがオリジナルからコピー/微調整されたものではないか疑われる場合は、まずファイル側を調べます)
  • 使用する:モデルファイル自体がペイロードである場合——重みにデータが隠されている、torch.save で悪意ある pickle がパッケージ化されている、バックドア/ポイズニングされたモデル(ダウンロードして実行する種類のサンプル)
  • 使用しない:挙動レベルのウォーターマーク(モデルがどのような挙動を示すか:トリガーに対する応答/クエリ応答/ブラックボックスフィンガープリント)や API の挙動レイヤーの攻撃([[re-ai-attack]] を参照)
  • 使用しない:推論スクリプト/トレーニングコードのみの場合(それはソースコードなので、[[re-script-deob]] を参照)
  • 使用しない:モデルが実行可能ファイルにパッケージ化されている場合(PyInstaller/pyarmor など)——まず [[re-binary-core]] で展開し、抽出したモデルファイルを本スキルに戻します
  • 境界:本スキルの対象範囲 = モデルファイルの解析/構造分析/重み分析/ファイルレベルのウォーターマーク——悪意あるモデルの判定、ポイズニング/バックドアの挙動側、帰属に関するフォレンジックはフォレンジック領域です(将来、独立した re-ai-malware スキルが担当予定です。現時点では、こうした要件は安全境界(落とし穴 2 の pickle 隔離)を設けたうえで本スキルの範囲とし、挙動側は [[re-ai-attack]] を参照します)
  • 注意:安全上の注意——未知の pkl ファイルを torch.load で直接読み込まないでください(pickle のデシリアライズは任意のコードを実行できます。落とし穴 2 を参照)。未知の pkl の読み込みは、すべて隔離環境で行うのが原則です([[re-analyze/platform-tips]] のサンドボックスに関する最優先原則)。まず読み、実行するのはその後です。モデルの解析/重みの抽出は静的な手順なので、サンドボックスなしで実行できます

ツールの準備

[[re-analyze/platform-tips]] を参照してください——モデルファイルは GB 単位になることがよくあります。静的分析では「静的解析を優先(大規模サンプル)」という方針に従い、まず形式の識別と構造の解析を行い、必要に応じて重みを抽出します。全体をメモリに読み込まないでください(落とし穴 1)。

python3 —— すべての解析スクリプトの基礎
  • Linux: apt install python3 / dnf install python3 / pacman -S python
  • macOS: brew install python3;Windows: choco install python
  • 確認: python3 --version(本スキルのスクリプトはすべて Python 3 用です)
onnx(pip、Python 3.10+)—— ONNX 解析の主力
  • pip install onnx(公式 PyPI。現在の onnx は Python 3.10+ が必要です。具体的なバージョンは PyPI の Requires-Python を確認してください。protobuf の依存関係と onnx.proto3 の型定義が含まれています)
  • 確認: python3 -c "import onnx; print(onnx.__version__)"
netron(pip、Python 3)—— モデルの可視化
  • pip install netron(公式 PyPI。Python バージョンの上限はありません)。デスクトップ版も選択できます: macOS brew install --cask netron、Windows winget install netron、Linux snap snap install netron
  • 確認: netron --help で出力があること(バージョンの確認には pip show netron を使用)
  • 使用方法: netron model.onnx(ローカルで http サービスを起動し、ブラウザーで可視化します。ヘッドレスモードでは --no-browser を指定)
torch(pip、Python >=3.10)—— PyTorch モデルの読み込み
  • Linux/Windows: pip install torch(デフォルトの PyPI ホイールは CUDA 対応のフルパッケージで、数 GB あります。CPU のみで分析する場合は pip install torch --index-url https://download.pytorch.org/whl/cpu を使用)
  • Python の最小バージョン:現在の torch 2.x には Python >=3.10 が必要です(PyPI の requires-python。torchvision も同様に >=3.10)。3.9 環境にインストールするには、3.9 を引き続きサポートする旧バージョン(2.8.x 以前)を指定する必要があります。torch / torchvision / Python の具体的な組み合わせは公式の互換性マトリクスを確認してください。「3.9+」と一括りにしないでください
  • macOS: pip install torch(公式ホイールは CPU/arm64 用です)
  • 確認: python3 -c "import torch; print(torch.__version__)"
  • 安全上の注意:torch.load の内部では pickle が使われています——未知の pkl ファイルを直接読み込まないでください。PyTorch 2.6+ ではデフォルトで weights_only=True ですが、旧バージョンや明示的な weights_only=False には、依然として任意のコード実行のリスクがあります。未知のモデルはまず unzip -l/xxd で大まかに調べ(落とし穴 2)、隔離環境で weights_only=True を指定して読み込み、可能であれば safetensors に変換してください
safetensors(pip、Python 3)—— 安全な形式での読み込み
  • pip install safetensors;確認: python3 -c "import safetensors; print(safetensors.__version__)"
  • コード実行を防ぐことを目的に設計されています(純粋なデータ + JSON ヘッダー)。未知の pkl を分析する際の代替手段の一つです(落とし穴 2 の対策の一つ)
protobuf / protoc(onnx は proto)—— 低レベル形式
  • Python バインディング: pip install protobuf(現在の 7.x には Python >=3.10 が必要です。6.x 以降も最小バージョンが引き上げられています。3.8/3.9 では対応する旧バージョンを指定してください。onnx には依存関係が含まれているため、通常は個別にインストールする必要はありません)
  • protoc コンパイルツール: Debian/Ubuntu apt install protobuf-compiler、Fedora dnf install protobuf-compiler、Arch pacman -S protobuf、macOS brew install protobuf
  • 確認: protoc --version;python3 -c "import google.protobuf; print(google.protobuf.__version__)"

操作步骤

按顺序执行,每步产物(模型哈希/结构摘要/权重清单)存档 sha256 + 路径([[re-ioc]] 证据链)。

  1. 模型格式识别(onnx / safetensors / pytorch pkl / tflite):

    file model.bin && sha256sum model.bin > model.sha256 && xxd model.bin | head -2
    
    • ONNX:无固定魔数——protobuf 流,首字段为 ir_version(tag 0x08 + 变长值;值随 ONNX 版本递增:IR 8 对应 ONNX 1.10,当前发布已到 IR 13——所以第二个字节不是常量 08,不能拿 08 08 12 当识别特征);producer_name(tag 0x12,长度前缀)按规范只 SHOULD 出现,不保证紧跟其后;onnx.checker.check_model 可验证合法性
    • Safetensors:前 8 字节 = 小端 u64 头长度(struct.unpack("<Q", data[:8])),随后是 JSON 头(张量名/形状/dtype/偏移)
    • PyTorch:file 显示 Zip archive(PK\x03\x04 头)——unzip -l model.pt 看条目(state_dict 含 data.pkl;torch.jit.script 含 data.pkl/constants.pkl/bytecode.pkl);老式纯 pkl 是裸 pickle 流(无 PK 头)——不直接 load,先 xxd/strings 粗看(坑 2)
    • TFLite:flatbuffers 流——无固定魔数,但字节 4–7 为文件标识符 TFL3(对应 schema 的 __model_identifier 字段),xxd/strings 可见;结构解析可用 Netron(支持 .tflite 可视化),权重提取分支思路同 onnx(flatbuffers 解析,超出本技能深度时标注"结构化 dump 为准")
    • 判定后按格式走对应分支;拿不准先 [[re-triage]] 初勘(熵/strings 特征)
  2. 结构解析(图/层/算子):

    netron model.onnx --no-browser          # 可视化(有图形界面再开浏览器)
    python3 - <<'PY'
    import onnx
    m = onnx.load("model.onnx")
    print("producer:", m.producer_name, m.producer_version)   # 框架/优化器指纹(见坑 3)
    g = m.graph
    print("inputs:", [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in g.input])
    print("nodes:", len(g.node), "initializers:", len(g.initializer), "outputs:", [o.name for o in g.output])
    for n in g.node[:20]:
        print(n.op_type, n.name, list(n.input), "->", list(n.output))
    PY
    
    • PyTorch 侧:torch.jit.load 得 ScriptModule 可打印 model.graph(TorchScript 结构);torch.load 的 state_dict 只有张量没有网络结构——结构在训练/推理脚本里,需配合源码还原(见跨域 [[re-script-deob]])
    • 关注点:算子序列(卷积/注意力等结构指纹)、输入输出张量形状、常量节点位置(权重藏在哪)
  3. 权重提取(张量 dump):

    mkdir -p weights
    # ONNX:逐 tensor 惰性落地(大模型见坑 1)
    # 注意:这不是"流式"——to_array() 对使用 external data 的 tensor 会先把该张量【完整】读进 ndarray,
    #       只是"一次一个张量"而不是"一次整个模型";显存/内存峰值取决于最大的那个张量
    python3 - <<'PY'
    import onnx, numpy as np
    m = onnx.load("model.onnx", load_external_data=False)     # 只载图结构,不载外部权重
    for init in m.graph.initializer:
        # base_dir 默认为空字符串——外部权重不在 CWD 时必须显式指定其所在目录
        arr = onnx.numpy_helper.to_array(init, base_dir=".")   # 该张量在此处被完整 materialize
        np.save(f"weights/{init.name.replace('/', '_')}.npy", arr)
        print(init.name, arr.shape, arr.dtype)
        del arr                                                # 及时释放,为下一张量让出内存
    PY
    # Safetensors:惰性按张量读取(不整载内存)
    python3 - <<'PY'
    from safetensors import safe_open
    with safe_open("model.safetensors", framework="numpy") as f:
        print(len(f.keys()), "tensors")
        for k in list(f.keys())[:10]:
            t = f.get_tensor(k); print(k, t.shape, t.dtype)
    PY
    # PyTorch state_dict(weights_only 安全加载,见坑 2)
    python3 - <<'PY'
    import torch
    sd = torch.load("model.pth", weights_only=True)
    for k, v in list(sd.items())[:10]:
        print(k, tuple(v.shape) if hasattr(v, "shape") else type(v))
    PY
    
    • 产出:权重清单(张量名/形状/dtype/数值摘要)+ npy 存档——这是水印检测与窃取判定的原料
  4. 模型水印/指纹检测(嵌入权重):

    • 权重级:逐张量统计(min/max/mean/std、直方图分桶)与疑似原版模型比对;水印常嵌在特定层(首层卷积 bias、归一化 scale、embedding 矩阵行),多为低比特位扰动——检查关键张量的低比特位模式与数值分布异常,而非精确相等(坑 4)
    • 指纹级:全权重 sha256 摘要、逐层张量 hash 序列;与候选原版逐层距离(L2/余弦)比对,输出"每层距离热点图"
    • 行为级:同一测试输入集跑两模型推理,比较 logits 与激活分布——重训练/蒸馏窃取者权重不同但行为接近
    • 产出:相似度矩阵 + 热点图,结论注明判定依据与阈值
  5. 模型窃取判定(架构相似度):

    python3 - <<'PY'
    import onnx
    from difflib import SequenceMatcher
    m1 = onnx.load("suspect.onnx"); m2 = onnx.load("original.onnx")
    s1 = [n.op_type for n in m1.graph.node]; s2 = [n.op_type for n in m2.graph.node]
    print("op-seq similarity:", SequenceMatcher(None, s1, s2).ratio())   # 算子序列相似度
    # 层对齐后逐层权重余弦/L2 距离:见步骤 4 的每层距离矩阵
    PY
    
    • 维度:算子序列(SequenceMatcher)、输入/输出形状、逐层权重余弦/L2、输出 logits 距离
    • 判定纪律:单一维度是弱证据(同架构不同训练=正常);多维度一致且权重分布高度接近(如逐层余弦 >0.99)才可主张窃取,结论标注各维度数值
    • 隐藏载荷检查:m.metadata_props(ONNX metadata 藏字符串/代码)、zip 条目中多余文件(坑 2 相关)、权重中形状/数值分布特异的异常张量

跨域联合

  • [[re-managed]]:本网关「识别运行时」识别到 AI 模型文件后固定调用本技能(模型是"代码在数据里"的托管域分支)
  • [[re-binary-core]]:模型内嵌代码、模型被打包进可执行文件(PyInstaller/pyarmor 打包的推理程序)——先二进制域拆解(格式解析/反编译),拆出的模型文件回本技能
  • [[re-sandbox]]:一切未知 pkl 的 load 默认隔离环境([[re-analyze/platform-tips]] 最高原则)
  • [[re-malware]]:恶意模型载荷(pickle 恶意代码、后门权重、投毒模型分发)的行为与情报侧
  • [[re-ioc]]:模型指纹(sha256/张量 hash/水印模式)进 IOC;[[re-triage]]:模型文件初勘入口与哈希存档
  • [[re-script-deob]]:PyTorch 推理/训练脚本还原(state_dict 无结构时的补全路径)
  • 引用 [[re-analyze/platform-tips]] 静态优先(大型样本)与沙箱最高原则分支
  • [[re-ai-attack]]:文件级证据 → 行为级验证(单向数据流)——本技能产出的文件侧结论(结构 / 权重 / 文件级 fingerprint)作为 re-ai-attack 行为一致性验证的输入。示例:拿到 suspect.pt → 本技能提取结构+权重生成 fingerprint → 转 re-ai-attack 做行为一致性验证(API 侧比对)
  • 只有 API(无文件)的评估直接进 [[re-ai-attack]],不经本技能

常见坑与陷阱

  • 大模型文件巨大(GB 级):现象——onnx.load/torch.load 吃满内存卡死,netron 打开超时,np.save 批量写盘满;原因——权重数 GB,一次性整体加载到内存;对策——分析前先 du -sh/sha256sum 存档;ONNX 用 onnx.load(..., load_external_data=False) 只载图结构,再逐 tensor 取值——注意 numpy_helper.to_array() 对该 tensor 是完整读出(不是流式),所以峰值取决于最大的张量,且用 external data 时要显式给 base_dir(默认空字符串,权重不在 CWD 时直接失败);Safetensors 用 safe_open 惰性按张量读;PyTorch 大模型 torch.load(..., mmap=True);处理对象是"结构摘要 + 定向张量",不是整个文件
  • pkl 反序列化风险(不要直接 torch.load 未知 pkl):现象——load 后进程反弹 shell/文件被删,或报诡异 AttributeError/ModuleNotFoundError;原因——pickle 协议可注入任意代码(__reduce__/__setstate__),torch.load 底层就是 pickle,恶意模型是投毒载荷载体;对策——安全提示:未知模型绝不直接 torch.load;先 unzip -l/xxd/strings 粗查(zip 头 PK vs 裸 pickle、条目有无可疑模块名),用 weights_only=True(PyTorch 2.6+ 默认)加载,需要全功能加载时在隔离环境([[re-sandbox]])执行;可转 safetensors 的样本直接转(纯数据无代码执行)
  • 图优化混淆层结构:现象——onnx-simplifier/TensorRT/onnxruntime 优化后的模型算子序列与训练态对不上(Conv+BN 融合成一个 Conv、常量折叠、名字全改),结构相似度误判;原因——优化器做算子融合/常量折叠,图结构与训练态不同,producer 字段会变;对策——先读 m.producer_name/m.producer_version 识别优化器与版本(融合 ConvBN 的特征:BN 层消失且 scale 并入 conv 权重);架构比较前先规范化算子序列(按算子类别抽象,忽略名字与常量差异);可用 onnx-simplifier/onnxruntime graph_optimization_level 对比优化前后 diff 还原原始层
  • 水印鲁棒性(剪枝后仍存):现象——精确值比对未命中就下"无水印"结论,或两个无关模型在个别层数值巧合相似被误判"窃取";原因——鲁棒水印经剪枝/量化/重训练后仍存活(设计目标),精确匹配必漏;正常模型同架构同数据权重分布相似,单层巧合是假阳性;对策——水印检测用"统计异常"(低比特位扰动/数值分布特异层)而非"精确相等",窃取判定用多维度证据 + 阈值(步骤 5 纪律),结论标注置信度与证据强度

コメント