diff --git a/app/simple_downloader.py b/app/simple_downloader.py index 2b18d15c3..00ba0bd76 100644 --- a/app/simple_downloader.py +++ b/app/simple_downloader.py @@ -23,11 +23,15 @@ class SimpleDownloader: task_id = str(uuid.uuid4()) # SECURITY: Validate and sanitize inputs to prevent path traversal - # Sanitize model_type to prevent directory traversal - model_type = os.path.basename(model_type).replace('..', '').replace('/', '').replace('\\', '') + # Sanitize model_type - remove dangerous characters but keep underscores + import re + model_type = re.sub(r'[./\\]', '', model_type) + model_type = model_type.replace('..', '') - # Sanitize filename to prevent path traversal - filename = os.path.basename(filename).replace('..', '') + # Sanitize filename - use os.path.basename and remove traversal attempts + filename = os.path.basename(filename) + if '..' in filename or filename.startswith('.'): + raise ValueError("Invalid filename - no hidden files or path traversal") # Validate filename has allowed extension allowed_extensions = ['.safetensors', '.ckpt', '.pt', '.pth', '.bin', '.sft'] diff --git a/server.py b/server.py index 01f4a6926..0c4015e03 100644 --- a/server.py +++ b/server.py @@ -822,10 +822,11 @@ class PromptServer(): if not filename: filename = "model.safetensors" - # SECURITY: Sanitize filename + # SECURITY: Sanitize filename - allow more characters but still safe import os filename = os.path.basename(filename) - if not re.match(r'^[a-zA-Z0-9_.-]+$', filename): + # Block path traversal attempts but allow spaces, parens, brackets, etc + if '..' in filename or '/' in filename or '\\' in filename or filename.startswith('.'): return web.json_response({"error": "Invalid filename format"}, status=400) # Create download task (simple_downloader now has additional validation)