Sub2Socks/fetcher.py

#!/usr/bin/env python3
"""Subscription fetching and VLESS URL parsing."""

import base64
import logging
import time
import urllib.error
import urllib.parse as up
import urllib.request
from typing import Dict, List, Optional

from models import ProxyInfo

logger = logging.getLogger(__name__)


def fetch_subscription(url: str, max_retries: int = 3, backoff_base: float = 1.0) -> str:
    """Download subscription content with retry and exponential backoff.

    Args:
        url: Subscription URL to download.
        max_retries: Maximum number of retry attempts (default 3).
        backoff_base: Base delay in seconds for exponential backoff (default 1.0s -> 1, 2, 4).

    Returns:
        Decoded subscription text content.

    Raises:
        urllib.error.URLError: If all retries fail.
        Exception: On non-recoverable errors.
    """
    last_error: Optional[Exception] = None

    for attempt in range(max_retries):
        try:
            logger.info("Fetching subscription from %s (attempt %d/%d)", url, attempt + 1, max_retries)
            req = urllib.request.Request(url, headers={"User-Agent": "clash.meta"})
            with urllib.request.urlopen(req, timeout=30) as resp:
                raw = resp.read()
            try:
                content = raw.decode("utf-8")
            except UnicodeDecodeError:
                logger.debug("Raw data is base64 encoded, decoding...")
                content = base64.b64decode(raw).decode("utf-8")

            if attempt > 0:
                logger.info("Fetch succeeded on attempt %d", attempt + 1)
            return content

        except urllib.error.URLError as e:
            last_error = e
            wait_time = backoff_base * (2 ** attempt)
            logger.warning("Fetch failed (attempt %d/%d): %s. Retrying in %.1fs...",
                           attempt + 1, max_retries, e, wait_time)
            if attempt < max_retries - 1:
                time.sleep(wait_time)

        except Exception as e:
            logger.error("Non-recoverable fetch error: %s", e)
            raise

    logger.error("All %d fetch attempts failed for %s", max_retries, url)
    raise last_error  # type: ignore[misc]


def extract_vless_lines(text: str) -> List[str]:
    """Extract VLESS URI lines from subscription text.

    Handles both plain and base64-encoded line formats.

    Args:
        text: Raw subscription content.

    Returns:
        List of VLESS URL strings.
    """
    result: List[str] = []
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        if stripped.lower().startswith("vless://"):
            result.append(stripped)
            continue
        try:
            decoded = base64.b64decode(stripped).decode("utf-8")
            if decoded.lower().startswith("vless://"):
                result.append(decoded)
        except Exception:
            pass
    return result


def parse_vless_url(url_str: str) -> Dict[str, str]:
    """Parse a VLESS URL string into a component dictionary.

    Args:
        url_str: VLESS subscription URI (e.g. vless://uuid@host:port?params).

    Returns:
        Dictionary with all parsed proxy fields.
    """
    parsed = up.urlparse(url_str.strip())
    query_params = up.parse_qs(parsed.query)

    def get_param(key: str, default: Optional[str] = None) -> Optional[str]:
        if key in query_params:
            return query_params[key][0]
        return default

    return {
        "uuid":          parsed.username or "",
        "server":        parsed.hostname or "",
        "port":          parsed.port or (443 if parsed.scheme == "vless+tls" else 80),
        "security":      get_param("security", "tls"),
        "type":          get_param("type", "tcp"),
        "sni":           get_param("sni") or "",
        "alpn":          get_param("alpn") or "",
        "pbk":           get_param("public_key") or get_param("pbk") or "",
        "sid":           get_param("sid") or "",
        "spx":           get_param("spx") or "",
        "path":          get_param("path") or "/",
        "serviceName":   get_param("serviceName") or get_param("service-name") or "",
        "host":          get_param("host") or "",
        "flow":          get_param("flow") or "",
    }


def url_to_proxy_info(url_str: str) -> ProxyInfo:
    """Convert a VLESS URL directly to a ProxyInfo dataclass.

    Args:
        url_str: VLESS subscription URI.

    Returns:
        ProxyInfo instance with parsed fields.
    """
    raw = parse_vless_url(url_str)
    return ProxyInfo(
        uuid=raw["uuid"],
        server=raw["server"],
        port=int(raw["port"]),
        security=raw["security"],
        type=raw["type"],
        sni=raw["sni"],
        alpn=raw["alpn"],
        pbk=raw["pbk"],
        sid=raw["sid"],
        spx=raw["spx"],
        path=raw["path"],
        serviceName=raw["serviceName"],
        host=raw["host"],
        flow=raw["flow"],
    )