diff --git a/py/api/__init__.py b/py/api/__init__.py index 90d4695c..c0c424dc 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -557,27 +557,49 @@ class PromptManagerAPI( return None def _get_all_output_dirs(self): - """Get all configured output directories, falling back to auto-detect. + """Get the directories images may be served from. + + Configured gallery roots are filtered through + ``GalleryConfig.validate_gallery_root`` so a root that was hand-edited + into config.json (or persisted before validation existed) can never + expose files outside ComfyUI's directories. Returns: - List[Path]: Valid output directory paths, possibly empty. + List[Path]: Validated directory paths, possibly empty. """ from ..config import GalleryConfig + configured = list(GalleryConfig.MONITORING_DIRECTORIES) output_dirs = [] - if GalleryConfig.MONITORING_DIRECTORIES: - for d in GalleryConfig.MONITORING_DIRECTORIES: - p = Path(d).resolve() - if p.is_dir(): - output_dirs.append(p) + for d in configured: + ok, reason = GalleryConfig.validate_gallery_root(d) + if ok: + output_dirs.append(Path(os.path.realpath(d))) + else: + self.logger.warning(f"Ignoring configured gallery root: {reason}") - # Fallback to auto-detect if no configured dirs are valid - if not output_dirs: - fallback = self._find_comfyui_output_dir() - if fallback: - output_dirs.append(Path(fallback)) + if output_dirs: + return output_dirs - return output_dirs + if configured: + # Every configured root was rejected: never let auto-detection + # resurrect it, serve ComfyUI's own output directory instead. + return self._comfyui_output_dir_list() + + fallback = self._find_comfyui_output_dir() + return [Path(fallback)] if fallback else [] + + def _comfyui_output_dir_list(self): + """ComfyUI's output directory as a one-element list, or [] if unknown.""" + try: + import folder_paths + + output_dir = folder_paths.get_output_directory() + except (ImportError, AttributeError, OSError): + return [] + if isinstance(output_dir, str) and os.path.isdir(output_dir): + return [Path(os.path.realpath(output_dir))] + return [] def _extract_comfyui_metadata(self, image_path): """Extract ComfyUI workflow metadata from PNG image files.""" diff --git a/py/api/admin.py b/py/api/admin.py index 7c2fb8e4..12e68e6b 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -402,12 +402,10 @@ class AdminRoutesMixin: async def save_settings(self, request): """Save settings.""" try: - from ..config import PromptManagerConfig, GalleryConfig + from ..config import PromptManagerConfig data = await request.json() - restart_required = False - # Update in-memory config if "result_timeout" in data: try: PromptManagerConfig.RESULT_TIMEOUT = validate_result_timeout( @@ -420,124 +418,15 @@ class AdminRoutesMixin: if "webui_display_mode" in data: PromptManagerConfig.WEBUI_DISPLAY_MODE = data["webui_display_mode"] - # Blocked system directories (shared by both path handlers) - blocked = [ - "/etc", - "/usr", - "/bin", - "/sbin", - "/boot", - "/proc", - "/sys", - "/dev", - "/var/log", - "/root", - "C:\\Windows", - "C:\\Program Files", - ] + new_roots, error = self._parse_gallery_roots(data) + if error: + return web.json_response({"success": False, "error": error}, status=400) - # Handle gallery root paths (array — preferred) - if "gallery_root_paths" in data: - new_paths = data["gallery_root_paths"] - if not isinstance(new_paths, list): - return web.json_response( - { - "success": False, - "error": "gallery_root_paths must be a list", - }, - status=400, - ) + restart_required = False + if new_roots is not None: + restart_required = self._apply_gallery_roots(new_roots) - validated_paths = [] - for path_str in new_paths: - path_str = path_str.strip() - if not path_str: - continue - resolved = Path(path_str).resolve() - if not resolved.is_dir(): - return web.json_response( - { - "success": False, - "error": f"Path does not exist or is not a directory: {path_str}", - }, - status=400, - ) - for b in blocked: - if str(resolved).startswith(b): - return web.json_response( - { - "success": False, - "error": f"Cannot use system directory: {path_str}", - }, - status=400, - ) - validated_paths.append(path_str) - - old_paths = list(GalleryConfig.MONITORING_DIRECTORIES) - if validated_paths != old_paths: - GalleryConfig.MONITORING_DIRECTORIES = validated_paths - self._cached_output_dir = None - self._gallery_cache = {} - restart_required = True - - elif "gallery_root_path" in data: - # Backward compat: single path string - new_path = data["gallery_root_path"].strip() - old_path = ( - GalleryConfig.MONITORING_DIRECTORIES[0] - if GalleryConfig.MONITORING_DIRECTORIES - else "" - ) - - if new_path != old_path: - if new_path: - resolved = Path(new_path).resolve() - if not resolved.is_dir(): - return web.json_response( - { - "success": False, - "error": f"Gallery path does not exist or is not a directory: {new_path}", - }, - status=400, - ) - for b in blocked: - if str(resolved).startswith(b): - return web.json_response( - { - "success": False, - "error": "Gallery path cannot point to a system directory", - }, - status=400, - ) - GalleryConfig.MONITORING_DIRECTORIES = [new_path] - else: - GalleryConfig.MONITORING_DIRECTORIES = [] - self._cached_output_dir = None - self._gallery_cache = {} - restart_required = True - - # Save to config file for persistence - config_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) - config_file = os.path.join(config_dir, "config.json") - - config_data = { - "web_ui": { - "result_timeout": PromptManagerConfig.RESULT_TIMEOUT, - "webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE, - }, - "gallery": { - "monitoring": {"directories": GalleryConfig.MONITORING_DIRECTORIES} - }, - } - - try: - with open(config_file, "w") as f: - json.dump(config_data, f, indent=2) - self.logger.info(f"Settings saved to {config_file}") - except Exception as save_err: - self.logger.warning(f"Could not save config file: {save_err}") + self._persist_settings() return web.json_response( { @@ -547,11 +436,82 @@ class AdminRoutesMixin: } ) except Exception as e: + self.logger.error(f"Save settings error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to save settings: {str(e)}"}, + {"success": False, "error": "Failed to save settings"}, status=500, ) + def _parse_gallery_roots(self, data): + """Extract and validate gallery roots from a settings payload. + + Accepts ``gallery_root_paths`` (list, preferred) or the legacy single + ``gallery_root_path`` string. Blank entries are dropped. + + Returns: + (roots, error): ``roots`` is the validated list, or ``None`` when + the payload carries no gallery root key; ``error`` is a message + when validation failed (and ``roots`` is then ``None``). + """ + if "gallery_root_paths" in data: + raw = data["gallery_root_paths"] + if not isinstance(raw, list): + return None, "gallery_root_paths must be a list" + elif "gallery_root_path" in data: + raw = [data["gallery_root_path"]] + else: + return None, None + + from ..config import GalleryConfig + + roots = [] + for entry in raw: + if not isinstance(entry, str): + return None, "Gallery root paths must be strings" + entry = entry.strip() + if not entry: + continue + ok, reason = GalleryConfig.validate_gallery_root(entry) + if not ok: + return ( + None, + f"Invalid gallery root '{os.path.basename(entry)}': {reason}", + ) + roots.append(entry) + return roots, None + + def _apply_gallery_roots(self, roots): + """Install validated gallery roots; returns True when they changed.""" + from ..config import GalleryConfig + + if roots == list(GalleryConfig.MONITORING_DIRECTORIES): + return False + GalleryConfig.MONITORING_DIRECTORIES = roots + self._cached_output_dir = None + self._gallery_cache = {} + return True + + def _persist_settings(self): + """Write the user-editable settings to the configured config.json.""" + from ..config import PromptManagerConfig, GalleryConfig + + config_file = PromptManagerConfig.get_config_path() + config_data = { + "web_ui": { + "result_timeout": PromptManagerConfig.RESULT_TIMEOUT, + "webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE, + }, + "gallery": { + "monitoring": {"directories": GalleryConfig.MONITORING_DIRECTORIES} + }, + } + try: + with open(config_file, "w") as f: + json.dump(config_data, f, indent=2) + self.logger.info(f"Settings saved to {config_file}") + except OSError as save_err: + self.logger.warning(f"Could not save config file: {save_err}") + async def run_diagnostics(self, request): """Run comprehensive system diagnostics and health checks.""" try: diff --git a/py/config.py b/py/config.py index 88ef6b93..00e6e3f2 100644 --- a/py/config.py +++ b/py/config.py @@ -31,7 +31,8 @@ routes = server_instance.routes extension_uri = None # Will be set in __init__.py import os -from typing import Dict, Any, List +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple # Import logging system try: @@ -46,6 +47,58 @@ except ImportError: # Initialize logger for config operations config_logger = get_logger("prompt_manager.config") +# Environment variable listing extra directories (os.pathsep-separated) that +# may be used as gallery roots in addition to ComfyUI's own directories. +EXTRA_GALLERY_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS" + +# Environment variable overriding where config.json is read from and saved to. +CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH" + +# ComfyUI folder_paths getters whose directories are valid gallery parents. +_COMFYUI_DIRECTORY_GETTERS = ( + "get_output_directory", + "get_input_directory", + "get_temp_directory", + "get_user_directory", +) + + +def _canonical_path(path: str) -> str: + """Return a path normalised for comparison (realpath + normcase).""" + return os.path.normcase(os.path.realpath(path)) + + +def _is_filesystem_root(path: str) -> bool: + """True for '/' on POSIX and drive roots such as 'C:\\' on Windows.""" + return os.path.dirname(path) == path + + +def _comfyui_directories() -> List[str]: + """Directories reported by ComfyUI's folder_paths module, if importable.""" + try: + import folder_paths + except ImportError: + return [] + + found = [] + for getter_name in _COMFYUI_DIRECTORY_GETTERS: + getter = getattr(folder_paths, getter_name, None) + if getter is None: + continue + try: + directory = getter() + except Exception: + continue + if isinstance(directory, str) and directory: + found.append(directory) + return found + + +def _extra_gallery_roots() -> List[str]: + """Directories the user opted in through PROMPT_MANAGER_EXTRA_GALLERY_ROOTS.""" + raw = os.environ.get(EXTRA_GALLERY_ROOTS_ENV, "") + return [entry.strip() for entry in raw.split(os.pathsep) if entry.strip()] + class GalleryConfig: """Configuration class for the gallery monitoring and image processing system. @@ -200,6 +253,61 @@ class GalleryConfig: if "metadata_extraction_timeout" in performance: cls.METADATA_EXTRACTION_TIMEOUT = performance["metadata_extraction_timeout"] + @classmethod + def allowed_gallery_parents(cls) -> List[str]: + """Canonical directories under which gallery roots may live. + + Filesystem roots are never allowed as parents, even when listed in + the environment, because that would re-open the whole disk. + """ + parents = [] + for candidate in _comfyui_directories() + _extra_gallery_roots(): + try: + canonical = _canonical_path(candidate) + except (OSError, ValueError): + continue + if _is_filesystem_root(canonical) or not os.path.isdir(canonical): + continue + if canonical not in parents: + parents.append(canonical) + return parents + + @classmethod + def validate_gallery_root(cls, path: Any) -> Tuple[bool, str]: + """Check whether ``path`` may be used as a gallery root. + + Returns: + (True, "") when the directory lies inside one of + :meth:`allowed_gallery_parents`; otherwise (False, reason). + """ + if not isinstance(path, str) or not path.strip(): + return False, "Gallery root must be a non-empty path" + + try: + canonical = _canonical_path(path.strip()) + except (OSError, ValueError): + return False, "Gallery root could not be resolved" + + if _is_filesystem_root(canonical): + return False, "Gallery root cannot be a filesystem root" + if canonical == _canonical_path(os.path.expanduser("~")): + return False, "Gallery root cannot be the home directory" + if not os.path.exists(canonical): + return False, "Gallery root does not exist" + if not os.path.isdir(canonical): + return False, "Gallery root is not a directory" + + candidate = Path(canonical) + for parent in cls.allowed_gallery_parents(): + if candidate.is_relative_to(Path(parent)): + return True, "" + + return ( + False, + "Gallery root must be inside a ComfyUI directory " + f"(output, input, temp, user) or one listed in {EXTRA_GALLERY_ROOTS_ENV}", + ) + class IntegrationConfig: """Configuration for third-party extension integrations. @@ -315,7 +423,20 @@ class PromptManagerConfig: } @classmethod - def load_from_file(cls, config_path: str): + def get_config_path(cls) -> str: + """Path of the persisted config.json. + + Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file + lives at the repository root next to ``pyproject.toml``. + """ + override = os.environ.get(CONFIG_PATH_ENV, "").strip() + if override: + return override + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + return os.path.join(repo_root, "config.json") + + @classmethod + def load_from_file(cls, config_path: Optional[str] = None): """Load configuration settings from a JSON file. Reads configuration from the specified JSON file and updates the current @@ -324,7 +445,8 @@ class PromptManagerConfig: Args: config_path (str): Path to the JSON configuration file to load. - Can be relative or absolute path. + Can be relative or absolute path. Defaults to + :meth:`get_config_path`. Raises: The method handles all exceptions internally and logs errors rather @@ -336,6 +458,9 @@ class PromptManagerConfig: """ import json + if config_path is None: + config_path = cls.get_config_path() + if os.path.exists(config_path): try: with open(config_path, "r") as f: @@ -348,7 +473,7 @@ class PromptManagerConfig: config_logger.info(f"Config file not found: {config_path}, using defaults") @classmethod - def save_to_file(cls, config_path: str): + def save_to_file(cls, config_path: Optional[str] = None): """Save the current configuration to a JSON file. Serializes the complete configuration (including gallery settings) to @@ -357,6 +482,7 @@ class PromptManagerConfig: Args: config_path (str): Path where the JSON configuration file should be saved. Parent directories will be created if they don't exist. + Defaults to :meth:`get_config_path`. Raises: The method handles all exceptions internally and logs errors rather @@ -368,9 +494,14 @@ class PromptManagerConfig: """ import json + if config_path is None: + config_path = cls.get_config_path() + try: config = cls.get_config() - os.makedirs(os.path.dirname(config_path), exist_ok=True) + parent_dir = os.path.dirname(config_path) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) with open(config_path, "w") as f: json.dump(config, f, indent=2) @@ -434,8 +565,6 @@ class PromptManagerConfig: # Load configuration on import try: - config_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - config_file = os.path.join(config_dir, "config.json") - PromptManagerConfig.load_from_file(config_file) + PromptManagerConfig.load_from_file() except Exception as e: config_logger.error(f"Error during config initialization: {e}") diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py new file mode 100644 index 00000000..bcf74ba1 --- /dev/null +++ b/tests/test_api_admin.py @@ -0,0 +1,225 @@ +"""Route tests for the admin API (py/api/admin.py). + +Uses aiohttp's test client against a PromptManagerAPI wired to a temporary +database and a stub ``folder_paths`` module, so the tests run without ComfyUI. +""" + +import json +import os +import shutil +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules["server"] = _mock_server + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from aiohttp import web # noqa: E402 +from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 + +from database.operations import PromptDatabase # noqa: E402 +from py.api import PromptManagerAPI # noqa: E402 +from py.config import GalleryConfig, PromptManagerConfig # noqa: E402 + +CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH" +EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS" + + +class AdminAPITestCase(AioHTTPTestCase): + """Test app with PromptManager routes, temp DB, temp config and fake ComfyUI tree.""" + + async def get_application(self): + self.tmpdir = tempfile.mkdtemp() + self.comfy_dir = Path(self.tmpdir) / "ComfyUI" + self.output_dir = self.comfy_dir / "output" + self.output_dir.mkdir(parents=True) + self.outside_dir = Path(self.tmpdir) / "outside" + self.outside_dir.mkdir() + + self.config_path = os.path.join(self.tmpdir, "config.json") + self._orig_env = { + key: os.environ.pop(key, None) for key in (CONFIG_PATH_ENV, EXTRA_ROOTS_ENV) + } + os.environ[CONFIG_PATH_ENV] = self.config_path + + self._orig_folder_paths = sys.modules.get("folder_paths") + sys.modules["folder_paths"] = types.SimpleNamespace( + base_path=str(self.comfy_dir), + get_output_directory=lambda: str(self.output_dir), + ) + + self._orig_gallery_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) + self._orig_pm_config = PromptManagerConfig.get_config() + GalleryConfig.MONITORING_DIRECTORIES = [] + + db_path = os.path.join(self.tmpdir, "prompts.db") + app = web.Application() + routes = web.RouteTableDef() + self.api = PromptManagerAPI() + self.api.db = PromptDatabase(db_path) + self.api.add_routes(routes) + app.router.add_routes(routes) + return app + + async def tearDownAsync(self): + GalleryConfig.MONITORING_DIRECTORIES = self._orig_gallery_dirs + PromptManagerConfig.update_config(self._orig_pm_config) + for key, value in self._orig_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + if self._orig_folder_paths is None: + sys.modules.pop("folder_paths", None) + else: + sys.modules["folder_paths"] = self._orig_folder_paths + shutil.rmtree(self.tmpdir, ignore_errors=True) + + # ── helpers ──────────────────────────────────────────────────────── + + def _write_config(self, content): + with open(self.config_path, "w") as f: + json.dump(content, f) + + def _read_config(self): + with open(self.config_path) as f: + return json.load(f) + + +class TestSaveSettingsGalleryRoots(AdminAPITestCase): + + async def test_rejects_filesystem_root_and_keeps_config(self): + self._write_config({"marker": "untouched"}) + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": [os.path.abspath(os.sep)]}, + ) + + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertEqual(self._read_config(), {"marker": "untouched"}) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, []) + + async def test_rejects_directory_outside_comfyui(self): + self._write_config({"marker": "untouched"}) + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": [str(self.outside_dir)]}, + ) + + self.assertEqual(resp.status, 400) + self.assertEqual(self._read_config(), {"marker": "untouched"}) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, []) + + async def test_one_bad_root_rejects_whole_request(self): + good = self.output_dir / "good" + good.mkdir() + self._write_config({"marker": "untouched"}) + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": [str(good), str(self.outside_dir)]}, + ) + + self.assertEqual(resp.status, 400) + self.assertEqual(self._read_config(), {"marker": "untouched"}) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, []) + + async def test_legacy_single_path_rejected_too(self): + self._write_config({"marker": "untouched"}) + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_path": os.path.expanduser("~")}, + ) + + self.assertEqual(resp.status, 400) + self.assertEqual(self._read_config(), {"marker": "untouched"}) + + async def test_accepts_subdir_of_output_and_persists(self): + good = self.output_dir / "renders" + good.mkdir() + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": [str(good)]}, + ) + + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertTrue(data["success"]) + self.assertTrue(data["restart_required"]) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(good)]) + saved = self._read_config() + self.assertEqual(saved["gallery"]["monitoring"]["directories"], [str(good)]) + + async def test_accepts_env_listed_directory(self): + os.environ[EXTRA_ROOTS_ENV] = str(self.outside_dir) + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": [str(self.outside_dir)]}, + ) + + self.assertEqual(resp.status, 200) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.outside_dir)]) + + async def test_legacy_single_path_accepted(self): + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_path": str(self.output_dir)}, + ) + + self.assertEqual(resp.status, 200) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.output_dir)]) + + async def test_empty_paths_clear_roots(self): + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": ["", " "]}, + ) + + self.assertEqual(resp.status, 200) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, []) + self.assertEqual( + self._read_config()["gallery"]["monitoring"]["directories"], [] + ) + + async def test_paths_must_be_a_list(self): + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": str(self.output_dir)}, + ) + self.assertEqual(resp.status, 400) + + async def test_path_entries_must_be_strings(self): + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={"gallery_root_paths": [42]}, + ) + self.assertEqual(resp.status, 400) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_config.py b/tests/test_config.py index 8f908625..09444b4e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -291,5 +291,59 @@ class TestIntegrationConfig(unittest.TestCase): self.assertEqual(config1, config2) +class TestConfigPathOverride(unittest.TestCase): + """PROMPT_MANAGER_CONFIG_PATH relocates the persisted config.json.""" + + ENV = "PROMPT_MANAGER_CONFIG_PATH" + + def setUp(self): + self._orig_env = os.environ.pop(self.ENV, None) + self._orig = PromptManagerConfig.get_config() + + def tearDown(self): + if self._orig_env is None: + os.environ.pop(self.ENV, None) + else: + os.environ[self.ENV] = self._orig_env + PromptManagerConfig.update_config(self._orig) + + def test_default_config_path_is_repo_config_json(self): + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + expected = os.path.join(repo_root, "config.json") + self.assertEqual( + os.path.normcase(PromptManagerConfig.get_config_path()), + os.path.normcase(expected), + ) + + def test_env_override_is_used(self): + with tempfile.TemporaryDirectory() as tmpdir: + custom = os.path.join(tmpdir, "custom.json") + os.environ[self.ENV] = custom + self.assertEqual(PromptManagerConfig.get_config_path(), custom) + + def test_save_without_path_uses_env_override(self): + with tempfile.TemporaryDirectory() as tmpdir: + custom = os.path.join(tmpdir, "nested", "custom.json") + os.environ[self.ENV] = custom + PromptManagerConfig.RESULT_TIMEOUT = 42 + + PromptManagerConfig.save_to_file() + + with open(custom) as f: + saved = json.load(f) + self.assertEqual(saved["web_ui"]["result_timeout"], 42) + + def test_load_without_path_uses_env_override(self): + with tempfile.TemporaryDirectory() as tmpdir: + custom = os.path.join(tmpdir, "custom.json") + with open(custom, "w") as f: + json.dump({"web_ui": {"result_timeout": 77}}, f) + os.environ[self.ENV] = custom + + PromptManagerConfig.load_from_file() + + self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 77) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_gallery_roots.py b/tests/test_gallery_roots.py new file mode 100644 index 00000000..61199f32 --- /dev/null +++ b/tests/test_gallery_roots.py @@ -0,0 +1,261 @@ +"""Tests for gallery root validation (GalleryConfig.validate_gallery_root). + +A gallery root may only point inside one of ComfyUI's own directories +(output/input/temp/user, as reported by ``folder_paths``) or inside a +directory listed in ``PROMPT_MANAGER_EXTRA_GALLERY_ROOTS``. Anything else, +including filesystem roots and the home directory, is rejected so the +image-serving routes can never be pointed at arbitrary files. +""" + +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules["server"] = _mock_server + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from py.config import GalleryConfig # noqa: E402 + +EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS" + + +class GalleryRootTestCase(unittest.TestCase): + """Shared fixture: a fake ComfyUI tree exposed through a stub folder_paths.""" + + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(self._rmtree, self.tmpdir) + + self.comfy_dir = Path(self.tmpdir) / "ComfyUI" + self.output_dir = self.comfy_dir / "output" + self.input_dir = self.comfy_dir / "input" + self.output_dir.mkdir(parents=True) + self.input_dir.mkdir(parents=True) + + self._orig_folder_paths = sys.modules.get("folder_paths") + sys.modules["folder_paths"] = types.SimpleNamespace( + base_path=str(self.comfy_dir), + get_output_directory=lambda: str(self.output_dir), + get_input_directory=lambda: str(self.input_dir), + get_temp_directory=self._raise_runtime_error, + # get_user_directory intentionally absent + ) + self.addCleanup(self._restore_folder_paths) + + self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None) + self.addCleanup(self._restore_extra_env) + + @staticmethod + def _raise_runtime_error(): + raise RuntimeError("temp directory unavailable") + + @staticmethod + def _rmtree(path): + import shutil + + shutil.rmtree(path, ignore_errors=True) + + def _restore_folder_paths(self): + if self._orig_folder_paths is None: + sys.modules.pop("folder_paths", None) + else: + sys.modules["folder_paths"] = self._orig_folder_paths + + def _restore_extra_env(self): + if self._orig_extra is None: + os.environ.pop(EXTRA_ROOTS_ENV, None) + else: + os.environ[EXTRA_ROOTS_ENV] = self._orig_extra + + +class TestValidateGalleryRootRejects(GalleryRootTestCase): + + def test_filesystem_root_rejected(self): + ok, reason = GalleryConfig.validate_gallery_root(os.path.abspath(os.sep)) + self.assertFalse(ok) + self.assertIn("root", reason.lower()) + + @unittest.skipUnless(os.name == "nt", "drive roots only exist on Windows") + def test_windows_drive_root_rejected(self): + ok, _ = GalleryConfig.validate_gallery_root("C:\\") + self.assertFalse(ok) + + def test_home_directory_rejected(self): + ok, reason = GalleryConfig.validate_gallery_root(os.path.expanduser("~")) + self.assertFalse(ok) + self.assertIn("home", reason.lower()) + + def test_directory_outside_comfyui_rejected(self): + outside = Path(self.tmpdir) / "elsewhere" + outside.mkdir() + ok, _ = GalleryConfig.validate_gallery_root(str(outside)) + self.assertFalse(ok) + + def test_nonexistent_path_rejected(self): + missing = self.output_dir / "does_not_exist" + ok, reason = GalleryConfig.validate_gallery_root(str(missing)) + self.assertFalse(ok) + self.assertIn("exist", reason.lower()) + + def test_file_rejected(self): + file_path = self.output_dir / "image.png" + file_path.write_bytes(b"PNG") + ok, reason = GalleryConfig.validate_gallery_root(str(file_path)) + self.assertFalse(ok) + self.assertIn("directory", reason.lower()) + + def test_empty_and_non_string_rejected(self): + self.assertFalse(GalleryConfig.validate_gallery_root("")[0]) + self.assertFalse(GalleryConfig.validate_gallery_root(" ")[0]) + self.assertFalse(GalleryConfig.validate_gallery_root(None)[0]) + self.assertFalse(GalleryConfig.validate_gallery_root(42)[0]) + + def test_sibling_with_shared_prefix_rejected(self): + """'output2' must not pass because it string-starts with 'output'.""" + sibling = self.comfy_dir / "output2" + sibling.mkdir() + ok, _ = GalleryConfig.validate_gallery_root(str(sibling)) + self.assertFalse(ok) + + def test_traversal_out_of_output_rejected(self): + outside = Path(self.tmpdir) / "elsewhere" + outside.mkdir() + sneaky = os.path.join(str(self.output_dir), "..", "..", "elsewhere") + ok, _ = GalleryConfig.validate_gallery_root(sneaky) + self.assertFalse(ok) + + def test_symlink_escaping_output_rejected(self): + outside = Path(self.tmpdir) / "elsewhere" + outside.mkdir() + link = self.output_dir / "escape" + try: + os.symlink(str(outside), str(link), target_is_directory=True) + except (OSError, NotImplementedError): + self.skipTest("symlinks not available on this platform") + ok, _ = GalleryConfig.validate_gallery_root(str(link)) + self.assertFalse(ok) + + def test_env_root_that_is_filesystem_root_is_ignored(self): + """Listing '/' in the env var must not re-open the whole disk.""" + os.environ[EXTRA_ROOTS_ENV] = os.path.abspath(os.sep) + outside = Path(self.tmpdir) / "elsewhere" + outside.mkdir() + ok, _ = GalleryConfig.validate_gallery_root(str(outside)) + self.assertFalse(ok) + + +class TestValidateGalleryRootAccepts(GalleryRootTestCase): + + def test_output_directory_itself_accepted(self): + ok, reason = GalleryConfig.validate_gallery_root(str(self.output_dir)) + self.assertTrue(ok, reason) + + def test_subdirectory_of_output_accepted(self): + sub = self.output_dir / "renders" / "2026" + sub.mkdir(parents=True) + ok, reason = GalleryConfig.validate_gallery_root(str(sub)) + self.assertTrue(ok, reason) + + def test_subdirectory_of_input_accepted(self): + sub = self.input_dir / "refs" + sub.mkdir() + ok, reason = GalleryConfig.validate_gallery_root(str(sub)) + self.assertTrue(ok, reason) + + def test_directory_listed_in_env_accepted(self): + extra_a = Path(self.tmpdir) / "extra_a" + extra_b = Path(self.tmpdir) / "extra_b" + extra_a.mkdir() + extra_b.mkdir() + os.environ[EXTRA_ROOTS_ENV] = os.pathsep.join([str(extra_a), str(extra_b)]) + + ok_a, reason_a = GalleryConfig.validate_gallery_root(str(extra_a)) + ok_b, reason_b = GalleryConfig.validate_gallery_root(str(extra_b / ".")) + self.assertTrue(ok_a, reason_a) + self.assertTrue(ok_b, reason_b) + + def test_relative_path_is_resolved_before_checking(self): + sub = self.output_dir / "rel" + sub.mkdir() + orig_cwd = os.getcwd() + os.chdir(str(self.output_dir)) + self.addCleanup(os.chdir, orig_cwd) + ok, reason = GalleryConfig.validate_gallery_root("rel") + self.assertTrue(ok, reason) + + def test_accepts_without_folder_paths_when_env_lists_dir(self): + sys.modules.pop("folder_paths", None) + extra = Path(self.tmpdir) / "extra" + extra.mkdir() + os.environ[EXTRA_ROOTS_ENV] = str(extra) + ok, reason = GalleryConfig.validate_gallery_root(str(extra)) + self.assertTrue(ok, reason) + + def test_rejects_everything_without_folder_paths_and_env(self): + sys.modules.pop("folder_paths", None) + ok, _ = GalleryConfig.validate_gallery_root(str(self.output_dir)) + self.assertFalse(ok) + + +class TestGetAllOutputDirsFiltersRoots(GalleryRootTestCase): + """PromptManagerAPI._get_all_output_dirs() only returns validated roots.""" + + def setUp(self): + super().setUp() + from py.api import PromptManagerAPI + + self.api = PromptManagerAPI() + self.api._cached_output_dir = None + self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) + self.addCleanup(self._restore_dirs) + + def _restore_dirs(self): + GalleryConfig.MONITORING_DIRECTORIES = self._orig_dirs + + def _dirs(self): + return [os.path.normcase(str(p)) for p in self.api._get_all_output_dirs()] + + def test_filesystem_root_is_never_served(self): + GalleryConfig.MONITORING_DIRECTORIES = [os.path.abspath(os.sep)] + dirs = self._dirs() + self.assertNotIn(os.path.normcase(os.path.abspath(os.sep)), dirs) + self.assertEqual(dirs, [os.path.normcase(os.path.realpath(self.output_dir))]) + + def test_outside_directory_is_dropped_and_valid_one_kept(self): + outside = Path(self.tmpdir) / "elsewhere" + outside.mkdir() + good = self.output_dir / "good" + good.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(outside), str(good)] + self.assertEqual(self._dirs(), [os.path.normcase(os.path.realpath(good))]) + + def test_all_rejected_falls_back_to_comfyui_output_only(self): + outside = Path(self.tmpdir) / "elsewhere" + outside.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(outside)] + self.assertEqual( + self._dirs(), [os.path.normcase(os.path.realpath(self.output_dir))] + ) + + def test_no_configured_dirs_uses_auto_detection(self): + GalleryConfig.MONITORING_DIRECTORIES = [] + self.api._find_comfyui_output_dir = lambda: str(self.output_dir) + self.assertEqual(self._dirs(), [os.path.normcase(str(self.output_dir))]) + + def test_nothing_available_returns_empty_list(self): + sys.modules.pop("folder_paths", None) + outside = Path(self.tmpdir) / "elsewhere" + outside.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(outside)] + self.assertEqual(self._dirs(), []) + + +if __name__ == "__main__": + unittest.main()