mirror/app/services/tester.py

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