From 3e18f2f09ca7f079c617ed42a882cb0f84781e21 Mon Sep 17 00:00:00 2001 From: JeLuF Date: Wed, 16 Nov 2022 22:34:02 +0100 Subject: [PATCH 1/2] Picklescan of model files During getModel(), the server picklescans the model files for potential malicious code in the pickled python objects. If a malicious file is found, the web UI will show a big red error message, the makeImage button will be disabled, and the user must remove the malicious file and reload the UI page. --- scripts/on_sd_start.bat | 10 ++++++++++ scripts/on_sd_start.sh | 7 +++++++ ui/media/js/main.js | 10 +++++++++- ui/server.py | 9 +++++++++ 4 files changed, 35 insertions(+), 1 deletion(-) diff --git a/scripts/on_sd_start.bat b/scripts/on_sd_start.bat index 51a51549..df28e57a 100644 --- a/scripts/on_sd_start.bat +++ b/scripts/on_sd_start.bat @@ -191,6 +191,16 @@ call WHERE uvicorn > .tmp exit /b ) +@>nul 2>nul call python -m picklescan --help +@if "%ERRORLEVEL%" NEQ "0" ( + @echo. & echo Picklescan not found. Installing + @call pip install picklescan || ( + echo "Error installing the picklescan package necessary for Stable Diffusion UI. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" + pause + exit /b + ) +) + @>nul findstr /m "conda_sd_ui_deps_installed" ..\scripts\install_status.txt @if "%ERRORLEVEL%" NEQ "0" ( @echo conda_sd_ui_deps_installed >> ..\scripts\install_status.txt diff --git a/scripts/on_sd_start.sh b/scripts/on_sd_start.sh index d4bb1ad1..adcba1f2 100755 --- a/scripts/on_sd_start.sh +++ b/scripts/on_sd_start.sh @@ -156,6 +156,13 @@ else echo conda_sd_ui_deps_installed >> ../scripts/install_status.txt fi +if python -m picklescan --help >/dev/null 2>&1; then + echo "Picklescan is already installed." +else + echo "Picklescan not found, installing." + pip install picklescan || fail "Picklescan installation failed." +fi + mkdir -p "../models/stable-diffusion" diff --git a/ui/media/js/main.js b/ui/media/js/main.js index afb0925d..eee7d8b9 100644 --- a/ui/media/js/main.js +++ b/ui/media/js/main.js @@ -1141,8 +1141,16 @@ async function getModels() { let res = await fetch('/get/models') const models = await res.json() - console.log('get models response', models) + console.log('got models response', models) + if ( "scan-error" in models ) { + // let previewPane = document.getElementById('tab-content-wrapper') + let previewPane = document.getElementById('preview') + previewPane.style.background="red" + previewPane.style.textAlign="center" + previewPane.innerHTML = '

🔥Malware alert!🔥

The file ' + models['scan-error'] + ' in your models/stable-diffusion folder is probably malware infected.

Please delete this file from the folder before proceeding!

After deleting the file, reload this page.

' + makeImageBtn.disabled = true + } let modelOptions = models['options'] let stableDiffusionOptions = modelOptions['stable-diffusion'] let vaeOptions = modelOptions['vae'] diff --git a/ui/server.py b/ui/server.py index 8b54f847..154fc6cd 100644 --- a/ui/server.py +++ b/ui/server.py @@ -7,6 +7,8 @@ import traceback import sys import os +import picklescan.scanner +import rich SD_DIR = os.getcwd() print('started in ', SD_DIR) @@ -206,6 +208,13 @@ def getModels(): os.makedirs(models_dir) for file in os.listdir(models_dir): + scan_result = picklescan.scanner.scan_file_path( os.path.join(models_dir, file)) + if ( scan_result.issues_count >0 or scan_result.infected_files >0): + rich.print(":warning: [bold red]Scan %s: %d scanned, %d issue, %d infected.[/bold red]" % ( file, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files) ) + models['scan-error'] = file + return models + else: + rich.print("Scan %s: [green]%d scanned, %d issue, %d infected.[/green]" % ( file, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files ) ) for model_extension in model_extensions: if file.endswith(model_extension): model_name = file[:-len(model_extension)] From 7b4cfbeeaa80d69655885fdeb2a81fa4c259ceca Mon Sep 17 00:00:00 2001 From: JeLuF Date: Wed, 16 Nov 2022 23:47:24 +0100 Subject: [PATCH 2/2] Update CHANGES.md --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index c9553b22..05562b34 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -2,6 +2,7 @@ ### 2.4.6 * 16 Nov 2022 - Fix a regression in VRAM usage during startup, which caused 'Out of Memory' errors when starting on GPUs with 4gb (or less) VRAM +* 16 Nov 2022 - Add Picklescan - a scanner for malicious model files. If it finds a malicious file, it will halt the web application and alert the user ### 2.4.5 * 16 Nov 2022 - Add checkbox for "Open browser on startup".