Handle invalid PyTorch files (#19)

README.md

@@ -70,6 +70,12 @@ picklescan -l DEBUG -p downloads/pytorch_model.bin
 picklescan -l DEBUG -u https://huggingface.co/prajjwal1/bert-tiny/resolve/main/pytorch_model.bin
 ```
 
+Lint the code:
+```
+black src tests
+flake8 src tests --count --show-source
+```
+
 Publish the package to [PyPI](https://pypi.org/project/picklescan/): bump the package version in `setup.cfg` and create a GitHub release. This triggers the `publish` workflow.
 
 Alternative manual steps to publish the package:

src/picklescan/torch.py

@@ -74,8 +74,11 @@ def _is_zipfile(f) -> bool:
 
 
 def get_magic_number(data: IO[bytes]) -> Optional[int]:
-    for opcode, args, _pos in genops(data):
-        if "INT" in opcode.name or "LONG" in opcode.name:
-            data.seek(0)
-            return int(args)
+    try:
+        for opcode, args, _pos in genops(data):
+            if "INT" in opcode.name or "LONG" in opcode.name:
+                data.seek(0)
+                return int(args)
+    except ValueError:
+        return None
     return None

tests/data/bad_pytorch.pt

@@ -0,0 +1,2 @@
+�PNG
+

tests/test_scanner.py

@@ -217,6 +217,9 @@ def initialize_pickle_files():
 
     initialize_numpy_file(f"{_root_path}/data/object_array.npy")
 
+    # Fake PyTorch file (PNG file format) simulating https://huggingface.co/RectalWorm/loras_new/blob/main/Owl_Mage_no_background.pt
+    initialize_data_file(f"{_root_path}/data/bad_pytorch.pt", b"\211PNG\r\n\032\n")
+
 
 initialize_pickle_files()
 
@@ -410,6 +413,11 @@ def test_scan_file_path():
         scan_file_path(f"{_root_path}/data/malicious9.pkl"), malicious9
     )
 
+    bad_pytorch = ScanResult([], 0, 0, 0, True)
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/bad_pytorch.pt"), bad_pytorch
+    )
+
 
 def test_scan_directory_path():
     sr = ScanResult(