Bump version to 0.0.22 and fix tests (#37)

Co-authored-by: Matthieu Maitre <mmaitre@microsoft.com>

Author: mmaitre314

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = picklescan
-version = 0.0.21
+version = 0.0.22
 author = Matthieu Maitre
 author_email = mmaitre314@users.noreply.github.com
 description = Security scanner detecting Python Pickle files performing suspicious actions

src/picklescan/scanner.py

@@ -130,11 +130,11 @@ def __str__(self) -> str:
     "bdb": "*",
     "pdb": "*",
     "asyncio": "*",
-    "pydoc": "pipepager", # pydoc.pipepager('help','echo pwned')
+    "pydoc": "pipepager",  # pydoc.pipepager('help','echo pwned')
     "venv": "*",
-    "torch.serialization": "load", # pickle could be used to load a different file
-    "functools": "partial", # functools.partial(os.system, "echo pwned")
-    "torch._inductor.codecache": "compile_file", # compile_file('', '', ['sh', '-c','$(echo pwned)'])
+    "torch.serialization": "load",  # pickle could be used to load a different file
+    "functools": "partial",  # functools.partial(os.system, "echo pwned")
+    "torch._inductor.codecache": "compile_file",  # compile_file('', '', ['sh', '-c','$(echo pwned)'])
     "pip": "*",
 }
 
@@ -159,7 +159,14 @@ def __str__(self) -> str:
 _pytorch_file_extensions = {".bin", ".pt", ".pth", ".ckpt"}
 _pickle_file_extensions = {".pkl", ".pickle", ".joblib", ".dat", ".data"}
 _zip_file_extensions = {".zip", ".npz", ".7z"}
-_pickle_magic_bytes = {b"\x80\x00", b"\x80\x01", b"\x80\x02", b"\x80\x03", b"\x80\x04", b"\x80\x05"}
+_pickle_magic_bytes = {
+    b"\x80\x00",
+    b"\x80\x01",
+    b"\x80\x02",
+    b"\x80\x03",
+    b"\x80\x04",
+    b"\x80\x05",
+}
 
 
 def _is_7z_file(f: IO[bytes]) -> bool:
@@ -355,6 +362,7 @@ def scan_7z_bytes(data: IO[bytes], file_id) -> ScanResult:
 
             return result
 
+
 def get_magic_bytes_from_zipfile(zip: zipfile.ZipFile, num_bytes=8):
     magic_bytes = {}
     for file_info in zip.infolist():
@@ -374,11 +382,15 @@ def scan_zip_bytes(data: IO[bytes], file_id) -> ScanResult:
         for file_name in file_names:
             magic_number = magic_bytes.get(file_name, b"")
             file_ext = os.path.splitext(file_name)[1]
-            if file_ext in _pickle_file_extensions or any(magic_number.startswith(mn) for mn in _pickle_magic_bytes):
+            if file_ext in _pickle_file_extensions or any(
+                magic_number.startswith(mn) for mn in _pickle_magic_bytes
+            ):
                 _log.debug("Scanning file %s in zip archive %s", file_name, file_id)
                 with zip.open(file_name, "r") as file:
                     result.merge(scan_pickle_bytes(file, f"{file_id}:{file_name}"))
-            elif file_ext in _numpy_file_extensions or magic_number.startswith(b"\x93NUMPY"):
+            elif file_ext in _numpy_file_extensions or magic_number.startswith(
+                b"\x93NUMPY"
+            ):
                 _log.debug("Scanning file %s in zip archive %s", file_name, file_id)
                 with zip.open(file_name, "r") as file:
                     result.merge(scan_numpy(file, f"{file_id}:{file_name}"))

tests/test_scanner.py

@@ -6,16 +6,19 @@
 import os
 import pathlib
 import pickle
+import pip
 import py7zr
+import pydoc
 import pytest
 import requests
 import runpy
 import socket
 import subprocess
 import sys
-from unittest import TestCase
+import venv
 import zipfile
