Skip to content

Commit 0e64ac2

Browse files
committed
1 parent b27724a commit 0e64ac2

File tree

3 files changed

+17
-11
lines changed

3 files changed

+17
-11
lines changed

src/picklescan/scanner.py

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -493,19 +493,23 @@ def scan_bytes(data: IO[bytes], file_id, file_ext: Optional[str] = None) -> Scan
493493
try:
494494
return scan_pytorch(data, file_id)
495495
except InvalidMagicError as e:
496-
_log.error(f"ERROR: Invalid magic number for file {e}")
497-
return ScanResult([], scan_err=True)
498-
elif file_ext is not None and file_ext in _numpy_file_extensions:
496+
_log.warning(
497+
f"WARNING: Invalid PyTorch magic number for file {e}. Trying to scan as non-PyTorch file.",
498+
exc_info=_log.isEnabledFor(logging.DEBUG),
499+
)
500+
data.seek(0)
501+
502+
if file_ext is not None and file_ext in _numpy_file_extensions:
499503
return scan_numpy(data, file_id)
504+
505+
is_zip = zipfile.is_zipfile(data)
506+
data.seek(0)
507+
if is_zip:
508+
return scan_zip_bytes(data, file_id)
509+
elif _is_7z_file(data):
510+
return scan_7z_bytes(data, file_id)
500511
else:
501-
is_zip = zipfile.is_zipfile(data)
502-
data.seek(0)
503-
if is_zip:
504-
return scan_zip_bytes(data, file_id)
505-
elif _is_7z_file(data):
506-
return scan_7z_bytes(data, file_id)
507-
else:
508-
return scan_pickle_bytes(data, file_id)
512+
return scan_pickle_bytes(data, file_id)
509513

510514

511515
def scan_huggingface_model(repo_id):
92 Bytes
Binary file not shown.

tests/test_scanner.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -765,6 +765,7 @@ def initialize_pickle_files():
765765
initialize_pickle_file_from_reduce("GHSA-9w88-8rmg-7g2p.pkl", reduce_GHSA_9w88_8rmg_7g2p)
766766
initialize_pickle_file_from_reduce("GHSA-49gj-c84q-6qm9.pkl", reduce_GHSA_49gj_c84q_6qm9)
767767
initialize_pickle_file_from_reduce("GHSA-q77w-mwjj-7mqx.pkl", reduce_GHSA_q77w_mwjj_7mqx)
768+
initialize_pickle_file_from_reduce("GHSA-jgw4-cr84-mqxg.bin", reduce_GHSA_q77w_mwjj_7mqx)
768769

769770

770771
initialize_pickle_files()
@@ -1043,6 +1044,7 @@ def test_scan_file_path():
10431044
assert_scan("GHSA-9w88-8rmg-7g2p.pkl", [Global("cProfile", "runctx", SafetyLevel.Dangerous)])
10441045
assert_scan("GHSA-49gj-c84q-6qm9.pkl", [Global("cProfile", "run", SafetyLevel.Dangerous)])
10451046
assert_scan("GHSA-q77w-mwjj-7mqx.pkl", [Global("asyncio.unix_events", "_UnixSubprocessTransport._start", SafetyLevel.Dangerous)])
1047+
assert_scan("GHSA-jgw4-cr84-mqxg.bin", [Global("asyncio.unix_events", "_UnixSubprocessTransport._start", SafetyLevel.Dangerous)])
10461048
assert_scan("malicious1_crc.zip", [Global("builtins", name="eval", safety=SafetyLevel.Dangerous)])
10471049

10481050

0 commit comments

Comments
 (0)