Nexus/server/infrastructure/subscription_parser.py

"""Parse proxy subscription URLs and extract node IP addresses.

Supports: SS (ShadowSocks), VMess, VLESS, Trojan, Hysteria2.
Used for building login IP allowlists from proxy service subscriptions.
"""

import base64
import json
import logging
import re
from typing import Optional
from urllib.parse import urlparse

logger = logging.getLogger("nexus.subscription_parser")

# Simple IPv4 pattern
_IPV4_RE = re.compile(r"^(\d{1,3}\.){3}\d{1,3}$")
# Simple IPv6 — we accept the host field as-is if not IPv4/hostname
_DOMAIN_RE = re.compile(r"^[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?(\.[a-zA-Z]{2,})+$")


def _decode_b64(s: str) -> str:
    """Best-effort Base64 decode with padding fix."""
    s = s.strip().rstrip("=")
    padding = 4 - len(s) % 4 if len(s) % 4 else 0
    return base64.b64decode(s + "=" * padding, altchars=b"-_").decode("utf-8", errors="ignore")


def _extract_host_from_url(raw: str) -> Optional[str]:
    """Parse a URL and return the hostname."""
    try:
        p = urlparse(raw)
        return p.hostname or None
    except Exception:
        return None


def _parse_ss(line: str) -> Optional[str]:
    """Parse ShadowSocks URI: ss://BASE64#tag  or  ss://BASE64@host:port#tag"""
    body = line[5:].split("#")[0]   # strip prefix and fragment
    if "@" in body:
        # New format: ss://METHOD:PWD@HOST:PORT
        return _extract_host_from_url(line)
    # Old format: ss://BASE64(method:pwd@host:port)
    try:
        decoded = _decode_b64(body)
        if "@" in decoded:
            host_part = decoded.split("@", 1)[1].split(":")[0]
            return host_part
    except Exception:  # noqa: S110 — non-critical URL parse fallback
        pass
    return None


def _parse_vmess(line: str) -> Optional[str]:
    """Parse VMess URI: vmess://BASE64(JSON) — JSON has 'add' key for host."""
    try:
        body = line[8:]
        json_str = _decode_b64(body)
        data = json.loads(json_str)
        return data.get("add") or data.get("host") or None
    except Exception:
        return None


def _parse_simple_url(line: str) -> Optional[str]:
    """Schemes where host is directly in URL: vless, trojan, hysteria2, hy2, tuic."""
    return _extract_host_from_url(line)


def _is_valid_host(host: Optional[str]) -> bool:
    if not host:
        return False
    host = host.strip()
    if not host or host in ("localhost", "127.0.0.1", "::1"):
        return False
    if _IPV4_RE.match(host):
        return True
    if _DOMAIN_RE.match(host):
        return True
    # IPv6 literal (bracketed already stripped by urlparse)
    if ":" in host:
        return True
    return False


def parse_subscription_content(raw_content: str) -> list[str]:
    """Given subscription raw text (or base64 of it), return unique host/IP list."""
    # Try to decode as base64 first (many subscriptions are base64-encoded)
    content = raw_content.strip()
    if "\n" not in content and len(content) > 50:
        try:
            decoded = _decode_b64(content)
            if any(decoded.startswith(p) for p in ("ss://", "vmess://", "vless://", "trojan://")):
                content = decoded
        except Exception:  # noqa: S110 — non-critical parse, treat as non-encoded
            pass

    hosts: list[str] = []
    seen: set[str] = set()

    for line in content.splitlines():
        line = line.strip()
        if not line:
            continue

        host: Optional[str] = None

        if line.startswith("ss://"):
            host = _parse_ss(line)
        elif line.startswith("vmess://"):
            host = _parse_vmess(line)
        elif line.startswith(("vless://", "trojan://", "hysteria2://", "hy2://", "tuic://")):
            host = _parse_simple_url(line)

        if host and _is_valid_host(host) and host not in seen:
            seen.add(host)
            hosts.append(host)

    logger.info("Subscription parsed: %d nodes → %d unique hosts", len(content.splitlines()), len(hosts))
    return hosts