From 340fbdca597116e38004afbc3d9781e5a84cb73d Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sun, 2 Aug 2026 16:21:58 +0200 Subject: [PATCH] refactor: extract downloader validation helpers --- core/download/validation.py | 147 +++++++++++++++++++++++++++++++++++ core/downloader.py | 151 +++--------------------------------- 2 files changed, 158 insertions(+), 140 deletions(-) create mode 100644 core/download/validation.py diff --git a/core/download/validation.py b/core/download/validation.py new file mode 100644 index 0000000..4502a53 --- /dev/null +++ b/core/download/validation.py @@ -0,0 +1,147 @@ +"""Validation and sanitization helpers for model downloads.""" + +import os +import re +from typing import Any, Dict, Optional +from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse + +from ..network_utils import host_matches_domain +from ..path_utils import get_filename_from_path +from ..type_utils import DEFAULT_BROWSER_USER_AGENT, MODEL_EXTENSIONS + +DOWNLOAD_USER_AGENT = DEFAULT_BROWSER_USER_AGENT + +SENSITIVE_METADATA_KEYS = { + "authorization", + "headers", + "hf_token", + "civitai_key", + "api_key", + "apikey", + "access_token", + "token", + "session", + "session_token", + "cookie", + "cookies", +} + +SENSITIVE_QUERY_KEYS = { + "authorization", + "auth", + "hf_token", + "civitai_key", + "api_key", + "apikey", + "access_token", + "token", + "session", + "sessionid", + "cookie", +} + +_INVALID_DOWNLOAD_FILENAME_RE = re.compile(r'[<>:"/\\|?*\x00-\x1f]+') +_HTTP_URL_IN_TEXT_RE = re.compile(r"https?://[^\s]+", re.IGNORECASE) + + +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_allowed_model_download_filename(filename: Any) -> bool: + """Return True only for model file extensions supported by the resolver.""" + safe_name = sanitize_download_filename(filename) + return bool(safe_name and os.path.splitext(safe_name)[1].lower() in MODEL_EXTENSIONS) + + +def _is_sensitive_metadata_key(key: Any) -> bool: + key_text = str(key or "").strip().lower() + return ( + key_text in SENSITIVE_METADATA_KEYS + or "token" in key_text + or "authorization" in key_text + or "cookie" in key_text + ) + + +def _strip_sensitive_url_params(value: str) -> str: + parsed = urlparse(value) + if parsed.scheme not in {"http", "https"} or not parsed.query: + return value + + filtered = [] + changed = False + for key, item_value in parse_qsl(parsed.query, keep_blank_values=True): + key_lower = key.lower() + if key_lower in SENSITIVE_QUERY_KEYS or "token" in key_lower: + changed = True + continue + filtered.append((key, item_value)) + + if not changed: + return value + return urlunparse(parsed._replace(query=urlencode(filtered, doseq=True))) + + +def _sanitize_download_error(value: Any) -> str: + """Remove signed query strings and credentials from errors shown or logged.""" + text = str(value or "") + + def redact_url(match: re.Match) -> str: + raw_url = match.group(0) + try: + parsed = urlparse(raw_url) + if parsed.query: + return urlunparse(parsed._replace(query="", fragment="")) + except Exception: + pass + return raw_url + + return _HTTP_URL_IN_TEXT_RE.sub(redact_url, text) + + +def _clean_http_header_value(value: Any) -> str: + return str(value or "").replace("\r", "").replace("\n", "").strip() + + +def _get_header_value(headers: Dict[str, str], key: str) -> str: + key_lower = key.lower() + for existing_key, value in headers.items(): + if str(existing_key).lower() == key_lower: + return str(value or "") + return "" + + +def _set_header_default(headers: Dict[str, str], key: str, value: str) -> None: + if not _get_header_value(headers, key): + headers[key] = value + + +def build_download_headers( + url: str, + headers: Optional[Dict[str, str]] = None, +) -> Dict[str, str]: + """Build request headers shared by the Python and aria2 download backends.""" + request_headers: Dict[str, str] = {} + for key, value in (headers or {}).items(): + clean_key = _clean_http_header_value(key) + clean_value = _clean_http_header_value(value) + if clean_key and clean_value: + request_headers[clean_key] = clean_value + + _set_header_default(request_headers, "User-Agent", DOWNLOAD_USER_AGENT) + _set_header_default(request_headers, "Accept", "*/*") + _set_header_default(request_headers, "Accept-Encoding", "identity") + + host = urlparse(str(url or "")).hostname + if host_matches_domain(host, "civitai.com", "civitai.red"): + _set_header_default(request_headers, "Referer", "https://civitai.com/") + _set_header_default(request_headers, "Origin", "https://civitai.com") + + return request_headers diff --git a/core/downloader.py b/core/downloader.py index 62c2609..f2f97b8 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -6,7 +6,6 @@ Handles downloading models from various sources with progress tracking. import hashlib import os -import re import secrets import shutil import socket @@ -18,7 +17,7 @@ from collections import deque from io import BytesIO from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple -from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse +from urllib.parse import urlparse, urlunparse import requests @@ -47,8 +46,6 @@ from .path_utils import ( from .resolver import invalidate_local_hash_match_cache, normalize_sha256 from .scanner import invalidate_model_files_cache from .type_utils import ( - DEFAULT_BROWSER_USER_AGENT, - MODEL_EXTENSIONS, as_dict, as_list, extract_response_file_size, @@ -88,7 +85,6 @@ ARIA2_RPC_TIMEOUT = (2, 5) # local JSON-RPC should respond quickly ARIA2_STATUS_RPC_RETRIES = 4 ARIA2_STATUS_RPC_RETRY_DELAY = 0.15 ARIA2_IDLE_STOP_SECONDS = 5 * 60 -DOWNLOAD_USER_AGENT = DEFAULT_BROWSER_USER_AGENT MODEL_PREVIEW_WIDTH = 480 MODEL_PREVIEW_MAX_HEIGHT = 4096 MODEL_PREVIEW_QUALITY = 85 @@ -115,6 +111,16 @@ HF_XET_ARIA2_AUTH_HOSTS = { "cas-bridge-direct.xethub-eu.hf.co", } +from .download.validation import ( + DOWNLOAD_USER_AGENT, + _get_header_value, + _is_sensitive_metadata_key, + _sanitize_download_error, + _strip_sensitive_url_params, + build_download_headers, + is_allowed_model_download_filename, + sanitize_download_filename, +) from .settings import ( load_settings, normalize_download_backend, @@ -122,141 +128,6 @@ from .settings import ( normalize_relative_subfolder, ) -SENSITIVE_METADATA_KEYS = { - "authorization", - "headers", - "hf_token", - "civitai_key", - "api_key", - "apikey", - "access_token", - "token", - "session", - "session_token", - "cookie", - "cookies", -} - -SENSITIVE_QUERY_KEYS = { - "authorization", - "auth", - "hf_token", - "civitai_key", - "api_key", - "apikey", - "access_token", - "token", - "session", - "sessionid", - "cookie", -} - -_INVALID_DOWNLOAD_FILENAME_RE = re.compile(r'[<>:"/\\|?*\x00-\x1f]+') -_HTTP_URL_IN_TEXT_RE = re.compile(r"https?://[^\s]+", re.IGNORECASE) - - -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_allowed_model_download_filename(filename: Any) -> bool: - """Return True only for model file extensions supported by the resolver.""" - safe_name = sanitize_download_filename(filename) - return bool(safe_name and os.path.splitext(safe_name)[1].lower() in MODEL_EXTENSIONS) - - -def _is_sensitive_metadata_key(key: Any) -> bool: - key_text = str(key or "").strip().lower() - return ( - key_text in SENSITIVE_METADATA_KEYS - or "token" in key_text - or "authorization" in key_text - or "cookie" in key_text - ) - - -def _strip_sensitive_url_params(value: str) -> str: - parsed = urlparse(value) - if parsed.scheme not in {"http", "https"} or not parsed.query: - return value - - filtered = [] - changed = False - for key, item_value in parse_qsl(parsed.query, keep_blank_values=True): - key_lower = key.lower() - if key_lower in SENSITIVE_QUERY_KEYS or "token" in key_lower: - changed = True - continue - filtered.append((key, item_value)) - - if not changed: - return value - return urlunparse(parsed._replace(query=urlencode(filtered, doseq=True))) - - -def _sanitize_download_error(value: Any) -> str: - """Remove signed query strings and credentials from errors shown or logged.""" - text = str(value or "") - - def redact_url(match: re.Match) -> str: - raw_url = match.group(0) - try: - parsed = urlparse(raw_url) - if parsed.query: - return urlunparse(parsed._replace(query="", fragment="")) - except Exception: - pass - return raw_url - - return _HTTP_URL_IN_TEXT_RE.sub(redact_url, text) - - -def _clean_http_header_value(value: Any) -> str: - return str(value or "").replace("\r", "").replace("\n", "").strip() - - -def _get_header_value(headers: Dict[str, str], key: str) -> str: - key_lower = key.lower() - for existing_key, value in headers.items(): - if str(existing_key).lower() == key_lower: - return str(value or "") - return "" - - -def _set_header_default(headers: Dict[str, str], key: str, value: str) -> None: - if not _get_header_value(headers, key): - headers[key] = value - - -def build_download_headers( - url: str, - headers: Optional[Dict[str, str]] = None, -) -> Dict[str, str]: - """Build request headers shared by the Python and aria2 download backends.""" - request_headers: Dict[str, str] = {} - for key, value in (headers or {}).items(): - clean_key = _clean_http_header_value(key) - clean_value = _clean_http_header_value(value) - if clean_key and clean_value: - request_headers[clean_key] = clean_value - - _set_header_default(request_headers, "User-Agent", DOWNLOAD_USER_AGENT) - _set_header_default(request_headers, "Accept", "*/*") - _set_header_default(request_headers, "Accept-Encoding", "identity") - - host = urlparse(str(url or "")).hostname - if host_matches_domain(host, "civitai.com", "civitai.red"): - _set_header_default(request_headers, "Referer", "https://civitai.com/") - _set_header_default(request_headers, "Origin", "https://civitai.com") - - return request_headers - def _resolve_download_url_for_aria2( url: str,