Limit B614 to torch.load deserializers (#1348)

* Limit B614 to torch.load deserializers

Avoids false positives for torch.*.load helpers such as
torch.utils.cpp_extension.load while preserving checks for
torch.load and torch.serialization.load.
Updated docstrings and example to reflect expected behavior.

Resolves: #1343

* Update examples/pytorch_load.py

---------

Co-authored-by: Eric Brown <ericwb@users.noreply.github.com>

Author: dibussoc

bandit/plugins/pytorch_load.py

@@ -6,9 +6,10 @@
 B614: Test for unsafe PyTorch load
 ==================================
 
-This plugin checks for unsafe use of `torch.load`. Using `torch.load` with
-untrusted data can lead to arbitrary code execution. There are two safe
-alternatives:
+This plugin checks for unsafe use of `torch.load` and
+`torch.serialization.load`. Using `torch.load` or
+`torch.serialization.load` with untrusted data can lead to arbitrary
+code execution. There are two safe alternatives:
 
 1. Use `torch.load` with `weights_only=True` where only tensor data is
    extracted, and no arbitrary Python objects are deserialized
@@ -24,7 +25,7 @@
 
         >> Issue: Use of unsafe PyTorch load
         Severity: Medium   Confidence: High
-        CWE: CWE-94 (https://cwe.mitre.org/data/definitions/94.html)
+        CWE: CWE-502 (https://cwe.mitre.org/data/definitions/502.html)
         Location: examples/pytorch_load_save.py:8
         7    loaded_model.load_state_dict(torch.load('model_weights.pth'))
         8    another_model.load_state_dict(torch.load('model_weights.pth',
@@ -34,7 +35,7 @@
 
 .. seealso::
 
-     - https://cwe.mitre.org/data/definitions/94.html
+     - https://cwe.mitre.org/data/definitions/502.html
      - https://pytorch.org/docs/stable/generated/torch.load.html#torch.load
      - https://github.com/huggingface/safetensors
 
@@ -50,23 +51,18 @@
 @test.test_id("B614")
 def pytorch_load(context):
     """
-    This plugin checks for unsafe use of `torch.load`. Using `torch.load`
-    with untrusted data can lead to arbitrary code execution. The safe
-    alternative is to use `weights_only=True` or the safetensors library.
+    This plugin checks for unsafe use of `torch.load` and
+    `torch.serialization.load`. Using `torch.load` or
+    `torch.serialization.load` with untrusted data can lead to
+    arbitrary code execution. The safe alternative is to use
+    `weights_only=True` or the safetensors library.
     """
     imported = context.is_module_imported_exact("torch")
     qualname = context.call_function_name_qual
     if not imported and isinstance(qualname, str):
         return
 
-    qualname_list = qualname.split(".")
-    func = qualname_list[-1]
-    if all(
-        [
-            "torch" in qualname_list,
-            func == "load",
-        ]
-    ):
+    if qualname in {"torch.load", "torch.serialization.load"}:
         # For torch.load, check if weights_only=True is specified
         weights_only = context.get_call_arg_value("weights_only")
         if weights_only == "True" or weights_only is True:

examples/pytorch_load.py

@@ -24,3 +24,7 @@
 # Example of loading with both map_location and weights_only=True (should NOT trigger B614)
 safe_cpu_model = models.resnet18()
 safe_cpu_model.load_state_dict(torch.load('model_weights.pth', map_location='cpu', weights_only=True))
+
+# Example of a torch.*.load call that should NOT trigger B614
+# Only pickle deserializers should trigger B614
+torch.utils.cpp_extension.load(name="example_ext", sources=[])