diff --git a/__init__.py b/__init__.py index 393e96b..9117115 100644 --- a/__init__.py +++ b/__init__.py @@ -190,6 +190,7 @@ class ModelResolverExtension: from .core.scanner import get_model_files, invalidate_model_files_cache, find_local_file_path from .core.path_utils import ( get_filename_from_path, + is_path_in_configured_model_roots, read_json_safe, write_json_atomic, ) @@ -235,6 +236,7 @@ class ModelResolverExtension: get_metadata_sidecar_path, normalize_download_category, write_lora_manager_metadata, + sanitize_download_filename, ) from .core.aria2_installer import Aria2InstallError, install_aria2_engine from .core.sources.popular import ( @@ -733,16 +735,20 @@ class ModelResolverExtension: {"error": "path is required"}, status=400 ) - normalized_path = os.path.normpath(target_path) + normalized_path = os.path.abspath(os.path.normpath(target_path)) if not os.path.exists(normalized_path): return web.json_response( {"error": "path does not exist"}, status=404 ) + if not is_path_in_configured_model_roots(normalized_path): + return web.json_response( + {"error": "path is outside configured model directories"}, + status=403, + ) if os.path.isfile(normalized_path): - absolute_path = os.path.abspath(normalized_path) subprocess.Popen( - ["explorer.exe", "/select,", absolute_path], + ["explorer.exe", "/select,", normalized_path], shell=False, ) else: @@ -769,6 +775,8 @@ class ModelResolverExtension: normalized_path = _os.path.abspath(_os.path.normpath(file_path)) if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path): return "", "file does not exist" + if not is_path_in_configured_model_roots(normalized_path): + return "", "file is outside configured model directories" return normalized_path, "" def write_calculated_hash_metadata(normalized_path, metadata_path, sha256): @@ -901,8 +909,9 @@ class ModelResolverExtension: {"error": "file_path is required"}, status=400 ) if error: + status = 403 if error == "file is outside configured model directories" else 404 return web.json_response( - {"error": error}, status=404 + {"error": error}, status=status ) sha256 = calculate_sha256_with_progress(normalized_path) @@ -941,8 +950,9 @@ class ModelResolverExtension: {"error": "file_path is required"}, status=400 ) if error: + status = 403 if error == "file is outside configured model directories" else 404 return web.json_response( - {"error": error}, status=404 + {"error": error}, status=status ) self.hash_tracker.cleanup() @@ -4055,6 +4065,7 @@ class ModelResolverExtension: parsed = urlparse(url) filename = unquote(parsed.path.split("/")[-1]) + filename = sanitize_download_filename(filename) if not filename: return web.json_response( diff --git a/core/downloader.py b/core/downloader.py index bff1788..aa55569 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -7,6 +7,7 @@ Handles downloading models from various sources with progress tracking. import os import json import hashlib +import re import secrets import shutil import socket @@ -55,7 +56,7 @@ ARIA2_RPC_TIMEOUT = (2, 5) # local JSON-RPC should respond quickly ARIA2_IDLE_STOP_SECONDS = 5 * 60 DOWNLOAD_USER_AGENT = DEFAULT_BROWSER_USER_AGENT -from .settings import CATEGORY_MAP, load_settings, normalize_download_backend, normalize_download_category +from .settings import CATEGORY_MAP, load_settings, normalize_download_backend, normalize_download_category, normalize_relative_subfolder SENSITIVE_METADATA_KEYS = { "authorization", @@ -86,6 +87,18 @@ SENSITIVE_QUERY_KEYS = { "cookie", } +_INVALID_DOWNLOAD_FILENAME_RE = re.compile(r'[<>:"/\\|?*\x00-\x1f]+') + + +def sanitize_download_filename(filename: Any) -> str: + """Return a safe basename for a downloaded model file.""" + text = get_filename_from_path(str(filename or "")).strip() + text = _INVALID_DOWNLOAD_FILENAME_RE.sub("_", text) + text = re.sub(r"\s+", " ", text).strip(" .") + if text in {"", ".", ".."}: + return "" + return text + def _is_sensitive_metadata_key(key: Any) -> bool: key_text = str(key or "").strip().lower() @@ -1670,6 +1683,15 @@ def download_model( if download_id is None: download_id = generate_download_id() + filename = sanitize_download_filename(filename) + if not filename: + return { + "success": False, + "download_id": download_id, + "error": "Invalid filename", + } + subfolder = normalize_relative_subfolder(subfolder) + # Get destination directory dest_dir = get_download_directory(category, base_directory) if not dest_dir: @@ -1681,9 +1703,16 @@ def download_model( # Add subfolder if specified if subfolder: - dest_dir = os.path.join(dest_dir, subfolder) + dest_dir = os.path.join(dest_dir, *subfolder.split("/")) - dest_path = os.path.join(dest_dir, filename) + dest_dir = os.path.abspath(os.path.normpath(dest_dir)) + dest_path = os.path.abspath(os.path.normpath(os.path.join(dest_dir, filename))) + if not is_path_within(dest_path, dest_dir): + return { + "success": False, + "download_id": download_id, + "error": "Download target is outside the selected model directory", + } # Check if file already exists if os.path.exists(dest_path): @@ -1976,10 +2005,12 @@ def start_background_download( download_id for tracking progress """ download_id = generate_download_id() + filename = sanitize_download_filename(filename) + subfolder = normalize_relative_subfolder(subfolder) initial_directory = get_download_directory(category, base_directory) or "" if initial_directory and subfolder: - initial_directory = os.path.join(initial_directory, subfolder) - initial_path = os.path.join(initial_directory, filename) if initial_directory else "" + initial_directory = os.path.join(initial_directory, *subfolder.split("/")) + initial_path = os.path.join(initial_directory, filename) if initial_directory and filename else "" # Pre-initialize progress dict so it's always available for polling # even if download fails before download_file is called diff --git a/core/path_utils.py b/core/path_utils.py index 71f290e..1bd0a80 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -8,7 +8,7 @@ and directory matching to be shared across all modules. import os import json import hashlib -from typing import Any, Tuple, Optional, Callable +from typing import Any, Tuple, Optional, Callable, Iterable, List _log = None @@ -58,13 +58,100 @@ def is_path_within(path_value: Any, root_value: Any) -> bool: if not path_value or not root_value: return False try: - path_key = get_path_key(path_value) - root_key = get_path_key(root_value) + path_key = get_path_identity(str(path_value)) + root_key = get_path_identity(str(root_value)) return os.path.commonpath([path_key, root_key]) == root_key except Exception: return False +def _get_folder_paths_module(folder_paths_module: Optional[Any] = None) -> Optional[Any]: + if folder_paths_module is not None: + return folder_paths_module + try: + import folder_paths + return folder_paths + except ImportError: + return None + + +def _coerce_folder_paths(value: Any) -> List[str]: + if isinstance(value, str): + return [value] + if not isinstance(value, (list, tuple)): + return [] + if value and isinstance(value[0], (list, tuple, set)): + return [str(path) for path in value[0] if path] + return [str(path) for path in value if isinstance(path, str) and path] + + +def get_configured_model_roots( + folder_paths_module: Optional[Any] = None, + *, + skip_keys: Optional[Iterable[str]] = None, +) -> List[str]: + """Return all configured ComfyUI model roots, including external model folders.""" + fp = _get_folder_paths_module(folder_paths_module) + if fp is None: + return [] + + skipped = {str(key).strip().lower() for key in (skip_keys or {"custom_nodes", "configs"})} + names = [] + folder_names_and_paths = getattr(fp, "folder_names_and_paths", None) + if isinstance(folder_names_and_paths, dict): + names = list(folder_names_and_paths.keys()) + else: + try: + get_folder_names = getattr(fp, "get_folder_names", None) + if callable(get_folder_names): + names = list(get_folder_names() or []) + except Exception: + names = [] + + roots: List[str] = [] + seen = set() + for raw_name in names: + name = str(raw_name or "").strip() + if not name or name.lower() in skipped: + continue + + paths = [] + try: + get_folder_paths = getattr(fp, "get_folder_paths", None) + if callable(get_folder_paths): + paths = list(get_folder_paths(name) or []) + except Exception: + paths = [] + + if not paths and isinstance(folder_names_and_paths, dict): + paths = _coerce_folder_paths(folder_names_and_paths.get(name)) + + for path in paths: + if not path: + continue + absolute_path = os.path.abspath(os.path.normpath(str(path))) + identity = get_path_identity(absolute_path) + if identity in seen: + continue + seen.add(identity) + roots.append(absolute_path) + + return roots + + +def is_path_in_configured_model_roots( + path_value: Any, + folder_paths_module: Optional[Any] = None, +) -> bool: + """Return True when a path is inside a configured ComfyUI model root.""" + if not path_value: + return False + return any( + is_path_within(path_value, root) + for root in get_configured_model_roots(folder_paths_module) + ) + + def prefer_local_base_directory( candidate: str, current: str, diff --git a/metadata/base-models.meta.json b/metadata/base-models.meta.json index 3486d72..a208557 100644 --- a/metadata/base-models.meta.json +++ b/metadata/base-models.meta.json @@ -1,11 +1,7 @@ { - "updated_at": "2026-06-21T00:24:23+00:00", - "last_checked_at": "2026-06-21T00:24:23+00:00", + "updated_at": "2026-07-04T21:52:15+00:00", + "last_checked_at": "2026-07-04T21:52:15+00:00", "local_count": 94, - "new_models_added": 3, - "new_models_added_list": [ - "Ideogram 4.0", - "Boogu", - "PolyGen" - ] + "new_models_added": 0, + "new_models_added_list": [] } diff --git a/metadata/model-list.meta.json b/metadata/model-list.meta.json index 3dbaf00..7acd99a 100644 --- a/metadata/model-list.meta.json +++ b/metadata/model-list.meta.json @@ -5,5 +5,5 @@ "sha": "5d9308f50a08cf3cfb47e6a25cc525425714d3ad", "size": 274526, "model_count": 538, - "updated_at": "2026-07-03T22:59:41+00:00" + "updated_at": "2026-07-04T21:52:13+00:00" }