from __future__ import annotations import socket import time from urllib.parse import urlparse import requests from app.extensions import db from app.models import Mirror, TestResult from app.services.host import get_ubuntu_codename, run_on_host, validate_codename def _record(mirror: Mirror, test_type: str, success: bool, latency_ms: float | None, error: str | None) -> TestResult: result = TestResult( mirror_id=mirror.id, test_type=test_type, success=success, latency_ms=latency_ms, error=error, ) db.session.add(result) db.session.commit() return result def test_dns(mirror: Mirror) -> TestResult: ips = mirror.get_ips() if not ips: return _record(mirror, "dns_resolve", False, None, "IP تعریف نشده") start = time.perf_counter() try: for domain in ("github.com", "docker.io"): socket.getaddrinfo(domain, 443, type=socket.SOCK_STREAM) latency = (time.perf_counter() - start) * 1000 return _record(mirror, "dns_resolve", True, latency, None) except socket.gaierror as exc: latency = (time.perf_counter() - start) * 1000 return _record(mirror, "dns_resolve", False, latency, str(exc)) def test_tcp(url: str, timeout: int = 10) -> tuple[bool, float | None, str | None]: parsed = urlparse(url if "://" in url else f"https://{url}") host = parsed.hostname port = parsed.port or (443 if parsed.scheme == "https" else 80) if not host: return False, None, "host نامعتبر" start = time.perf_counter() try: sock = socket.create_connection((host, port), timeout=timeout) sock.close() return True, (time.perf_counter() - start) * 1000, None except OSError as exc: return False, (time.perf_counter() - start) * 1000, str(exc) def test_http(url: str, timeout: int = 15) -> tuple[bool, float | None, str | None]: start = time.perf_counter() try: response = requests.head(url.rstrip("/"), timeout=timeout, allow_redirects=True) if response.status_code >= 500: return False, (time.perf_counter() - start) * 1000, f"HTTP {response.status_code}" return True, (time.perf_counter() - start) * 1000, None except requests.RequestException as exc: return False, (time.perf_counter() - start) * 1000, str(exc) def test_apt(mirror: Mirror) -> TestResult: if not mirror.url: return _record(mirror, "apt_release", False, None, "URL تعریف نشده") codename = validate_codename(get_ubuntu_codename()) release_url = f"{mirror.url.rstrip('/')}/dists/{codename}/Release" ok, latency, error = test_http(release_url) return _record(mirror, "apt_release", ok, latency, error) def test_docker(mirror: Mirror) -> TestResult: if not mirror.url: return _record(mirror, "docker_registry", False, None, "URL تعریف نشده") api_url = f"{mirror.url.rstrip('/')}/v2/" ok, latency, error = test_http(api_url) return _record(mirror, "docker_registry", ok, latency, error) def test_github(mirror: Mirror) -> TestResult: meta = mirror.get_meta() prefix = meta.get("instead_prefix") or mirror.url if not prefix: return _record(mirror, "github_ls_remote", False, None, "prefix تعریف نشده") test_url = f"{prefix.rstrip('/')}/octocat/Hello-World.git" start = time.perf_counter() result = run_on_host(["git", "ls-remote", test_url, "HEAD"], timeout=45) latency = (time.perf_counter() - start) * 1000 if result.returncode == 0: return _record(mirror, "github_ls_remote", True, latency, None) return _record(mirror, "github_ls_remote", False, latency, (result.stderr or result.stdout)[:300]) def test_pip(mirror: Mirror) -> TestResult: if not mirror.url: return _record(mirror, "pip_index", False, None, "URL تعریف نشده") ok, latency, error = test_http(mirror.url) return _record(mirror, "pip_index", ok, latency, error) def test_npm(mirror: Mirror) -> TestResult: if not mirror.url: return _record(mirror, "npm_registry", False, None, "URL تعریف نشده") ok, latency, error = test_http(mirror.url) return _record(mirror, "npm_registry", ok, latency, error) def test_mirror(mirror: Mirror) -> list[TestResult]: testers = { "dns": [test_dns], "apt": [test_apt], "docker": [test_docker], "github": [test_github], "pip": [test_pip], "npm": [test_npm], } results = [] for fn in testers.get(mirror.category, []): results.append(fn(mirror)) return results def test_all_enabled(category: str | None = None) -> list[TestResult]: query = Mirror.query.filter_by(enabled=True) if category: query = query.filter_by(category=category) all_results: list[TestResult] = [] for mirror in query.order_by(Mirror.priority).all(): all_results.extend(test_mirror(mirror)) return all_results def get_best_mirror(category: str) -> Mirror | None: mirrors = Mirror.query.filter_by(category=category, enabled=True).order_by(Mirror.priority).all() best: Mirror | None = None best_latency = float("inf") for mirror in mirrors: results = test_mirror(mirror) if not results: continue if all(r.success for r in results): avg_latency = sum(r.latency_ms or 9999 for r in results) / len(results) if avg_latency < best_latency: best_latency = avg_latency best = mirror return best