Sub2Socks/main.py

#!/usr/bin/env python3
"""Main orchestrator for Sub2SOCKS."""

import json
import logging
import os
import random
import sys
from typing import List, Optional

import fetcher
import config_builder

logger = logging.getLogger(__name__)


def setup_logging(level: str = "info") -> None:
    """Configure root logger with proper formatting.

    Args:
        level: Log level string (debug, info, warn, error).
    """
    numeric_level = getattr(logging, level.upper(), logging.INFO)
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(logging.Formatter(
        "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    ))
    root = logging.getLogger()
    root.setLevel(numeric_level)
    root.addHandler(handler)


def pick_node(vless_urls: List[str], strategy: str = "first") -> Optional[str]:
    """Select a VLESS URL from the list based on the given strategy.

    Args:
        vless_urls: List of VLESS URL strings.
        strategy: 'first' or 'random'. Defaults to 'first'.

    Returns:
        Selected VLESS URL string, or None if list is empty.
    """
    if not vless_urls:
        return None
    if strategy == "random":
        chosen = random.choice(vless_urls)
    else:
        chosen = vless_urls[0]
    logger.debug("Chosen node URL (truncated): %s", chosen[:60] + "...")
    return chosen


def run(
    sub_url: str,
    socks_port: int = 1080,
    socks_user: str = "",
    socks_pass: str = "",
    dns_server: str = "udp://1.1.1.1",
    pick_strategy: str = "first",
    log_level: str = "info",
    config_file_path: str = "/tmp/.sub2socks.json",
) -> None:
    """Execute the full Sub2SOCKS startup sequence and exec into sing-box.

    Args:
        sub_url: Subscription URL containing VLESS nodes.
        socks_port: Local SOCKS5 listen port.
        socks_user: Optional auth username.
        socks_pass: Optional auth password.
        dns_server: Upstream DNS resolver address.
        pick_strategy: Node selection strategy ('first' or 'random').
        log_level: Logging verbosity level.
        config_file_path: Path where sing-box JSON config is written.
    """
    setup_logging(log_level)
    print_banner(sub_url, socks_port, dns_server, pick_strategy, log_level)

    logger.info("Fetching subscription from %s", sub_url)
    try:
        sub_data = fetcher.fetch_subscription(sub_url)
    except Exception as exc:
        logger.error("Failed to reach %s: %s", sub_url, exc)
        sys.exit(1)

    vless_lines = fetcher.extract_vless_lines(sub_data)
    if not vless_lines:
        logger.error("No VLESS nodes found in subscription!")
        sys.stdout.write(sub_data[:500] + "\n")
        sys.exit(1)

    logger.info("Found %d VLESS node(s), picking one with strategy '%s'", len(vless_lines), pick_strategy)

    chosen = pick_node(vless_lines, pick_strategy)
    if not chosen:
        logger.error("Node selection failed.")
        sys.exit(1)

    proxy = fetcher.url_to_proxy_info(chosen)

    logger.info("Selected node: %s:%d", proxy.server, proxy.port)

    singbox_config = config_builder.build_singbox_config(
        proxy=proxy,
        socks_port=socks_port,
        socks_user=socks_user,
        socks_pass=socks_pass,
        dns_server=dns_server,
        log_level=log_level,
    )

    with open(config_file_path, "w") as fp:
        json.dump(singbox_config, fp, indent=2)

    logger.info("Config written to %s", config_file_path)

    os.execve(
        "/usr/local/bin/sing-box",
        ["sing-box", "run", "-c", config_file_path],
        dict(os.environ),
    )


def print_banner(
    sub_url: str,
    socks_port: int,
    dns_server: str,
    pick_strategy: str,
    log_level: str,
) -> None:
    """Print the ASCII startup banner to stdout.

    Args:
        sub_url: Subscription URL being used.
        socks_port: SOCKS listen port.
        dns_server: DNS upstream address.
        pick_strategy: Node pick strategy in use.
        log_level: Active log level.
    """
    banner = (
        "\n"
        "+-------------------------------------------------------+\n"
        "|   Sub2SOCKS -- bootstrapping                          |\n"
        "+-------------------------------------------------------+\n"
        f"|  SUB_URL       : {sub_url}\n"
        f"|  SOCKS port    : {socks_port:>6d}\n"
        f"|  DNS upstream  : {dns_server}\n"
        f"|  PICK_STRATEGY : {pick_strategy}\n"
        f"|  LOG_LEVEL     : {log_level}\n"
        "+-------------------------------------------------------+\n"
    )
    sys.stdout.write(banner)


def main() -> None:
    """Entry-point that reads env vars and calls run()."""
    run(
        sub_url=os.environ["SUB_URL"],
        socks_port=int(os.environ.get("SOCKS_PORT", "1080")),
        socks_user=os.environ.get("SOCKS_USER", ""),
        socks_pass=os.environ.get("SOCKS_PASS", ""),
        dns_server=os.environ.get("DNS_SERVER", "udp://1.1.1.1"),
        pick_strategy=os.environ.get("PICK_STRATEGY", "first"),
        log_level=os.environ.get("LOG_LEVEL", "info"),
    )


if __name__ == "__main__":
    main()