nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: extract downloader validation helpers

Browse Source
This commit is contained in:
Dariusz L
2026-08-02 16:21:58 +02:00
parent 1ce9b77de3
commit 340fbdca59
2 changed files with 158 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/download/validation.py
+147
View File
@@ -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
core/downloader.py
+11 -140
View File
@@ -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,