diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 99b117a..320a0a5 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -9,7 +9,7 @@ import json import re import time from typing import Any, Callable, Dict, List, Optional, Tuple -from urllib.parse import urljoin, urlparse +from urllib.parse import urljoin from ..log_system import create_module_logger from ..matcher import ( @@ -27,7 +27,7 @@ from ..matcher import ( from ..matcher import ( normalize_base_model as _normalize_base_model, ) -from ..network_utils import host_matches_domain, request_source_response +from ..network_utils import request_source_response from ..path_utils import get_filename_from_path from ..progress import get_progress_reporter from ..type_utils import ( @@ -535,31 +535,11 @@ def _archive_link_is_dead(item: Dict[str, Any]) -> bool: ) -def _prepare_size_probe_url(url: Any) -> Optional[str]: +def _fetch_remote_file_size_bytes(url: Any, timeout: int = 15) -> Optional[int]: normalized = _normalize_download_url(url) if not normalized: return None - - normalized = prepare_remote_size_probe_url(normalized) - if not normalized: - return None - - parsed = urlparse(normalized) - host = parsed.hostname - path = parsed.path or "" - - if host_matches_domain(host, "civarchive.com") and not path.startswith("/api/download/"): - return None - if host_matches_domain(host, "civitai.com", "civitai.red") and not path.startswith("/api/download/"): - return None - - return normalized - - - - -def _fetch_remote_file_size_bytes(url: Any, timeout: int = 15) -> Optional[int]: - probe_url = _prepare_size_probe_url(url) + probe_url = prepare_remote_size_probe_url(normalized, ["civarchive.com", "civitai.com", "civitai.red"]) if not probe_url: return None return fetch_remote_file_size_cached(probe_url, headers=REQUEST_HEADERS, timeout=timeout) diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 8b31301..62d8d3e 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -28,7 +28,7 @@ from ..matcher import ( from ..matcher import ( normalize_base_model as _normalize_base_model, ) -from ..network_utils import execute_provider_json_request, request_source_json, request_source_response +from ..network_utils import execute_provider_json_request, request_source_response from ..path_utils import ( calculate_file_sha256, extract_safetensors_header_metadata, @@ -1056,12 +1056,12 @@ def search_civitai( if civitai_type: params["types"] = civitai_type - headers = {} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - - data = request_source_json( - f"{CIVITAI_API_URL}/models", params=params, headers=headers, timeout=15, log_name="CivitAI search models" + data = execute_provider_json_request( + "CivitAI search models", + f"{CIVITAI_API_URL}/models", + params=params, + api_key=api_key, + timeout=15, ) if data: diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index a0f8aa8..c50efc6 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -17,7 +17,7 @@ from ..progress import get_progress_reporter log = create_module_logger(__name__) from ..matcher import build_filename_search_queries -from ..network_utils import host_matches_domain, request_source_json, request_source_response +from ..network_utils import execute_provider_json_request, host_matches_domain, request_source_response from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic from ..type_utils import ( build_search_result, @@ -696,14 +696,19 @@ def search_huggingface_for_file( query_count=total_queries, ) search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20" - response = request_source_response(search_url, headers=headers, timeout=10, log_name="HuggingFace models search") - if response.status_code != 200: + repos = execute_provider_json_request( + "HuggingFace models search", + search_url, + api_key=token, + timeout=10, + ) + if repos is None: log.warning( - f"HuggingFace API status={response.status_code} query={search_query}" + f"HuggingFace API query={search_query} failed" ) continue - - repos = response.json() + if not repos: + continue log.info(f"HuggingFace API query={search_query} repos={len(repos)}") for repo in repos: @@ -970,14 +975,14 @@ def search_huggingface( results = [] try: - headers = {} - if token: - headers["Authorization"] = f"Bearer {token}" - params = {"search": query, "limit": limit, "full": "true"} - data = request_source_json( - f"{HF_API_URL}/models", params=params, headers=headers, timeout=15, log_name="HuggingFace general search" + data = execute_provider_json_request( + "HuggingFace general search", + f"{HF_API_URL}/models", + params=params, + api_key=token, + timeout=15, ) if data: diff --git a/core/type_utils.py b/core/type_utils.py index 3b3b132..316e8bf 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -893,7 +893,7 @@ def extract_trained_words(*values: Any) -> List[str]: _remote_size_cache: Dict[tuple[str, Optional[str]], Optional[int]] = {} -def prepare_remote_size_probe_url(url: Any) -> Optional[str]: +def prepare_remote_size_probe_url(url: Any, allowed_domains: Optional[List[str]] = None) -> Optional[str]: """ Normalize a remote URL to be suitable for file size probing. Converts HuggingFace '/blob/' URLs to '/resolve/' URLs. @@ -916,6 +916,12 @@ def prepare_remote_size_probe_url(url: Any) -> Optional[str]: except Exception: return None + if allowed_domains and host: + if not host_matches_domain(host, *allowed_domains): + return None + if host_matches_domain(host, "civarchive.com", "civitai.com", "civitai.red") and not path.startswith("/api/download/"): + return None + if host and host_matches_domain(host, "huggingface.co") and "/blob/" in path: value = value.replace("/blob/", "/resolve/", 1) diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index 7bfc7ef..5bd834d 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -283,6 +283,44 @@ class TestRefactoringUnification(unittest.TestCase): self.assertEqual(data, {"key": "value"}) mock_request_source_json.assert_called_once() + def test_prepare_remote_size_probe_url_with_allowed_domains(self): + from core.type_utils import prepare_remote_size_probe_url + + # Test default behavior + self.assertEqual( + prepare_remote_size_probe_url("https://huggingface.co/user/repo/blob/main/model.safetensors"), + "https://huggingface.co/user/repo/resolve/main/model.safetensors" + ) + + # Test allowed domains filtering + self.assertEqual( + prepare_remote_size_probe_url("https://civitai.com/api/download/models/123", ["civitai.com", "civitai.red"]), + "https://civitai.com/api/download/models/123" + ) + + self.assertIsNone( + prepare_remote_size_probe_url("https://otherdomain.com/model.safetensors", ["civitai.com", "civitai.red"]) + ) + + @patch("core.network_utils.request_source_json") + def test_execute_provider_json_request(self, mock_request_source_json): + from core.network_utils import execute_provider_json_request + mock_request_source_json.return_value = {"response": "ok"} + + # Test Bearer Authorization and Cookie header injection + res = execute_provider_json_request( + "Test Provider", + "https://example.com/api", + api_key="my_secret_token", + session_token="my_session" + ) + + self.assertEqual(res, {"response": "ok"}) + mock_request_source_json.assert_called_once() + call_args = mock_request_source_json.call_args + headers = call_args[1]["headers"] + self.assertEqual(headers["Authorization"], "Bearer my_secret_token") + self.assertIn("__Secure-civ-token=my_session", headers["Cookie"]) if __name__ == "__main__": diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 0730bf6..8091ac1 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -28,8 +28,8 @@ from core.sources.civarchive import ( _collect_download_urls_unified, _normalize_archive_version, parse_civarchive_url, - _prepare_size_probe_url, ) +from core.type_utils import prepare_remote_size_probe_url # --------------------------------------------------------------------------- # civitai internals @@ -462,24 +462,27 @@ class ParseCivarchiveUrlTests(unittest.TestCase): # =========================================================================== class PrepareSizeProbeUrlTests(unittest.TestCase): + def _prepare_size_probe_url(self, url): + return prepare_remote_size_probe_url(url, ["civarchive.com", "civitai.com", "civitai.red"]) + def test_civarchive_api_download_url_passes(self): url = "https://civarchive.com/api/download/models/77.safetensors" - self.assertEqual(_prepare_size_probe_url(url), url) + self.assertEqual(self._prepare_size_probe_url(url), url) def test_civarchive_non_download_url_blocked(self): url = "https://civarchive.com/models/77" - self.assertIsNone(_prepare_size_probe_url(url)) + self.assertIsNone(self._prepare_size_probe_url(url)) def test_civitai_non_download_url_blocked(self): url = "https://civitai.com/models/123" - self.assertIsNone(_prepare_size_probe_url(url)) + self.assertIsNone(self._prepare_size_probe_url(url)) def test_civitai_api_download_url_passes(self): url = "https://civitai.com/api/download/models/99.safetensors" - self.assertEqual(_prepare_size_probe_url(url), url) + self.assertEqual(self._prepare_size_probe_url(url), url) def test_none_returns_none(self): - self.assertIsNone(_prepare_size_probe_url(None)) + self.assertIsNone(self._prepare_size_probe_url(None)) diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index 5940349..d249004 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -5,34 +5,7 @@ import { getSvgIcon } from "../../utils/icon_utils.js"; import { LOG_LEVEL as DEFAULT_FRONTEND_LOG_LEVEL } from "../../log_system/config.js"; import { logger as frontendLogger } from "../../log_system/logger.js"; import { pollBackgroundTask, safeStorage } from "../utils/html_utils.js"; -const SETTINGS_MAP = [ - { serverKey: 'civitai_key', localKey: 'ModelResolver.civitaiApiKey', type: 'string', default: '' }, - { serverKey: 'civitai_session_token', localKey: 'ModelResolver.civitaiSessionToken', type: 'string', default: '' }, - { serverKey: 'civitai_use_trpc_search', localKey: 'ModelResolver.civitaiUseTrpcSearch', type: 'boolean', default: true }, - { serverKey: 'civitai_use_api_search', localKey: 'ModelResolver.civitaiUseApiSearch', type: 'boolean', default: true }, - { serverKey: 'civitai_use_html_fallback', localKey: 'ModelResolver.civitaiUseHtmlFallback', type: 'boolean', default: true }, - { serverKey: 'hf_token', localKey: 'ModelResolver.huggingFaceToken', type: 'string', default: '' }, - { serverKey: 'brave_search_api_key', localKey: 'ModelResolver.braveSearchApiKey', type: 'string', default: '' }, - { serverKey: 'hf_use_api_search', localKey: 'ModelResolver.hfUseApiSearch', type: 'boolean', default: true }, - { serverKey: 'hf_use_comfy_org_fallback', localKey: 'ModelResolver.hfUseComfyOrgFallback', type: 'boolean', default: true }, - { serverKey: 'hf_use_brave_fallback', localKey: 'ModelResolver.hfUseBraveFallback', type: 'boolean', default: true }, - { serverKey: 'auto_fill_base_model', localKey: 'ModelResolver.autoFillBaseModel', type: 'boolean', default: true }, - { serverKey: 'auto_fill_subfolder', localKey: 'ModelResolver.autoFillSubfolder', type: 'boolean', default: true }, - { serverKey: 'auto_refresh_comfy_models_after_apply', localKey: 'ModelResolver.autoRefreshComfyModelsAfterApply', type: 'boolean', default: true }, - { serverKey: 'workflow_hash_metadata_enabled', localKey: 'ModelResolver.workflowHashMetadataEnabled', type: 'boolean', default: true }, - { serverKey: 'workflow_dependency_marker_enabled', localKey: 'ModelResolver.workflowDependencyMarkerEnabled', type: 'boolean', default: false }, - { serverKey: 'download_backend', localKey: 'ModelResolver.downloadBackend', type: 'backend', default: 'python' }, - { serverKey: 'aria2c_path', localKey: 'ModelResolver.aria2cPath', type: 'string', default: '' }, - { serverKey: 'aria2_auto_stop_daemon', localKey: 'ModelResolver.aria2AutoStopDaemon', type: 'boolean', default: true }, - { serverKey: 'download_path_mode', localKey: 'ModelResolver.downloadPathMode', type: 'pathMode', default: 'suggested' }, - { serverKey: 'download_path_templates', localKey: 'ModelResolver.downloadPathTemplates', type: 'json', default: {} }, - { serverKey: 'base_model_path_mappings', localKey: 'ModelResolver.baseModelPathMappings', type: 'json', default: {} }, - { serverKey: 'frontend_logs_enabled', localKey: 'ModelResolver.frontendLogsEnabled', type: 'frontendLogsEnabled', default: true }, - { serverKey: 'backend_logs_enabled', localKey: 'ModelResolver.backendLogsEnabled', type: 'backendLogsEnabled', default: true }, - { serverKey: 'frontend_log_level', localKey: 'ModelResolver.frontendLogLevel', type: 'string', default: 'ERROR' }, - { serverKey: 'backend_log_level', localKey: 'ModelResolver.backendLogLevel', type: 'string', default: 'ERROR' }, - { serverKey: 'civitai_candidate_limit', localKey: 'ModelResolver.civitaiCandidateLimit', type: 'candidateLimit', default: 5 }, -]; +const SETTINGS_MAP = []; export const optionsMethods = { displayOptions() {