● Zion Boggan
repos/Cyclops V2/brain/semgrep-rules/ml-unsafe-load.yaml
zionboggan.com ↗
94 lines · yaml
History for this file →
1
rules:
2
 
3
  - id: cyc-ml-torch-load
4
    languages: [python]
5
    severity: ERROR
6
    message: >
7
      torch.load uses pickle by default; a crafted checkpoint runs arbitrary code on load. Even
8
      weights_only=False (default before torch 2.6) is exploitable. IMPACT GATE: Family 1 (malicious
9
      model file). Safe form: torch.load(..., weights_only=True).
10
    metadata:
11
      cwe: ["CWE-502: Deserialization of Untrusted Data"]
12
      confidence: HIGH
13
      archetype: [ml-model-loaders]
14
      impact: rce
15
    patterns:
16
      - pattern-either:
17
          - pattern: torch.load($P)
18
          - pattern: torch.load($P, weights_only=False, ...)
19
          - pattern: torch.jit.load($P, ...)
20
      - pattern-not: torch.load($P, weights_only=True, ...)
21
 
22
  - id: cyc-ml-joblib-numpy-pandas
23
    languages: [python]
24
    severity: ERROR
25
    message: >
26
      joblib.load / numpy.load(allow_pickle=True) / pandas.read_pickle deserialize via pickle ->
27
      arbitrary code from a malicious artifact. IMPACT GATE: Family 1.
28
    metadata:
29
      cwe: ["CWE-502: Deserialization of Untrusted Data"]
30
      confidence: HIGH
31
      archetype: [ml-model-loaders]
32
      impact: rce
33
    patterns:
34
      - pattern-either:
35
          - pattern: joblib.load(...)
36
          - pattern: numpy.load(..., allow_pickle=True, ...)
37
          - pattern: np.load(..., allow_pickle=True, ...)
38
          - pattern: pandas.read_pickle(...)
39
          - pattern: pd.read_pickle(...)
40
 
41
  - id: cyc-ml-keras-tf
42
    languages: [python]
43
    severity: ERROR
44
    message: >
45
      keras.models.load_model on an .h5/.keras file can execute a Lambda layer's pickled code, and
46
      load_model from an untrusted SavedModel can run code. IMPACT GATE: Family 1 (malicious model).
47
    metadata:
48
      cwe: ["CWE-502: Deserialization of Untrusted Data"]
49
      confidence: MEDIUM
50
      archetype: [ml-model-loaders]
51
      impact: rce
52
    patterns:
53
      - pattern-either:
54
          - pattern: keras.models.load_model(...)
55
          - pattern: tf.keras.models.load_model(...)
56
          - pattern: load_model(...)
57
 
58
  - id: cyc-ml-hf-trust-remote-code
59
    languages: [python]
60
    severity: ERROR
61
    message: >
62
      transformers from_pretrained / pipeline with trust_remote_code=True downloads and EXECUTES the
63
      model repo's custom Python (modeling_*.py) -> RCE from a malicious Hub model. IMPACT GATE:
64
      Family 1, the AI supply-chain RCE class.
65
    metadata:
66
      cwe: ["CWE-494: Download of Code Without Integrity Check", "CWE-829: Inclusion of Untrusted Functionality"]
67
      confidence: HIGH
68
      archetype: [ml-model-loaders, ai-agent-mcp, cicd-supplychain]
69
      impact: rce
70
    patterns:
71
      - pattern-either:
72
          - pattern: $M.from_pretrained(..., trust_remote_code=True, ...)
73
          - pattern: transformers.pipeline(..., trust_remote_code=True, ...)
74
          - pattern: AutoModel.from_pretrained(..., trust_remote_code=True, ...)
75
          - pattern: AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True, ...)
76
          - pattern: AutoTokenizer.from_pretrained(..., trust_remote_code=True, ...)
77
 
78
  - id: cyc-ml-pickle-direct
79
    languages: [python]
80
    severity: ERROR
81
    message: >
82
      Direct pickle/dill/cloudpickle load of model bytes -> RCE. Common in custom checkpoint loaders.
83
      IMPACT GATE: Family 1.
84
    metadata:
85
      cwe: ["CWE-502: Deserialization of Untrusted Data"]
86
      confidence: HIGH
87
      archetype: [ml-model-loaders]
88
      impact: rce
89
    patterns:
90
      - pattern-either:
91
          - pattern: pickle.load(...)
92
          - pattern: pickle.loads(...)
93
          - pattern: dill.load(...)
94
          - pattern: cloudpickle.load(...)