Merge pull request #492 from JeLuF/scanner

Picklescan of model files

CHANGES.md

@@ -20,6 +20,7 @@
 - A `What's New?` tab in the UI
 
 ### Detailed changelog
+* 2.4.9 - 18 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. Thanks @JeLuf
 * 2.4.8 - 18 Nov 2022 - A `Use as Input` button to use the settings from a previously generated image task. Thanks @patriceac
 * 2.4.7 - 18 Nov 2022 - Don't crash if a VAE file fails to load
 * 2.4.7 - 17 Nov 2022 - Fix a bug where Face Correction (GFPGAN) would fail on cuda:N (i.e. GPUs other than cuda:0), as well as fail on CPU if the system had an incompatible GPU.

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

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"

ui/media/js/main.js

@@ -1150,8 +1150,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 = '<H1>🔥Malware alert!🔥</H1><h2>The file <i>' + models['scan-error'] + '</i> in your <tt>models/stable-diffusion</tt> folder is probably malware infected.</h2><h2>Please delete this file from the folder before proceeding!</h2>After deleting the file, reload this page.<br><br><button onClick="window.location.reload();">Reload Page</button>'
+            makeImageBtn.disabled = true
+	}
         let modelOptions = models['options']
         let stableDiffusionOptions = modelOptions['stable-diffusion']
         let vaeOptions = modelOptions['vae']

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)]