zxix / stable-diffusion-pickle-scanner

187 stars 20 forks source link

Enhancement: Simple batch file to run scanner #2

Open TheEliteGeek opened 1 year ago

TheEliteGeek commented 1 year ago

For those who want a quicker way to use this instead of manually launching cmd and whatever each time, you can use a simple batch file.

Create a new text file (name it PickleScan or whatever, paste the following, save it, and change it to a .bat file.

Note: You MUST change the venv python path to YOUR venv python path.

@echo off

echo "Scanning...Please wait a moment..."

"F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe"  pickle_scan.py models > scan_output.txt

type scan_output.txt

pause

Running the .bat will launch the scanner and display the results in the cmd window. It'll still save the output to "scan_output.txt"

Someone might be able to make this a bit more fancy and have picklescan be what you run first, have it check for new models (comparing it to the previous scan output), scan those, and if all good, launch webui-user.bat.

I could probably do this but I'm too lazy right now...

Ruffy314 commented 1 year ago

Building on this

Checks if the venv path and model directory was set at top of script Allows to specify one additional folder to check (optional) Options to display the scan results in console or open in notepad Checks the output file for the words SCAN FAILED for quick check if anything was suspicious

@echo off

echo scan started on %date% %time% > scan_output.txt

:SETUP
rem scans SD models for malicious pickling
set VENV_PATH="F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe"
set SD_FOLDER="F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion"
set DOWNLOAD_FOLDER="F:\Whatever\Other\Folder"

rem how result details should be displayed ("yes" or "no"):
set SHOW_RESULT_IN_CONSOLE="yes"
set OPEN_RESULT_IN_NOTEPAD="yes"

:SCANNING
rem check if VENV_PATH was set
if %VENV_PATH% equ "F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe" (
echo ##### ERROR please set your VENV_PATH #####
goto EXIT
)

echo "Scanning...Please wait a moment..."

rem check if SD_FOLDER was set
if %SD_FOLDER% equ "F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion" (
echo ##### ERROR please set your SD_FOLDER #####
goto EXIT
)

echo "step 1: SD models folder"
echo ####################################################################### >> scan_output.txt
echo ##### scanning SD model folder "~~~webui\models\Stable-diffusion" ##### >> scan_output.txt
echo ####################################################################### >> scan_output.txt
%VENV_PATH% pickle_scan.py %SD_FOLDER% >> scan_output.txt

rem check if download folder was set, if not just skip instead of errorlevel
if %DOWNLOAD_FOLDER% equ "F:\Whatever\Other\Folder" (
echo "No download folder specified"
goto DISPLAY_RESULT
)
echo "step 2: download folder"
echo ##################################### >> scan_output.txt
echo ##### scanning download folder  ##### >> scan_output.txt
echo ##################################### >> scan_output.txt
%VENV_PATH% pickle_scan.py %DOWNLOAD_FOLDER% >> scan_output.txt

:DISPLAY_RESULT
if %SHOW_RESULT_IN_CONSOLE% equ "yes" (type scan_output.txt)
if %OPEN_RESULT_IN_NOTEPAD% equ "yes" (start notepad scan_output.txt)

echo "Number of failed scans (potentially malicious files):"
find /c "SCAN FAILED" scan_output.txt

:EXIT
pause
TheEliteGeek commented 1 year ago

Building on this

First off, you're incredible! This is amazing and I really appreciate it! Secondly, I'm building on your building on this by offering 3 options:

It seems to work (for me at least) but it might actually be broken ¯\(ツ)

If I wrote this in PowerShell (I never really learned batch), I would have been able to do more fancy things such as automatically moving checked models in the Downloads folder into the Models folder... Unfortunately PowerShell console doesn't seem to like running python scripts in venv (See https://github.com/zxix/stable-diffusion-pickle-scanner/issues/1) and I'm again too lazy to figure out why...

:: Updated by Ruffy314 on Github
:: Updated by MisterThief117 on Github

@echo off

echo scan started on %date% %time% > scan_output.txt

:SETUP
rem scans SD models for malicious pickling
set VENV_PATH="F:\Digital Art\AI Media\Stable Diffusion New\stable-diffusion-webui\venv\Scripts\Python.exe"
set SD_FOLDER="F:\Digital Art\AI Media\Stable Diffusion New\stable-diffusion-webui\models"
set DOWNLOAD_FOLDER="F:\Digital Art\AI Media\AI Models"

rem how result details should be displayed ("yes" or "no"):
set SHOW_RESULT_IN_CONSOLE="yes"
set OPEN_RESULT_IN_NOTEPAD="yes"

echo ------Welcome to PickleScanner EZ Launcher!------
Echo ----------Now with pointless complexity!----------
echo:
Echo Your VENV Path is currently: %VENV_PATH%
Echo Your MODELS Path is currently: %SD_FOLDER%

if %DOWNLOAD_FOLDER% equ "F:\Whatever\Other\Folder" (
echo Your DOWNLOAD folder is not set. This is optional.

GOTO BEGIN
) ELSE (
Echo Your DOWNLOAD Path is currently: %DOWNLOAD_FOLDER%
GOTO BEGIN
)

:BEGIN
echo:
echo Select a task:
echo =============
echo -
echo 1) Scan both MODELS and DOWNLOADS folder
echo 2) Scan DOWNLOADS folder only
echo 3) Scan MODELS folder only

