diff --git a/database/operations.py b/database/operations.py index 0c27b83..bceb9ad 100644 --- a/database/operations.py +++ b/database/operations.py @@ -499,12 +499,22 @@ def get_all_categories(self) -> List[str]: ) return [row["category"] for row in cursor.fetchall()] - def get_prompt_subfolders(self, root_dirs: Optional[List[str]] = None) -> List[str]: + def get_prompt_subfolders( + self, + root_dirs: Optional[List[str]] = None, + include_ancestors: bool = False, + ) -> List[str]: """ Get distinct subfolder paths from generated_images. Extracts the directory portion of image_path, made relative to root_dirs if provided. Returns sorted unique folder names. + + Args: + root_dirs: Gallery root directories for computing relative paths + include_ancestors: If True, also include all ancestor path segments + so that hierarchical tree navigation works. For example, a path + ``2026/08-Aug/2026-08-06`` also adds ``2026`` and ``2026/08-Aug``. """ with self.model.get_connection() as conn: cursor = conn.execute( @@ -533,9 +543,18 @@ def get_prompt_subfolders(self, root_dirs: Optional[List[str]] = None) -> List[s continue if not made_relative: - basename = os.path.basename(parent) - if basename: - folders.add(basename) + # Preserve full relative-style path instead of collapsing + # to just the basename, which loses hierarchy and creates + # ambiguity (e.g. foo/bar and baz/bar both become "bar"). + folders.add(parent) + + if include_ancestors: + ancestors = set() + for folder in folders: + parts = folder.replace("\\", "/").split("/") + for i in range(1, len(parts)): + ancestors.add("/".join(parts[:i])) + folders.update(ancestors) return sorted(folders) @@ -1088,10 +1107,10 @@ def find_duplicates(self) -> List[Dict[str, Any]]: # Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility # We'll sort the IDs manually after fetching cursor = conn.execute(""" - SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count, + SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count, GROUP_CONCAT(id) as ids, GROUP_CONCAT(created_at) as created_dates - FROM prompts + FROM prompts GROUP BY LOWER(TRIM(text)) HAVING COUNT(*) > 1 """) @@ -1160,10 +1179,10 @@ def cleanup_duplicates(self) -> int: # Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility # We'll sort the IDs manually after fetching cursor = conn.execute(""" - SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count, + SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count, GROUP_CONCAT(id) as ids, GROUP_CONCAT(created_at) as created_dates - FROM prompts + FROM prompts GROUP BY LOWER(TRIM(text)) HAVING COUNT(*) > 1 """) diff --git a/py/api/images.py b/py/api/images.py index e6b06e1..320a583 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -209,12 +209,12 @@ async def get_output_images(self, request): offset = int(request.query.get("offset", 0)) subfolder = request.query.get("subfolder", "").strip() - # Collect (path, mtime, root) from all directories + # Collect (path, mtime, root, root_index) from all directories all_images = [] - for output_path in output_dirs: + for root_idx, output_path in enumerate(output_dirs): dir_images = await self._get_gallery_files(output_path) for img_path, mtime in dir_images: - all_images.append((img_path, mtime, output_path)) + all_images.append((img_path, mtime, output_path, root_idx)) # Sort combined results by mtime (newest first) — no .stat() calls all_images.sort(key=lambda x: x[1], reverse=True) @@ -222,10 +222,10 @@ async def get_output_images(self, request): # Apply subfolder filter if provided if subfolder: filtered = [] - for img_path, mtime, root in all_images: + for img_path, mtime, root, ridx in all_images: rel_dir = str(img_path.relative_to(root).parent) if rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep): - filtered.append((img_path, mtime, root)) + filtered.append((img_path, mtime, root, ridx)) all_images = filtered total = len(all_images) @@ -235,7 +235,7 @@ async def get_output_images(self, request): def _format_page(): images = [] - for media_path, mtime, output_path in paginated: + for media_path, mtime, output_path, root_index in paginated: try: stat = media_path.stat() rel_path = media_path.relative_to(output_path) @@ -260,6 +260,7 @@ def _format_page(): "path": str(media_path), "relative_path": str(rel_path), "root_dir": str(output_path), + "root_index": root_index, "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", "thumbnail_url": thumbnail_url, "size": stat.st_size, @@ -364,6 +365,17 @@ async def serve_output_image(self, request): status=404, ) + # If a root index is provided, use only that root to avoid + # path collisions when multiple roots share the same relative path. + root_idx = request.query.get("root") + if root_idx is not None: + try: + idx = int(root_idx) + if 0 <= idx < len(output_dirs): + output_dirs = [output_dirs[idx]] + except (ValueError, IndexError): + pass + # Try each root directory until the file is found for output_path in output_dirs: image_path = (output_path / filepath).resolve() @@ -397,6 +409,17 @@ async def get_gallery_subfolders(self, request): if rel_dir and rel_dir != ".": subfolders.add(rel_dir) + include_ancestors = ( + request.query.get("include_ancestors", "").lower() == "true" + ) + if include_ancestors: + ancestors = set() + for folder in subfolders: + parts = folder.replace("\\", "/").split("/") + for i in range(1, len(parts)): + ancestors.add("/".join(parts[:i])) + subfolders.update(ancestors) + return web.json_response( {"success": True, "subfolders": sorted(subfolders)} ) diff --git a/py/api/prompts.py b/py/api/prompts.py index 1de4c32..bc5632b 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -170,8 +170,11 @@ async def get_subfolders(self, request): from ..config import GalleryConfig root_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) or None + include_ancestors = ( + request.query.get("include_ancestors", "").lower() == "true" + ) subfolders = await self._run_in_executor( - self.db.get_prompt_subfolders, root_dirs + self.db.get_prompt_subfolders, root_dirs, include_ancestors ) return web.json_response({"success": True, "subfolders": subfolders}) except Exception as e: diff --git a/tests/test_lora_database.py b/tests/test_lora_database.py index 946d6d8..2ede390 100644 --- a/tests/test_lora_database.py +++ b/tests/test_lora_database.py @@ -194,6 +194,37 @@ def test_with_root_dirs(self): self.assertIsInstance(folders, list) self.assertGreater(len(folders), 0) + def test_include_ancestors_adds_intermediate_paths(self): + pid = self._save("prompt") + self._link_image(pid, "/output/2026/08-Aug/2026-08-06/img.png") + + folders = self.db.get_prompt_subfolders( + root_dirs=["/output"], include_ancestors=True + ) + self.assertIn("2026", folders) + self.assertIn("2026/08-Aug", folders) + self.assertIn("2026/08-Aug/2026-08-06", folders) + + def test_include_ancestors_false_no_intermediates(self): + pid = self._save("prompt") + self._link_image(pid, "/output/2026/08-Aug/2026-08-06/img.png") + + folders = self.db.get_prompt_subfolders( + root_dirs=["/output"], include_ancestors=False + ) + self.assertIn("2026/08-Aug/2026-08-06", folders) + self.assertNotIn("2026", folders) + self.assertNotIn("2026/08-Aug", folders) + + def test_include_ancestors_sorted(self): + pid = self._save("prompt") + self._link_image(pid, "/output/a/b/c/img.png") + + folders = self.db.get_prompt_subfolders( + root_dirs=["/output"], include_ancestors=True + ) + self.assertEqual(folders, sorted(folders)) + # ── LoRA prompt workflow ────────────────────────────────────────────── diff --git a/tests/test_path_security.py b/tests/test_path_security.py new file mode 100644 index 0000000..1041756 --- /dev/null +++ b/tests/test_path_security.py @@ -0,0 +1,128 @@ +"""Tests for path traversal and symlink security in image serving.""" + +import os +import tempfile +import unittest +from pathlib import Path + + +class TestPathSecurity(unittest.TestCase): + """Verify that path resolution and boundary checks prevent escapes.""" + + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.output_dir = Path(self.tmpdir) / "output" + self.output_dir.mkdir() + self.secret_dir = Path(self.tmpdir) / "secret" + self.secret_dir.mkdir() + + # Create a legitimate image + (self.output_dir / "legit.png").write_bytes(b"PNG") + + # Create a secret file outside the output dir + (self.secret_dir / "password.txt").write_text("hunter2") + + def tearDown(self): + import shutil + + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def _is_safe_path(self, output_path, filepath): + """Reproduce the safety check from serve_output_image.""" + image_path = (output_path / filepath).resolve() + return image_path.is_relative_to(output_path.resolve()) and image_path.exists() + + def test_normal_path_allowed(self): + self.assertTrue(self._is_safe_path(self.output_dir, "legit.png")) + + def test_dot_dot_traversal_blocked(self): + self.assertFalse(self._is_safe_path(self.output_dir, "../secret/password.txt")) + + def test_encoded_dot_dot_blocked(self): + # Even if someone encodes ../ as %2e%2e%2f, Path resolution catches it + self.assertFalse( + self._is_safe_path(self.output_dir, "..%2fsecret%2fpassword.txt") + ) + + def test_absolute_path_blocked(self): + secret_path = str(self.secret_dir / "password.txt") + self.assertFalse(self._is_safe_path(self.output_dir, secret_path)) + + def test_symlink_escape_blocked(self): + # Create a symlink inside output_dir pointing outside + symlink_path = self.output_dir / "escape_link" + try: + symlink_path.symlink_to(self.secret_dir / "password.txt") + except OSError: + self.skipTest("Cannot create symlinks on this platform") + + # resolve() follows the symlink, so it should resolve to + # secret_dir which is NOT relative_to output_dir + image_path = symlink_path.resolve() + self.assertFalse(image_path.is_relative_to(self.output_dir.resolve())) + + def test_symlink_within_root_allowed(self): + # A symlink that stays within the output dir should be fine + subdir = self.output_dir / "subdir" + subdir.mkdir() + (subdir / "img.png").write_bytes(b"PNG") + + symlink_path = self.output_dir / "link_to_sub" + try: + symlink_path.symlink_to(subdir / "img.png") + except OSError: + self.skipTest("Cannot create symlinks on this platform") + + image_path = symlink_path.resolve() + self.assertTrue(image_path.is_relative_to(self.output_dir.resolve())) + + def test_nested_traversal_blocked(self): + # Deep nested traversal: sub/../../secret/password.txt + sub = self.output_dir / "sub" + sub.mkdir() + self.assertFalse( + self._is_safe_path(self.output_dir, "sub/../../secret/password.txt") + ) + + def test_null_byte_in_path(self): + # Null bytes should not bypass checks + try: + result = self._is_safe_path(self.output_dir, "legit.png\x00.txt") + except ValueError: + # Python raises ValueError for embedded null bytes — this is safe + result = False + self.assertFalse(result) + + +class TestRootIndexSecurity(unittest.TestCase): + """Verify that root_index parameter is validated safely.""" + + def test_negative_index_ignored(self): + """Negative root index should not select any directory.""" + output_dirs = [Path("/fake/dir1"), Path("/fake/dir2")] + idx = -1 + # The code checks 0 <= idx < len(output_dirs) + self.assertFalse(0 <= idx < len(output_dirs)) + + def test_out_of_range_index_ignored(self): + output_dirs = [Path("/fake/dir1")] + idx = 5 + self.assertFalse(0 <= idx < len(output_dirs)) + + def test_valid_index_selects_correct_root(self): + output_dirs = [Path("/fake/dir1"), Path("/fake/dir2"), Path("/fake/dir3")] + idx = 1 + self.assertTrue(0 <= idx < len(output_dirs)) + self.assertEqual(output_dirs[idx], Path("/fake/dir2")) + + def test_non_numeric_index_handled(self): + """Non-numeric root value should not crash.""" + try: + int("abc") + self.fail("Should have raised ValueError") + except ValueError: + pass # This is the expected behavior + + +if __name__ == "__main__": + unittest.main() diff --git a/web/gallery.html b/web/gallery.html index f33f104..b4878f2 100644 --- a/web/gallery.html +++ b/web/gallery.html @@ -38,6 +38,11 @@

PromptManager Gallery

images + | + + + +
-
-
- -
-
-
-

Loading gallery...

+
+
+ + + + - -
diff --git a/web/js/admin.js b/web/js/admin.js index 5910bb7..0a7d2ae 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -320,7 +320,7 @@ async loadSubfolders() { try { - const response = await fetch("/prompt_manager/subfolders"); + const response = await fetch("/prompt_manager/subfolders?include_ancestors=true"); if (response.ok) { const data = await response.json(); if (data.success) { @@ -344,7 +344,10 @@ this.subfolders.forEach((folder) => { const option = document.createElement("option"); option.value = folder; - option.textContent = folder; + const depth = folder.split("/").length - 1; + const indent = "\u00A0\u00A0".repeat(depth); + const label = folder.split("/").pop(); + option.textContent = indent + label; select.appendChild(option); }); select.value = current; diff --git a/web/js/gallery.js b/web/js/gallery.js index 49c927b..1598baf 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -6,10 +6,14 @@ this.total = 0; this.viewMode = 'grid'; // 'grid' or 'list' this.viewer = null; - + this.currentSubfolder = ''; + this.expandedFolders = new Set(); + this.thumbSize = 'md'; // sm, md, lg + this.initializeEventListeners(); + this.loadFolderTree(); this.loadImages(); - + // Check thumbnails at startup if enabled this.checkThumbnailsAtStartup(); } @@ -104,14 +108,183 @@ } }); }); + + // Thumbnail size buttons + ['Sm', 'Md', 'Lg'].forEach(size => { + const btn = document.getElementById('thumb' + size + 'Btn'); + if (btn) btn.addEventListener('click', () => this.setThumbSize(size.toLowerCase())); + }); + + // Folder tree events + document.getElementById('clearFolderFilterBtn').addEventListener('click', () => { + this.currentSubfolder = ''; + this.currentPage = 1; + this.renderFolderTree(); + this.loadImages(); + }); + + // Resize handle for folder panel + const resizeHandle = document.getElementById('folderResizeHandle'); + if (resizeHandle) { + let isResizing = false; + resizeHandle.addEventListener('mousedown', (e) => { + isResizing = true; + document.body.style.cursor = 'col-resize'; + document.body.style.userSelect = 'none'; + e.preventDefault(); + }); + document.addEventListener('mousemove', (e) => { + if (!isResizing) return; + const panel = document.getElementById('folderTreePanel'); + const newWidth = Math.max(160, Math.min(500, e.clientX - panel.getBoundingClientRect().left)); + panel.style.width = newWidth + 'px'; + }); + document.addEventListener('mouseup', () => { + if (isResizing) { + isResizing = false; + document.body.style.cursor = ''; + document.body.style.userSelect = ''; + } + }); + } + } + + async loadFolderTree() { + try { + const response = await fetch('/prompt_manager/gallery/subfolders?include_ancestors=true'); + if (!response.ok) return; + const data = await response.json(); + if (data.success && data.subfolders.length > 0) { + this.folderList = data.subfolders; + document.getElementById('folderTreePanel').classList.remove('hidden'); + const rh = document.getElementById('folderResizeHandle'); + if (rh) rh.classList.remove('hidden'); + this.renderFolderTree(); + } + } catch (error) { + console.error('Folder tree error:', error); + } + } + + renderFolderTree() { + const container = document.getElementById('folderTree'); + container.innerHTML = ''; + + // Build tree structure from flat sorted paths + const tree = {}; + for (const path of this.folderList) { + const parts = path.replace(/\\/g, '/').split('/'); + let node = tree; + for (const part of parts) { + if (!node[part]) node[part] = {}; + node = node[part]; + } + } + + const renderNode = (obj, parentPath, depth) => { + const entries = Object.keys(obj).sort(); + for (const name of entries) { + const fullPath = parentPath ? parentPath + '/' + name : name; + const hasChildren = Object.keys(obj[name]).length > 0; + const isExpanded = this.expandedFolders.has(fullPath); + const isActive = this.currentSubfolder === fullPath; + + const row = document.createElement('div'); + row.className = 'flex items-center px-2 py-1 cursor-pointer hover:bg-pm-hover transition-colors' + + (isActive ? ' bg-pm-accent/10 text-pm-accent font-medium' : ' text-pm'); + row.style.paddingLeft = (8 + depth * 16) + 'px'; + + // Toggle arrow + const arrow = document.createElement('span'); + arrow.className = 'inline-block w-4 text-center text-pm-secondary mr-1 select-none'; + if (hasChildren) { + arrow.textContent = isExpanded ? '\u25BE' : '\u25B8'; + arrow.addEventListener('click', (e) => { + e.stopPropagation(); + if (this.expandedFolders.has(fullPath)) { + this.expandedFolders.delete(fullPath); + } else { + this.expandedFolders.add(fullPath); + } + this.renderFolderTree(); + }); + } else { + arrow.textContent = ''; + } + row.appendChild(arrow); + + const label = document.createElement('span'); + label.textContent = name; + label.className = 'truncate'; + row.appendChild(label); + + // Click to filter + row.addEventListener('click', () => { + this.currentSubfolder = fullPath; + this.currentPage = 1; + this.renderFolderTree(); + this.loadImages(); + }); + + container.appendChild(row); + + if (hasChildren && isExpanded) { + renderNode(obj[name], fullPath, depth + 1); + } + } + }; + + // "All" option at top + const allRow = document.createElement('div'); + allRow.className = 'flex items-center px-2 py-1 cursor-pointer hover:bg-pm-hover transition-colors' + + (!this.currentSubfolder ? ' bg-pm-accent/10 text-pm-accent font-medium' : ' text-pm'); + allRow.style.paddingLeft = '8px'; + allRow.innerHTML = 'All Images'; + allRow.addEventListener('click', () => { + this.currentSubfolder = ''; + this.currentPage = 1; + this.renderFolderTree(); + this.loadImages(); + }); + container.appendChild(allRow); + + renderNode(tree, '', 0); + } + + setThumbSize(size) { + this.thumbSize = size; + const grid = document.getElementById('galleryGrid'); + const sizes = { + sm: { cols: 'repeat(auto-fill, minmax(100px, 1fr))', gap: '4px' }, + md: { cols: 'repeat(auto-fill, minmax(180px, 1fr))', gap: '8px' }, + lg: { cols: 'repeat(auto-fill, minmax(300px, 1fr))', gap: '12px' }, + }; + const s = sizes[size] || sizes.md; + grid.style.gridTemplateColumns = s.cols; + grid.style.gap = s.gap; + // Update button states + ['Sm', 'Md', 'Lg'].forEach(s => { + const btn = document.getElementById('thumb' + s + 'Btn'); + if (btn) { + if (s.toLowerCase() === size) { + btn.className = btn.className.replace('bg-pm-input', 'bg-pm-accent').replace('text-pm ', 'text-pm-accent-fg '); + } else { + btn.className = btn.className.replace('bg-pm-accent', 'bg-pm-input').replace('text-pm-accent-fg', 'text-pm'); + } + } + }); } async loadImages() { this.showLoading(); - + try { const offset = (this.currentPage - 1) * this.limit; - const response = await fetch(`/prompt_manager/images/output?limit=${this.limit}&offset=${offset}`); + let url = `/prompt_manager/images/output?limit=${this.limit}&offset=${offset}`; + if (this.currentSubfolder) { + url += `&subfolder=${encodeURIComponent(this.currentSubfolder)}`; + } + const response = await fetch(url); if (!response.ok) { throw new Error(`HTTP ${response.status}: ${response.statusText}`);