diff --git a/py/api/lora_integration.py b/py/api/lora_integration.py index 6275bfd..9630e78 100644 --- a/py/api/lora_integration.py +++ b/py/api/lora_integration.py @@ -76,7 +76,7 @@ class LoraIntegrationMixin: "enabled": config["enabled"], "path": config["path"], "trigger_words_enabled": config["trigger_words_enabled"], - "civitai_api_key": config.get("civitai_api_key", ""), + "has_civitai_api_key": bool(config.get("civitai_api_key")), "detected": detected_path is not None, "detected_path": detected_path or "", "trigger_cache_loaded": cache.is_loaded, diff --git a/py/lora_utils.py b/py/lora_utils.py index 62c925f..f23cc07 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -15,6 +15,7 @@ import threading import urllib.request from pathlib import Path from typing import Dict, List, Optional, Tuple +from urllib.parse import urlsplit try: from ..utils.logging_config import get_logger @@ -298,16 +299,67 @@ def get_preview_image_from_metadata( _THUMB_MAX_SIZE = 512 +# Only these hosts may ever receive the CivitAI API key, and only over HTTPS. +CIVITAI_HOSTS = frozenset( + {"civitai.com", "www.civitai.com", "api.civitai.com", "image.civitai.com"} +) +MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024 +_DOWNLOAD_CHUNK_BYTES = 64 * 1024 + + +def is_civitai_url(url: str) -> bool: + """True when ``url`` is an HTTPS URL on an allow-listed CivitAI host.""" + try: + parts = urlsplit(url) + except ValueError: + return False + hostname = (parts.hostname or "").lower() + return parts.scheme == "https" and hostname in CIVITAI_HOSTS + + +def _read_capped(resp, cap: int) -> Optional[bytes]: + """Read a response body, refusing anything larger than ``cap`` bytes.""" + declared = resp.headers.get("Content-Length") + if declared is not None: + try: + if int(declared) > cap: + return None + except ValueError: + pass + + chunks = [] + total = 0 + while True: + chunk = resp.read(_DOWNLOAD_CHUNK_BYTES) + if not chunk: + break + total += len(chunk) + if total > cap: + return None + chunks.append(chunk) + return b"".join(chunks) + def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]: - """Download a single image, resize to thumbnail, save as JPEG.""" + """Download a single CivitAI image, resize to thumbnail, save as JPEG. + + Refuses any URL that is not HTTPS on an allow-listed CivitAI host so the + API key is never sent elsewhere, and refuses bodies larger than + MAX_CIVITAI_DOWNLOAD_BYTES. + """ + if not is_civitai_url(url): + logger.warning(f"Refusing non-CivitAI download URL: {url}") + return None try: headers = {"User-Agent": "ComfyUI-PromptManager/1.0"} if api_key: headers["Authorization"] = f"Bearer {api_key}" req = urllib.request.Request(url, headers=headers) with urllib.request.urlopen(req, timeout=10) as resp: - raw = resp.read() + raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES) + if raw is None: + logger.warning(f"Refusing oversized download: {url}") + return None # Resize to thumbnail to save disk space from io import BytesIO @@ -361,7 +413,7 @@ def download_civitai_images( if not isinstance(img, dict): continue url = img.get("url", "") - if not isinstance(url, str) or not url.startswith("http"): + if not isinstance(url, str) or not is_civitai_url(url): continue url_hash = hashlib.md5(url.encode()).hexdigest()[:12] diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py new file mode 100644 index 0000000..90a74b2 --- /dev/null +++ b/tests/test_api_lora.py @@ -0,0 +1,119 @@ +"""Route tests for the LoraManager integration endpoints. + +Filesystem detection and network downloads are patched out so the tests run +anywhere without a LoraManager install or a CivitAI connection. +""" + +import json +import os +import shutil +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +if "folder_paths" not in sys.modules: + _stub = types.ModuleType("folder_paths") + _stub.models_dir = tempfile.gettempdir() + _stub.base_path = tempfile.gettempdir() + _stub.get_output_directory = lambda: tempfile.gettempdir() + _stub.get_folder_paths = lambda name: [] + sys.modules["folder_paths"] = _stub + +if "server" not in sys.modules: + _server = MagicMock() + _server.PromptServer.instance.routes = MagicMock() + sys.modules["server"] = _server + +from aiohttp import web +from aiohttp.test_utils import AioHTTPTestCase + +from database.operations import PromptDatabase +from py.api import PromptManagerAPI +from py.config import IntegrationConfig +from py.lora_utils import get_trigger_cache + +SECRET_KEY = "sk-civitai-super-secret-0123456789" + + +class LoraAPITestCase(AioHTTPTestCase): + """Stands up the PromptManager routes with a temp DB and neutral LoRA config.""" + + async def get_application(self): + self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") + self._temp_db.close() + self.tmp_root = Path(tempfile.mkdtemp()) + + self._config_patch = patch.multiple( + IntegrationConfig, + LORA_MANAGER_ENABLED=False, + LORA_MANAGER_PATH="", + LORA_TRIGGER_WORDS_ENABLED=False, + CIVITAI_API_KEY="", + ) + self._config_patch.start() + self._detect_patch = patch( + "py.lora_utils.detect_lora_manager", return_value=None + ) + self.detect_mock = self._detect_patch.start() + get_trigger_cache().clear() + + app = web.Application() + routes = web.RouteTableDef() + self.api = PromptManagerAPI() + self.api.db = PromptDatabase(self._temp_db.name) + self.api.add_routes(routes) + app.router.add_routes(routes) + return app + + async def tearDownAsync(self): + self._detect_patch.stop() + self._config_patch.stop() + get_trigger_cache().clear() + for path in ( + self._temp_db.name, + self._temp_db.name + "-wal", + self._temp_db.name + "-shm", + ): + if os.path.exists(path): + os.unlink(path) + shutil.rmtree(self.tmp_root, ignore_errors=True) + + async def _post_json(self, url, body): + return await self.client.request( + "POST", + url, + data=json.dumps(body), + headers={"Content-Type": "application/json"}, + ) + + +class TestLoraStatusNeverReturnsKey(LoraAPITestCase): + + async def test_status_reports_flag_without_the_key(self): + IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY + + resp = await self.client.request("GET", "/prompt_manager/lora/status") + + self.assertEqual(resp.status, 200) + body = await resp.text() + data = json.loads(body) + self.assertTrue(data["success"]) + self.assertIs(data["has_civitai_api_key"], True) + self.assertNotIn("civitai_api_key", data) + self.assertNotIn(SECRET_KEY, body) + + async def test_status_flag_is_false_without_a_key(self): + resp = await self.client.request("GET", "/prompt_manager/lora/status") + + data = await resp.json() + self.assertIs(data["has_civitai_api_key"], False) + self.assertNotIn("civitai_api_key", data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lora_utils.py b/tests/test_lora_utils.py index 3f355c3..ed22377 100644 --- a/tests/test_lora_utils.py +++ b/tests/test_lora_utils.py @@ -7,6 +7,7 @@ directory detection, TriggerWordCache, and image download logic. import json import os +import shutil import sys import tempfile import threading @@ -16,12 +17,17 @@ from unittest.mock import patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from py import lora_utils from py.lora_utils import ( + CIVITAI_HOSTS, + MAX_CIVITAI_DOWNLOAD_BYTES, TriggerWordCache, + download_civitai_images, get_civitai_image_urls, get_example_prompt_from_metadata, get_lora_image_cache_dir, get_trigger_words_from_metadata, + is_civitai_url, read_lora_metadata, ) @@ -224,6 +230,207 @@ class TestGetLoraImageCacheDir(unittest.TestCase): self.assertTrue(cache_dir.is_dir()) +# ── CivitAI download tests (no network) ─────────────────────────────── + + +def _png_bytes(): + """A valid 2x2 PNG so PIL can decode the fake download.""" + from io import BytesIO + + from PIL import Image + + buf = BytesIO() + Image.new("RGB", (2, 2), (10, 20, 30)).save(buf, "PNG") + return buf.getvalue() + + +class _FakeResponse: + """Minimal stand-in for the object returned by urllib.request.urlopen.""" + + def __init__(self, body, content_length=None): + self._body = body + self._pos = 0 + self.headers = {} + if content_length is not None: + self.headers["Content-Length"] = str(content_length) + + def read(self, amt=-1): + if amt is None or amt < 0: + chunk = self._body[self._pos :] + self._pos = len(self._body) + return chunk + chunk = self._body[self._pos : self._pos + amt] + self._pos += len(chunk) + return chunk + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +class _FakeOpener: + """Records every request and serves a canned response.""" + + def __init__(self, response): + self.response = response + self.requests = [] + + def __call__(self, req, timeout=None): + self.requests.append(req) + return self.response + + +class TestCivitaiUrlPolicy(unittest.TestCase): + + def test_allow_list_contains_expected_hosts(self): + self.assertEqual( + CIVITAI_HOSTS, + frozenset( + { + "civitai.com", + "www.civitai.com", + "api.civitai.com", + "image.civitai.com", + } + ), + ) + self.assertEqual(MAX_CIVITAI_DOWNLOAD_BYTES, 50 * 1024 * 1024) + + def test_https_civitai_hosts_are_allowed(self): + for host in CIVITAI_HOSTS: + self.assertTrue(is_civitai_url(f"https://{host}/api/download/1"), host) + + def test_other_hosts_and_schemes_are_refused(self): + for url in ( + "https://evil.example/x", + "http://civitai.com/api/download/1", + "https://civitai.com.evil.example/x", + "https://notcivitai.com/x", + "ftp://civitai.com/x", + "not a url", + "", + ): + self.assertFalse(is_civitai_url(url), url) + + +class TestDownloadOne(unittest.TestCase): + + def setUp(self): + self.tmpdir = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmpdir, True) + self.target = self.tmpdir / "out.jpg" + + def _run(self, url, opener, api_key="secret-key"): + with patch("urllib.request.urlopen", opener): + return lora_utils._download_one(url, self.target, api_key) + + def test_off_host_url_is_refused_without_a_request(self): + opener = _FakeOpener(_FakeResponse(_png_bytes())) + + result = self._run("https://evil.example/x", opener) + + self.assertIsNone(result) + self.assertEqual(opener.requests, []) + self.assertFalse(self.target.exists()) + + def test_plain_http_civitai_is_refused(self): + opener = _FakeOpener(_FakeResponse(_png_bytes())) + + result = self._run("http://civitai.com/api/download/1", opener) + + self.assertIsNone(result) + self.assertEqual(opener.requests, []) + + def test_civitai_download_sends_bearer_header(self): + opener = _FakeOpener(_FakeResponse(_png_bytes())) + + result = self._run("https://civitai.com/api/download/1", opener) + + self.assertEqual(os.path.realpath(result), os.path.realpath(str(self.target))) + self.assertEqual(len(opener.requests), 1) + req = opener.requests[0] + self.assertEqual(req.get_header("Authorization"), "Bearer secret-key") + self.assertTrue(self.target.exists()) + + def test_no_bearer_header_without_a_key(self): + opener = _FakeOpener(_FakeResponse(_png_bytes())) + + self._run("https://image.civitai.com/a.png", opener, api_key="") + + self.assertIsNone(opener.requests[0].get_header("Authorization")) + + def test_oversized_content_length_is_refused(self): + opener = _FakeOpener( + _FakeResponse(_png_bytes(), content_length=MAX_CIVITAI_DOWNLOAD_BYTES + 1) + ) + + result = self._run("https://civitai.com/api/download/1", opener) + + self.assertIsNone(result) + self.assertFalse(self.target.exists()) + + def test_reading_stops_past_cap_when_length_is_absent(self): + body = b"x" * (MAX_CIVITAI_DOWNLOAD_BYTES + 1) + response = _FakeResponse(body) + opener = _FakeOpener(response) + + with patch.object(lora_utils, "MAX_CIVITAI_DOWNLOAD_BYTES", 1024): + result = self._run("https://civitai.com/api/download/1", opener) + + self.assertIsNone(result) + self.assertLessEqual(response._pos, 1024 + lora_utils._DOWNLOAD_CHUNK_BYTES) + self.assertFalse(self.target.exists()) + + +class TestDownloadCivitaiImages(unittest.TestCase): + + def setUp(self): + self.tmpdir = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmpdir, True) + self.meta_path = self.tmpdir / "lora.safetensors.metadata.json" + self.meta_path.write_text("{}") + + def test_only_civitai_urls_are_fetched(self): + opener = _FakeOpener(_FakeResponse(_png_bytes())) + metadata = _make_metadata( + images=[ + {"url": "https://evil.example/steal"}, + {"url": "https://civitai.com/api/download/1"}, + {"url": "not-a-url"}, + ] + ) + + with patch("urllib.request.urlopen", opener): + paths = download_civitai_images( + metadata, self.meta_path, self.tmpdir / "cache", "key" + ) + + self.assertEqual(len(paths), 1) + self.assertEqual( + [r.full_url for r in opener.requests], + ["https://civitai.com/api/download/1"], + ) + + def test_cached_files_are_not_refetched(self): + opener = _FakeOpener(_FakeResponse(_png_bytes())) + metadata = _make_metadata(images=[{"url": "https://civitai.com/a.png"}]) + cache = self.tmpdir / "cache" + + with patch("urllib.request.urlopen", opener): + first = download_civitai_images(metadata, self.meta_path, cache, "") + second = download_civitai_images(metadata, self.meta_path, cache, "") + + self.assertEqual(first, second) + self.assertEqual(len(opener.requests), 1) + + def test_no_images_returns_empty(self): + self.assertEqual( + download_civitai_images({}, self.meta_path, self.tmpdir / "cache"), [] + ) + + # ── TriggerWordCache tests ────────────────────────────────────────────