echo -
set /p op=Type option:
if "%op%"=="1" goto op1
if "%op%"=="2" goto op2
if "%op%"=="3" goto op3

:op1 Scan MODELS and DOWNLOADS
:op3 Scan MODELS
rem check if VENV_PATH was set
if %VENV_PATH% equ "F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe" (
echo ----- ERROR please set your VENV_PATH -----
goto EXIT
)

rem check if SD_FOLDER was set
if %SD_FOLDER% equ "F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion" (
echo ----- ERROR please set your SD_FOLDER -----
goto EXIT
)

echo "Scanning...Please wait a moment..."
ECHO:
echo "Step 1: MODELS folder"
echo ------------------------------------------------------- >> scan_output.txt
echo ----- scanning SD model folder "\webui\models\" ------ >> scan_output.txt
echo ------------------------------------------------------- >> scan_output.txt
ECHO:

%VENV_PATH% pickle_scan.py %SD_FOLDER% >> scan_output.txt

if %SHOW_RESULT_IN_CONSOLE% equ "yes" type scan_output.txt

if %op% equ 3 GOTO DISPLAY_RESULT
if %op% equ 1 GOTO op2

:op2 Scan DOWNLOADS Folder
rem check if download folder was set, if not just skip instead of errorlevel
if %DOWNLOAD_FOLDER% equ "F:\Whatever\Other\Folder" (
echo "No download folder specified"
goto DISPLAY_RESULT
)

ECHO:
echo "Step 2: DOWNLOADS folder"
echo ------------------------------------- >> scan_output.txt
echo ----- Scanning download folder  ----- >> scan_output.txt
echo ------------------------------------- >> scan_output.txt
ECHO:
%VENV_PATH% pickle_scan.py %DOWNLOAD_FOLDER% >> scan_output.txt
if %SHOW_RESULT_IN_CONSOLE% equ "yes" (type scan_output.txt)
goto DISPLAY_RESULT

:DISPLAY_RESULT
ECHO:
echo -------------------------------------
echo -----Summary-----
if %OPEN_RESULT_IN_NOTEPAD% equ "yes" (start notepad scan_output.txt)
find /c "SCAN FAILED" scan_output.txt > Nul
if %errorlevel% equ 0 (
echo Possible Malicous Pickle Found. See logs for details
GOTO EXIT
) else (
ECHO No pickles found.
GOTO EXIT
)

GOTO EXIT

:EXIT
pause
zxix commented 1 year ago

This is great, definitely worth a PR!

Ruffy314 commented 1 year ago

We may have reached a point where it would be better use the .bat solely for setting variables and starting a separat python script for the logic 🙃

TheEliteGeek commented 1 year ago

I also used MidJourney to create this...logo? I initially made it to use as an icon for the bat file (requires making a shortcut of the bat and assigning the custom icon to the shortcut...), but I'll share it here for zxix (or anyone else) to use.

To use as an icon in Windows, it'll need to be converted into a .ico file using some converter.

PickleScan

Honestly the pickle looks more like a green turd than a pickle 🤣

Masqued-Man commented 1 year ago

Building on this

