474 lines
16 KiB
Python
474 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 .network_utils import fetch_json_from_public_url
|
|
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:
|
|
data = fetch_json_from_public_url(ARIA2_RELEASE_API_URL, headers=REQUEST_HEADERS, timeout=30)
|
|
except Exception 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
|
|
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 _validate_archive_member_path(destination_abs: Path, member_name: str) -> None:
|
|
target_abs = (destination_abs / member_name).resolve()
|
|
if target_abs != destination_abs and not is_path_within(
|
|
str(target_abs), str(destination_abs)
|
|
):
|
|
raise Aria2InstallError(f"Unsafe path in aria2 archive: {member_name}")
|
|
|
|
|
|
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}"
|
|
)
|
|
_validate_archive_member_path(destination_abs, 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}"
|
|
)
|
|
_validate_archive_member_path(destination_abs, 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,
|
|
}
|