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.
This commit is contained in:
JeLuF 2022-11-16 22:34:02 +01:00
parent 3555fa36aa
commit 3e18f2f09c
4 changed files with 35 additions and 1 deletions

View File

@ -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

View File

@ -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"

View File

@ -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 = '<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']

View File

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