Checks if the venv path and model directory was set at top of script Allows to specify one additional folder to check (optional) Options to display the scan results in console or open in notepad Checks the output file for the words SCAN FAILED for quick check if anything was suspicious

@echo off

echo scan started on %date% %time% > scan_output.txt

:SETUP
rem scans SD models for malicious pickling
set VENV_PATH="F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe"
set SD_FOLDER="F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion"
set DOWNLOAD_FOLDER="F:\Whatever\Other\Folder"

rem how result details should be displayed ("yes" or "no"):
set SHOW_RESULT_IN_CONSOLE="yes"
set OPEN_RESULT_IN_NOTEPAD="yes"

:SCANNING
rem check if VENV_PATH was set
if %VENV_PATH% equ "F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe" (
echo ##### ERROR please set your VENV_PATH #####
goto EXIT
)

echo "Scanning...Please wait a moment..."

rem check if SD_FOLDER was set
if %SD_FOLDER% equ "F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion" (
echo ##### ERROR please set your SD_FOLDER #####
goto EXIT
)

echo "step 1: SD models folder"
echo ####################################################################### >> scan_output.txt
echo ##### scanning SD model folder "~~~webui\models\Stable-diffusion" ##### >> scan_output.txt
echo ####################################################################### >> scan_output.txt
%VENV_PATH% pickle_scan.py %SD_FOLDER% >> scan_output.txt

rem check if download folder was set, if not just skip instead of errorlevel
if %DOWNLOAD_FOLDER% equ "F:\Whatever\Other\Folder" (
echo "No download folder specified"
goto DISPLAY_RESULT
)
echo "step 2: download folder"
echo ##################################### >> scan_output.txt
echo ##### scanning download folder  ##### >> scan_output.txt
echo ##################################### >> scan_output.txt
%VENV_PATH% pickle_scan.py %DOWNLOAD_FOLDER% >> scan_output.txt

:DISPLAY_RESULT
if %SHOW_RESULT_IN_CONSOLE% equ "yes" (type scan_output.txt)
if %OPEN_RESULT_IN_NOTEPAD% equ "yes" (start notepad scan_output.txt)

echo "Number of failed scans (potentially malicious files):"
find /c "SCAN FAILED" scan_output.txt

:EXIT
pause

So I'm trying to use this method, but it gives me the error saying that the venv path isn't set, even though I seem to have it set properly. Here's what I have in the file:

@echo off

echo scan started on %date% %time% > scan_output.txt

:SETUP rem scans SD models for malicious pickling set VENV_PATH="F:\Stable_Diffusion_2.0\stable-diffusion-webui\venv\Scripts\python.exe" set SD_FOLDER="F:\Stable_Diffusion_2.0\stable-diffusion-webui\models\Stable-diffusion" set DOWNLOAD_FOLDER="F:\Whatever\Other\Folder"

rem how result details should be displayed ("yes" or "no"): set SHOW_RESULT_IN_CONSOLE="yes" set OPEN_RESULT_IN_NOTEPAD="yes"