-import pip
+from functools import partial
+from unittest import TestCase
 
 from picklescan.cli import main
 from picklescan.scanner import (
@@ -33,11 +36,14 @@
     scan_numpy,
     scan_pytorch,
 )
-import torch
-import torch._inductor.codecache as codecache
-import pydoc
-import venv
-from functools import partial
+
+try:
+    import torch
+    import torch._inductor.codecache as codecache
+except ImportError:
+    # If PyTorch test files need to be regenerated, run 'pip install torch==2.6.0' first
+    torch = None
+    codecache = None
 
 _root_path = os.path.dirname(__file__)
 
@@ -105,20 +111,16 @@ def __reduce__(self):
             'import os\nos.system("whoami")',
         )
 
+
 class Malicious17:
     def __reduce__(self):
-        return codecache.compile_file, (
-            '', '', [
-                'sh',
-                '-c','$(echo "pwned")'
-            ]
-        )
+        return codecache.compile_file, ("", "", ["sh", "-c", '$(echo "pwned")'])
+
 
 class Malicious18:
     def __reduce__(self):
-        return pydoc.pipepager, (
-            '', 'echo "pwned"'
-        )
+        return pydoc.pipepager, ("", 'echo "pwned"')
+
 
 class Malicious19:
     def __init__(self, path, **kwargs):
@@ -128,11 +130,10 @@ def __init__(self, path, **kwargs):
     def __reduce__(self):
         return partial(torch.load, self.path, **self.kwargs), ()
 
+
 class Malicious20:
     def __reduce__(self):
-        return venv.create, (
-            'venv', False, False, True, False, "$(echo pwned)"
-        )
+        return venv.create, ("venv", False, False, True, False, "$(echo pwned)")
 
 
 class Malicious16:
@@ -469,8 +470,13 @@ def initialize_pickle_files():
     initialize_pickle_file(f"{_root_path}/data/malicious18.pkl", Malicious18(), 4)
 
     # This exploit serializes kwargs and passes them into a torch.load call
-    initialize_pickle_file(f"{_root_path}/data/malicious19.pkl",
-                           Malicious19("some_other_model.bin", pickle_file='config.json', weights_only=False), 4)
+    initialize_pickle_file(
+        f"{_root_path}/data/malicious19.pkl",
+        Malicious19(
+            "some_other_model.bin", pickle_file="config.json", weights_only=False
+        ),
+        4,
+    )
 
     initialize_pickle_file(f"{_root_path}/data/malicious20.pkl", Malicious20(), 4)
     initialize_7z_file(
@@ -486,7 +492,7 @@ def initialize_pickle_files():
 
     initialize_zip_file(
         f"{_root_path}/data/malicious1_wrong_ext.zip",
-        "data.txt", # Pickle file with a non-standard extension
+        "data.txt",  # Pickle file with a non-standard extension
         pickle.dumps(Malicious1(), protocol=4),
     )
 
@@ -640,9 +646,7 @@ def test_scan_file_path():
     compare_scan_results(
         scan_file_path(f"{_root_path}/data/malicious1.zip"), malicious1
     )
-    compare_scan_results(
-        scan_file_path(f"{_root_path}/data/malicious1.7z"), malicious1
-    )
+    compare_scan_results(scan_file_path(f"{_root_path}/data/malicious1.7z"), malicious1)
     compare_scan_results(
         scan_file_path(f"{_root_path}/data/malicious1_wrong_ext.zip"), malicious1
     )
@@ -830,10 +834,11 @@ def test_scan_directory_path():
             Global("torch.serialization", "load", SafetyLevel.Dangerous),
             Global("functools", "partial", SafetyLevel.Dangerous),
             Global("pip", "main", SafetyLevel.Dangerous),
+            Global("builtins", "eval", SafetyLevel.Dangerous),
         ],
-        scanned_files=37,
-        issues_count=38,
-        infected_files=31,
+        scanned_files=38,
+        issues_count=39,
+        infected_files=33,
         scan_err=True,
     )
     compare_scan_results(scan_directory_path(f"{_root_path}/data/"), sr)