Dariusz L
48db08ae1a
- Add core/network_utils.py with URL validation, public-IP enforcement, redirect validation, and safe request_public_url/validate_public_http_url helpers - Replace direct requests calls in downloader/type_utils/sources with request_public_url so redirects are validated before follow-up - Replace netloc string checks with host_matches_domain for CivitAI, CivArchive, and HuggingFace host matching - Add HuggingFace Xet download backend with progress reporting and cancellation support via hf_xet/huggingface_hub - Harden aria2 installer: reject symlinks/special files in archives, restrict custom aria2c paths to the managed install, and update UI copy - Normalize paths with realpath and enforce configured model roots in local hash/metadata routes, plus remove caller-selected metadata_path injection in sidecar writing - Sanitize download/network errors to strip signed query strings before logging or showing them to the user
473 lines
16 KiB
Python
473 lines
16 KiB
Python
"""Install aria2c from official aria2 GitHub releases."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import platform
|
|
import shutil
|
|
import stat
|
|
import tarfile
|
|
import tempfile
|
|
import zipfile
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
import requests
|
|
|
|
from .path_utils import is_path_within
|
|
|
|
|
|
ARIA2_RELEASE_API_URL = "https://api.github.com/repos/aria2/aria2/releases/latest"
|
|
REQUEST_HEADERS = {"User-Agent": "ComfyUI-Model-Resolver"}
|
|
CUSTOM_NODE_ROOT = Path(__file__).resolve().parents[1]
|
|
ARIA2_INSTALL_ROOT = CUSTOM_NODE_ROOT / "tools" / "aria2"
|
|
|
|
|
|
class Aria2InstallError(RuntimeError):
|
|
"""Raised when aria2 cannot be installed automatically."""
|
|
|
|
|
|
def _safe_name(value: Any) -> str:
|
|
text = str(value or "").strip()
|
|
safe = "".join(ch if ch.isalnum() or ch in {"-", "_", "."} else "_" for ch in text)
|
|
return safe.strip("._") or "aria2"
|
|
|
|
|
|
def _normalize_version(value: Any) -> str:
|
|
text = str(value or "").strip()
|
|
lower = text.lower()
|
|
if lower.startswith("release-"):
|
|
text = text[len("release-"):]
|
|
elif lower.startswith("v") and len(text) > 1 and text[1].isdigit():
|
|
text = text[1:]
|
|
return _safe_name(text or "latest")
|
|
|
|
|
|
def _canonical_install_dir(version: Any) -> Path:
|
|
return ARIA2_INSTALL_ROOT / _normalize_version(version)
|
|
|
|
|
|
def _machine_bits(machine: str) -> str:
|
|
machine = str(machine or "").lower()
|
|
if machine in {"amd64", "x86_64", "x64", "arm64", "aarch64"}:
|
|
return "64"
|
|
if machine in {"x86", "i386", "i686"}:
|
|
return "32"
|
|
return "64"
|
|
|
|
|
|
def _platform_tokens() -> Dict[str, str]:
|
|
system = platform.system().lower()
|
|
machine = platform.machine().lower()
|
|
|
|
if system == "windows":
|
|
bits = _machine_bits(machine)
|
|
return {
|
|
"system": "windows",
|
|
"machine": machine,
|
|
"label": f"Windows {bits}-bit",
|
|
"required": "win",
|
|
"arch": f"{bits}bit",
|
|
"exe_name": "aria2c.exe",
|
|
}
|
|
|
|
if system == "darwin":
|
|
return {
|
|
"system": "darwin",
|
|
"machine": machine,
|
|
"label": "macOS",
|
|
"required": "mac",
|
|
"arch": "",
|
|
"exe_name": "aria2c",
|
|
}
|
|
|
|
if system == "linux":
|
|
return {
|
|
"system": "linux",
|
|
"machine": machine,
|
|
"label": f"Linux {machine or 'unknown'}",
|
|
"required": "linux",
|
|
"arch": "",
|
|
"exe_name": "aria2c",
|
|
}
|
|
|
|
raise Aria2InstallError(
|
|
f"Automatic aria2 install is not supported on {platform.system() or 'this OS'}."
|
|
)
|
|
|
|
|
|
def _asset_name(asset: Dict[str, Any]) -> str:
|
|
return str(asset.get("name") or "").strip()
|
|
|
|
|
|
def _asset_url(asset: Dict[str, Any]) -> str:
|
|
return str(asset.get("browser_download_url") or "").strip()
|
|
|
|
|
|
def _is_source_archive(name: str) -> bool:
|
|
lower = name.lower()
|
|
return lower.endswith((".tar.gz", ".tar.bz2", ".tar.xz")) and "build" not in lower
|
|
|
|
|
|
def _score_asset(asset: Dict[str, Any], tokens: Dict[str, str]) -> int:
|
|
name = _asset_name(asset)
|
|
lower = name.lower()
|
|
if not name or not _asset_url(asset):
|
|
return -1
|
|
if _is_source_archive(name):
|
|
return -1
|
|
if tokens["system"] != "linux" and "android" in lower:
|
|
return -1
|
|
|
|
score = 0
|
|
arch = tokens.get("arch") or ""
|
|
|
|
if tokens["system"] == "windows":
|
|
if "win" not in lower:
|
|
return -1
|
|
score += 50
|
|
if arch and arch in lower:
|
|
score += 30
|
|
if lower.endswith(".zip"):
|
|
score += 10
|
|
return score
|
|
|
|
if tokens["system"] == "darwin":
|
|
if not any(token in lower for token in ("mac", "darwin", "osx", "os-x")):
|
|
return -1
|
|
score += 50
|
|
if lower.endswith((".zip", ".tar.gz", ".tar.xz", ".tar.bz2")):
|
|
score += 10
|
|
return score
|
|
|
|
if tokens["system"] == "linux":
|
|
if "linux" not in lower or "android" in lower:
|
|
return -1
|
|
score += 50
|
|
machine = tokens.get("machine") or ""
|
|
if machine and machine in lower:
|
|
score += 10
|
|
if lower.endswith((".zip", ".tar.gz", ".tar.xz", ".tar.bz2")):
|
|
score += 10
|
|
return score
|
|
|
|
return -1
|
|
|
|
|
|
def _select_release_asset(release: Dict[str, Any], tokens: Dict[str, str]) -> Dict[str, Any]:
|
|
assets = release.get("assets") or []
|
|
if not isinstance(assets, list) or not assets:
|
|
raise Aria2InstallError("The latest aria2 release does not expose downloadable assets.")
|
|
|
|
scored = [
|
|
(score, asset)
|
|
for asset in assets
|
|
for score in [_score_asset(asset, tokens)]
|
|
if score >= 0
|
|
]
|
|
if not scored:
|
|
release_name = release.get("tag_name") or release.get("name") or "latest release"
|
|
asset_names = ", ".join(_asset_name(asset) for asset in assets if _asset_name(asset))
|
|
raise Aria2InstallError(
|
|
f"No prebuilt aria2 desktop binary for {tokens['label']} was found in {release_name}. "
|
|
"Install aria2 with your system package manager so aria2c is available on PATH. "
|
|
f"Available assets: {asset_names or 'none'}"
|
|
)
|
|
|
|
scored.sort(key=lambda item: item[0], reverse=True)
|
|
return scored[0][1]
|
|
|
|
|
|
def _fetch_latest_release() -> Dict[str, Any]:
|
|
try:
|
|
response = requests.get(ARIA2_RELEASE_API_URL, headers=REQUEST_HEADERS, timeout=(10, 30))
|
|
response.raise_for_status()
|
|
except requests.RequestException as exc:
|
|
raise Aria2InstallError(
|
|
"Could not contact GitHub to check the latest aria2 release. "
|
|
"Check your internet connection or firewall/proxy settings. "
|
|
f"Details: {exc}"
|
|
) from exc
|
|
data = response.json()
|
|
if not isinstance(data, dict):
|
|
raise Aria2InstallError("GitHub returned an invalid aria2 release response.")
|
|
return data
|
|
|
|
|
|
def _assert_within_install_root(path: Path) -> None:
|
|
install_root = ARIA2_INSTALL_ROOT.resolve()
|
|
target = path.resolve()
|
|
if target == install_root:
|
|
return
|
|
if not is_path_within(str(target), str(install_root)):
|
|
raise Aria2InstallError(f"Refusing to write outside aria2 install root: {target}")
|
|
|
|
|
|
def _safe_extract_zip(archive_path: Path, destination: Path) -> None:
|
|
destination_abs = destination.resolve()
|
|
with zipfile.ZipFile(archive_path) as archive:
|
|
for member in archive.infolist():
|
|
unix_mode = member.external_attr >> 16
|
|
if unix_mode and stat.S_ISLNK(unix_mode):
|
|
raise Aria2InstallError(
|
|
f"Refusing symbolic link in aria2 archive: {member.filename}"
|
|
)
|
|
target = destination_abs / member.filename
|
|
target_abs = target.resolve()
|
|
if not is_path_within(str(target_abs), str(destination_abs)) and target_abs != destination_abs:
|
|
raise Aria2InstallError(f"Unsafe path in aria2 archive: {member.filename}")
|
|
archive.extractall(destination_abs)
|
|
|
|
|
|
def _safe_extract_tar(archive_path: Path, destination: Path) -> None:
|
|
destination_abs = destination.resolve()
|
|
with tarfile.open(archive_path) as archive:
|
|
for member in archive.getmembers():
|
|
if member.issym() or member.islnk():
|
|
raise Aria2InstallError(
|
|
f"Refusing link in aria2 archive: {member.name}"
|
|
)
|
|
if not (member.isfile() or member.isdir()):
|
|
raise Aria2InstallError(
|
|
f"Refusing special file in aria2 archive: {member.name}"
|
|
)
|
|
target = destination_abs / member.name
|
|
target_abs = target.resolve()
|
|
if not is_path_within(str(target_abs), str(destination_abs)) and target_abs != destination_abs:
|
|
raise Aria2InstallError(f"Unsafe path in aria2 archive: {member.name}")
|
|
archive.extractall(destination_abs)
|
|
|
|
|
|
def _extract_archive(archive_path: Path, destination: Path) -> None:
|
|
name = archive_path.name.lower()
|
|
if name.endswith(".zip"):
|
|
_safe_extract_zip(archive_path, destination)
|
|
return
|
|
if name.endswith((".tar.gz", ".tgz", ".tar.xz", ".tar.bz2")):
|
|
_safe_extract_tar(archive_path, destination)
|
|
return
|
|
raise Aria2InstallError(f"Unsupported aria2 archive format: {archive_path.name}")
|
|
|
|
|
|
def _find_executable(root: Path, exe_name: str) -> Optional[Path]:
|
|
if not root.exists():
|
|
return None
|
|
expected = exe_name.lower()
|
|
for path in root.rglob("*"):
|
|
if path.is_file() and path.name.lower() == expected:
|
|
return path
|
|
return None
|
|
|
|
|
|
def _find_existing_install(exe_name: str) -> Optional[Path]:
|
|
candidates = []
|
|
expected = exe_name.lower()
|
|
if not ARIA2_INSTALL_ROOT.exists():
|
|
return None
|
|
for path in ARIA2_INSTALL_ROOT.rglob("*"):
|
|
if path.is_file() and path.name.lower() == expected:
|
|
candidates.append(path)
|
|
if not candidates:
|
|
return None
|
|
|
|
def candidate_key(path: Path) -> tuple:
|
|
try:
|
|
relative_depth = len(path.relative_to(ARIA2_INSTALL_ROOT).parts)
|
|
except ValueError:
|
|
relative_depth = 999
|
|
try:
|
|
modified = path.stat().st_mtime
|
|
except OSError:
|
|
modified = 0
|
|
return (1 if relative_depth <= 2 else 0, modified)
|
|
|
|
candidates.sort(key=candidate_key, reverse=True)
|
|
return candidates[0]
|
|
|
|
|
|
def _existing_install_metadata(executable: Path) -> Dict[str, str]:
|
|
try:
|
|
relative = executable.relative_to(ARIA2_INSTALL_ROOT)
|
|
parts = relative.parts
|
|
except ValueError:
|
|
parts = ()
|
|
|
|
return {
|
|
"version": _normalize_version(parts[0] if parts else ""),
|
|
"asset": parts[1] if len(parts) > 2 else "",
|
|
"install_dir": str(executable.parent),
|
|
}
|
|
|
|
|
|
def _copy_directory_contents(source: Path, destination: Path) -> None:
|
|
source_abs = source.resolve()
|
|
destination_abs = destination.resolve()
|
|
if source_abs == destination_abs:
|
|
return
|
|
_assert_within_install_root(destination_abs)
|
|
destination_abs.mkdir(parents=True, exist_ok=True)
|
|
for child in source_abs.iterdir():
|
|
target = destination_abs / child.name
|
|
_assert_within_install_root(target)
|
|
if child.is_dir():
|
|
if target.exists() and not target.is_dir():
|
|
target.unlink()
|
|
shutil.copytree(child, target, dirs_exist_ok=True)
|
|
else:
|
|
if target.exists() and target.is_dir():
|
|
shutil.rmtree(target)
|
|
shutil.copy2(child, target)
|
|
|
|
|
|
def _cleanup_legacy_install_dirs() -> None:
|
|
if not ARIA2_INSTALL_ROOT.exists():
|
|
return
|
|
for child in ARIA2_INSTALL_ROOT.iterdir():
|
|
if not child.is_dir() or not child.name.lower().startswith("release-"):
|
|
continue
|
|
_assert_within_install_root(child)
|
|
shutil.rmtree(child, ignore_errors=True)
|
|
|
|
|
|
def _migrate_existing_install(executable: Path, exe_name: str) -> Path:
|
|
metadata = _existing_install_metadata(executable)
|
|
version = metadata.get("version") or "latest"
|
|
canonical_dir = _canonical_install_dir(version)
|
|
canonical_executable = canonical_dir / exe_name
|
|
try:
|
|
if executable.resolve() == canonical_executable.resolve():
|
|
return executable
|
|
except OSError:
|
|
pass
|
|
|
|
_assert_within_install_root(canonical_dir)
|
|
if canonical_dir.exists():
|
|
_assert_within_install_root(canonical_dir)
|
|
shutil.rmtree(canonical_dir)
|
|
canonical_dir.mkdir(parents=True, exist_ok=True)
|
|
_copy_directory_contents(executable.parent, canonical_dir)
|
|
|
|
migrated_executable = _find_executable(canonical_dir, exe_name)
|
|
if not migrated_executable:
|
|
raise Aria2InstallError(
|
|
f"Could not migrate existing aria2 install to {canonical_dir}."
|
|
)
|
|
_cleanup_legacy_install_dirs()
|
|
return migrated_executable
|
|
|
|
|
|
def _download_file(url: str, destination: Path) -> int:
|
|
total = 0
|
|
try:
|
|
with requests.get(url, headers=REQUEST_HEADERS, stream=True, timeout=(10, 120)) as response:
|
|
response.raise_for_status()
|
|
with destination.open("wb") as handle:
|
|
for chunk in response.iter_content(chunk_size=1024 * 1024):
|
|
if not chunk:
|
|
continue
|
|
handle.write(chunk)
|
|
total += len(chunk)
|
|
except requests.RequestException as exc:
|
|
raise Aria2InstallError(
|
|
"Could not download the aria2 release asset from GitHub. "
|
|
"Check your internet connection or firewall/proxy settings. "
|
|
f"Details: {exc}"
|
|
) from exc
|
|
return total
|
|
|
|
|
|
def _chmod_executable(path: Path) -> None:
|
|
if platform.system().lower() == "windows":
|
|
return
|
|
try:
|
|
mode = path.stat().st_mode
|
|
path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def install_aria2_engine(force: bool = False) -> Dict[str, Any]:
|
|
"""Download and install aria2c for the current platform."""
|
|
tokens = _platform_tokens()
|
|
if not force:
|
|
existing_executable = _find_existing_install(tokens["exe_name"])
|
|
if existing_executable and existing_executable.exists():
|
|
existing_executable = _migrate_existing_install(existing_executable, tokens["exe_name"])
|
|
_cleanup_legacy_install_dirs()
|
|
_chmod_executable(existing_executable)
|
|
existing_metadata = _existing_install_metadata(existing_executable)
|
|
return {
|
|
"success": True,
|
|
"already_installed": True,
|
|
"aria2c_path": str(existing_executable),
|
|
"version": existing_metadata["version"],
|
|
"asset": existing_metadata["asset"],
|
|
"platform": tokens["label"],
|
|
"install_dir": existing_metadata["install_dir"],
|
|
}
|
|
|
|
release = _fetch_latest_release()
|
|
asset = _select_release_asset(release, tokens)
|
|
tag_name = str(release.get("tag_name") or "latest").strip() or "latest"
|
|
version = _normalize_version(tag_name)
|
|
asset_name = _asset_name(asset)
|
|
asset_url = _asset_url(asset)
|
|
if not asset_url:
|
|
raise Aria2InstallError(f"aria2 release asset {asset_name or '<unknown>'} has no download URL.")
|
|
|
|
install_dir = _canonical_install_dir(version)
|
|
_assert_within_install_root(install_dir)
|
|
existing_executable = _find_executable(install_dir, tokens["exe_name"])
|
|
if existing_executable and existing_executable.exists() and not force:
|
|
_chmod_executable(existing_executable)
|
|
return {
|
|
"success": True,
|
|
"already_installed": True,
|
|
"aria2c_path": str(existing_executable),
|
|
"version": version,
|
|
"asset": asset_name,
|
|
"platform": tokens["label"],
|
|
"install_dir": str(install_dir),
|
|
}
|
|
|
|
ARIA2_INSTALL_ROOT.mkdir(parents=True, exist_ok=True)
|
|
install_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="model-resolver-aria2-") as tmp_dir_text:
|
|
tmp_dir = Path(tmp_dir_text)
|
|
archive_path = tmp_dir / asset_name
|
|
extract_dir = tmp_dir / "extract"
|
|
downloaded_bytes = _download_file(asset_url, archive_path)
|
|
extract_dir.mkdir(parents=True, exist_ok=True)
|
|
_extract_archive(archive_path, extract_dir)
|
|
extracted_executable = _find_executable(extract_dir, tokens["exe_name"])
|
|
if not extracted_executable:
|
|
raise Aria2InstallError(
|
|
f"Installed aria2 archive did not contain {tokens['exe_name']}."
|
|
)
|
|
|
|
if install_dir.exists():
|
|
_assert_within_install_root(install_dir)
|
|
shutil.rmtree(install_dir)
|
|
install_dir.mkdir(parents=True, exist_ok=True)
|
|
_copy_directory_contents(extracted_executable.parent, install_dir)
|
|
|
|
executable = _find_executable(install_dir, tokens["exe_name"])
|
|
if not executable:
|
|
raise Aria2InstallError(
|
|
f"Installed aria2 archive did not contain {tokens['exe_name']}."
|
|
)
|
|
_chmod_executable(executable)
|
|
_cleanup_legacy_install_dirs()
|
|
|
|
return {
|
|
"success": True,
|
|
"already_installed": False,
|
|
"aria2c_path": str(executable),
|
|
"version": version,
|
|
"asset": asset_name,
|
|
"platform": tokens["label"],
|
|
"install_dir": str(install_dir),
|
|
"downloaded_bytes": downloaded_bytes,
|
|
"download_url": asset_url,
|
|
}
|