:SCANNING rem check if VENV_PATH was set if %VENV_PATH% equ "F:\Stable_Diffusion_2.0\stable-diffusion-webui\venv\Scripts\python.exe" ( echo ##### ERROR please set your VENV_PATH ##### goto EXIT )

echo "Scanning...Please wait a moment..."

rem check if SD_FOLDER was set if %SD_FOLDER% equ "F:\Stable_Diffusion_2.0\stable-diffusion-webui\models\Stable-diffusion" ( echo ##### ERROR please set your SD_FOLDER ##### goto EXIT )

echo "step 1: SD models folder" echo ####################################################################### >> scan_output.txt echo ##### scanning SD model folder "~~~webui\models\Stable-diffusion" ##### >> scan_output.txt echo ####################################################################### >> scan_output.txt %VENV_PATH% pickle_scan.py %SD_FOLDER% >> scan_output.txt

rem check if download folder was set, if not just skip instead of errorlevel if %DOWNLOAD_FOLDER% equ "F:\Whatever\Other\Folder" ( echo "No download folder specified" goto DISPLAY_RESULT ) echo "step 2: download folder" echo ##################################### >> scan_output.txt echo ##### scanning download folder ##### >> scan_output.txt echo ##################################### >> scan_output.txt %VENV_PATH% pickle_scan.py %DOWNLOAD_FOLDER% >> scan_output.txt

:DISPLAY_RESULT if %SHOW_RESULT_IN_CONSOLE% equ "yes" (type scan_output.txt) if %OPEN_RESULT_IN_NOTEPAD% equ "yes" (start notepad scan_output.txt)

echo "Number of failed scans (potentially malicious files):" find /c "SCAN FAILED" scan_output.txt

:EXIT pause

But when I run it, it gives me the "ERROR please set your VENV_PATH" message. Am I doing something wrong?

saraprc-dev commented 1 year ago

I also used MidJourney to create this...logo? I initially made it to use as an icon for the bat file (requires making a shortcut of the bat and assigning the custom icon to the shortcut...), but I'll share it here for zxix (or anyone else) to use.

To use as an icon in Windows, it'll need to be converted into a .ico file using some converter.

PickleScan

Honestly the pickle looks more like a green turd than a pickle 🤣

Thanks for this. Goes well with my current setup. image

lopho commented 1 year ago

https://github.com/lopho/pickle_inspector can scan lists of pickles.

python scan_pickle.py --preset stable_diffusion_v1 --in modelA.ckpt modelB.ckpt modelC.ckpt sussy.ckpt sd-v1-4.ckpt

it works with globs as well:

python scan_pickle.py --preset stable_diffusion_v1 --in models/*.ckpt
Scanning file(s): ['models/sd-v1-4.ckpt', 'models/sus.ckpt', 'models/sussy.ckpt']
Using white list: ['collections.OrderedDict', 'torch._utils._rebuild_tensor_v2', 'torch.HalfStorage', 'torch.FloatStorage', 'torch.IntStorage', 'torch.LongStorage', 'pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint', 'numpy.core.multiarray.scalar', 'numpy.dtype', '_codecs.encode']
Reading models/sd-v1-4.ckpt
Found pickle in zip: archive/data.pkl
Scanning: archive/data.pkl
found: torch._utils._rebuild_tensor_v2
found: torch.FloatStorage
found: collections.OrderedDict
found: torch.IntStorage
found: torch.LongStorage
found: pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint
found: numpy.core.multiarray.scalar
found: numpy.dtype
found: _codecs.encode
Scan for models/sd-v1-4.ckpt PASSED ✅
Reading models/sus.ckpt
Found pickle in zip: archive/data.pkl
Scanning: archive/data.pkl
found: torch._utils._rebuild_tensor_v2
found: torch.FloatStorage
found: collections.OrderedDict
found: torch.IntStorage
found: torch.LongStorage
found: pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint
found: numpy.core.multiarray.scalar
found: numpy.dtype
found: _codecs.encode
Scan for models/sus.ckpt PASSED ✅
Reading models/sussy.ckpt
Found pickle in zip: archive/data.pkl
Scanning: archive/data.pkl
found: __builtin__.eval
BLOCKED: __builtin__.eval
Scan for models/sussy.ckpt FAILED! ⚠️
TheEliteGeek commented 1 year ago

Building on this Checks if the venv path and model directory was set at top of script Allows to specify one additional folder to check (optional) Options to display the scan results in console or open in notepad Checks the output file for the words SCAN FAILED for quick check if anything was suspicious

@echo off

echo scan started on %date% %time% > scan_output.txt

:SETUP
rem scans SD models for malicious pickling
set VENV_PATH="F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe"
set SD_FOLDER="F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion"
set DOWNLOAD_FOLDER="F:\Whatever\Other\Folder"

rem how result details should be displayed ("yes" or "no"):
set SHOW_RESULT_IN_CONSOLE="yes"
set OPEN_RESULT_IN_NOTEPAD="yes"

:SCANNING
rem check if VENV_PATH was set
if %VENV_PATH% equ "F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe" (
echo ##### ERROR please set your VENV_PATH #####
goto EXIT
)

echo "Scanning...Please wait a moment..."

rem check if SD_FOLDER was set
if %SD_FOLDER% equ "F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion" (
echo ##### ERROR please set your SD_FOLDER #####
goto EXIT
)

echo "step 1: SD models folder"
echo ####################################################################### >> scan_output.txt
echo ##### scanning SD model folder "~~~webui\models\Stable-diffusion" ##### >> scan_output.txt
echo ####################################################################### >> scan_output.txt
%VENV_PATH% pickle_scan.py %SD_FOLDER% >> scan_output.txt

rem check if download folder was set, if not just skip instead of errorlevel
if %DOWNLOAD_FOLDER% equ "F:\Whatever\Other\Folder" (
echo "No download folder specified"
goto DISPLAY_RESULT
)
echo "step 2: download folder"
echo ##################################### >> scan_output.txt
echo ##### scanning download folder  ##### >> scan_output.txt
echo ##################################### >> scan_output.txt
%VENV_PATH% pickle_scan.py %DOWNLOAD_FOLDER% >> scan_output.txt

:DISPLAY_RESULT
if %SHOW_RESULT_IN_CONSOLE% equ "yes" (type scan_output.txt)
if %OPEN_RESULT_IN_NOTEPAD% equ "yes" (start notepad scan_output.txt)

echo "Number of failed scans (potentially malicious files):"
find /c "SCAN FAILED" scan_output.txt

:EXIT
pause

So I'm trying to use this method, but it gives me the error saying that the venv path isn't set, even though I seem to have it set properly. Here's what I have in the file:

@echo off

echo scan started on %date% %time% > scan_output.txt

:SETUP rem scans SD models for malicious pickling set VENV_PATH="F:\Stable_Diffusion_2.0\stable-diffusion-webui\venv\Scripts\python.exe" set SD_FOLDER="F:\Stable_Diffusion_2.0\stable-diffusion-webui\models\Stable-diffusion" set DOWNLOAD_FOLDER="F:\Whatever\Other\Folder"

rem how result details should be displayed ("yes" or "no"): set SHOW_RESULT_IN_CONSOLE="yes" set OPEN_RESULT_IN_NOTEPAD="yes"

:SCANNING rem check if VENV_PATH was set if %VENV_PATH% equ "F:\Stable_Diffusion_2.0\stable-diffusion-webui\venv\Scripts\python.exe" ( echo ##### ERROR please set your VENV_PATH ##### goto EXIT )

echo "Scanning...Please wait a moment..."

rem check if SD_FOLDER was set if %SD_FOLDER% equ "F:\Stable_Diffusion_2.0\stable-diffusion-webui\models\Stable-diffusion" ( echo ##### ERROR please set your SD_FOLDER ##### goto EXIT )

echo "step 1: SD models folder" echo ####################################################################### >> scan_output.txt echo ##### scanning SD model folder "~~~webui\models\Stable-diffusion" ##### >> scan_output.txt echo ####################################################################### >> scan_output.txt %VENV_PATH% pickle_scan.py %SD_FOLDER% >> scan_output.txt

rem check if download folder was set, if not just skip instead of errorlevel if %DOWNLOAD_FOLDER% equ "F:\Whatever\Other\Folder" ( echo "No download folder specified" goto DISPLAY_RESULT ) echo "step 2: download folder" echo ##################################### >> scan_output.txt echo ##### scanning download folder ##### >> scan_output.txt echo ##################################### >> scan_output.txt %VENV_PATH% pickle_scan.py %DOWNLOAD_FOLDER% >> scan_output.txt

:DISPLAY_RESULT if %SHOW_RESULT_IN_CONSOLE% equ "yes" (type scan_output.txt) if %OPEN_RESULT_IN_NOTEPAD% equ "yes" (start notepad scan_output.txt)

echo "Number of failed scans (potentially malicious files):" find /c "SCAN FAILED" scan_output.txt

:EXIT pause

But when I run it, it gives me the "ERROR please set your VENV_PATH" message. Am I doing something wrong?

You only needed to update the paths under :SETUP It looks like you went a bit overboard and also set the path under :SCANNING. Go back and set the values under :SCANNING respectively. "F:\Whatever\Path\stable-diffusion-webui\venv\Scripts\Python.exe" "F:\Whatever\Path\stable-diffusion-webui\models\Stable-diffusion"

The reason is that "F:\Whatever\Path\stable-diffusion-webui..." was a placeholder path. The code under :SCANNING checks if it's the same value as the default path. If it's not, then it knows you changed the path to (hopefully) whatever is correct for you.

It's not the best written script and I guess a note should have been added to prevent this confusion :)

Masqued-Man commented 1 year ago

Yep, that fixed it. Thanks very much for the response.