zxix / stable-diffusion-pickle-scanner

187 stars 19 forks source link

BUG: pickle_scan stops if it encounters a non-picke file #14

Open phizch opened 1 year ago

phizch commented 1 year ago

When I tried scanning my downloads directory the scanning stopped before all files had been scanned.

The reason is that result = torch.load(path.as_posix(), pickle_module=pickle_inspector.pickle) threw an exception because one of the files in my download directory wasn't a valid checkpoint. (I think it was a misnamed safetensors file).

Python does produce an error message, but it's not saved in the scan_outputs.txt file, and it's not picked up on by the cmd script when it's reporting on potentially malicious files.

How to reproduce:

  1. Create an empty file named "0.bin" in one of the directories to be scanned.
  2. Scan

Expected behavior:

All files should be scanned.

Actual behavior:

The scan stops after trying to scan the empty file with an error like this:

...L:/StableDiffusion/Downloads/0.bin
Traceback (most recent call last):
  File "L:\StableDiffusion\stable-diffusion-pickle-scanner-main\pickle_scan.py", line 24, in <module>
    result = torch.load(path.as_posix(), pickle_module=pickle_inspector.pickle)
  File "L:\StableDiffusion\stable-diffusion-webui\venv\lib\site-packages\torch\serialization.py", line 713, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "L:\StableDiffusion\stable-diffusion-webui\venv\lib\site-packages\torch\serialization.py", line 920, in _legacy_load
    magic_number = pickle_module.load(f, **pickle_load_args)

Fix:

Do the scanning in try/except blocks.

for path in BASE_DIR.glob(r'**/*'):
  if path.suffix in EXTENSIONS:
    print("")
    print("..." + path.as_posix())
    try:
      result = torch.load(path.as_posix(), pickle_module=pickle_inspector.pickle)
      ..
    except:
      print("NOT A VAILD CHECKPOINT FILE")