linksocks 3.0.13__cp310-cp310-macosx_14_0_universal2.whl

linksocks/_server.py

@@ -0,0 +1,355 @@
+"""
+WebSocket SOCKS5 proxy server implementation.
+
+This module provides the Server class for running a SOCKS5 proxy server
+that communicates with clients over WebSocket connections.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Optional
+
+# Underlying Go bindings module (generated)
+from linksockslib import linksocks  # type: ignore
+
+from ._base import (
+    _SnakePassthrough,
+    _to_duration,
+    _logger,
+    BufferZerologLogger,
+    ReverseTokenResult,
+    DurationLike,
+)
+
+
+class Server(_SnakePassthrough):
+    """WebSocket SOCKS5 proxy server.
+    
+    The Server class manages WebSocket connections from clients and provides
+    SOCKS5 proxy functionality. It supports both forward and reverse proxy modes.
+    
+    In forward mode, clients connect to the server and the server makes outbound
+    connections on behalf of the clients.
+    
+    In reverse mode, the server runs a local SOCKS5 server and forwards connections
+    through WebSocket to clients, which then connect to targets directly.
+    """
+
+    def __init__(
+        self,
+        *,
+        logger: Optional[logging.Logger] = None,
+        ws_host: Optional[str] = None,
+        ws_port: Optional[int] = None,
+        socks_host: Optional[str] = None,
+        port_pool: Optional[Any] = None,
+        socks_wait_client: Optional[bool] = None,
+        buffer_size: Optional[int] = None,
+        api_key: Optional[str] = None,
+        channel_timeout: Optional[DurationLike] = None,
+        connect_timeout: Optional[DurationLike] = None,
+        fast_open: Optional[bool] = None,
+        upstream_proxy: Optional[str] = None,
+        upstream_username: Optional[str] = None,
+        upstream_password: Optional[str] = None,
+    ) -> None:
+        """Initialize the WebSocket SOCKS5 proxy server.
+        
+        Args:
+            logger: Python logger instance for this server
+            ws_host: WebSocket server listen address
+            ws_port: WebSocket server listen port
+            socks_host: SOCKS5 server listen address (for reverse mode)
+            port_pool: Pool of ports for SOCKS5 servers
+            socks_wait_client: Whether to wait for client connections before starting SOCKS5
+            buffer_size: Buffer size for data transfer
+            api_key: API key for HTTP management interface
+            channel_timeout: Timeout for WebSocket channels
+            connect_timeout: Timeout for outbound connections
+            fast_open: Assume connection success and allow data transfer immediately
+            upstream_proxy: Upstream proxy address for chaining
+            upstream_username: Username for upstream proxy authentication
+            upstream_password: Password for upstream proxy authentication
+        """
+        opt = linksocks.DefaultServerOption()
+        if logger is None:
+            logger = _logger
+        # Use buffer-based logger system
+        self._managed_logger = BufferZerologLogger(logger, f"server_{id(self)}")
+        opt.WithLogger(self._managed_logger.go_logger)
+        if ws_host is not None:
+            opt.WithWSHost(ws_host)
+        if ws_port is not None:
+            opt.WithWSPort(ws_port)
+        if socks_host is not None:
+            opt.WithSocksHost(socks_host)
+        if port_pool is not None:
+            opt.WithPortPool(port_pool)
+        if socks_wait_client is not None:
+            opt.WithSocksWaitClient(bool(socks_wait_client))
+        if buffer_size is not None:
+            opt.WithBufferSize(int(buffer_size))
+        if api_key is not None:
+            opt.WithAPI(api_key)
+        if channel_timeout is not None:
+            opt.WithChannelTimeout(_to_duration(channel_timeout))
+        if connect_timeout is not None:
+            opt.WithConnectTimeout(_to_duration(connect_timeout))
+        if fast_open is not None:
+            opt.WithFastOpen(bool(fast_open))
+        if upstream_proxy is not None:
+            opt.WithUpstreamProxy(upstream_proxy)
+        if upstream_username or upstream_password:
+            opt.WithUpstreamAuth(upstream_username or "", upstream_password or "")
+
+        self._raw = linksocks.NewLinkSocksServer(opt)
+        self._ctx = None
+
+    @property
+    def log(self) -> logging.Logger:
+        """Access the Python logger for this server instance."""
+        return self._managed_logger.py_logger
+
+    def add_forward_token(self, token: Optional[str] = None) -> str:
+        """Add a forward proxy token.
+        
+        Args:
+            token: Token string, auto-generated if not provided
+            
+        Returns:
+            The token string (generated or provided)
+        """
+        return self._raw.AddForwardToken(token or "")
+
+    async def async_add_forward_token(self, token: Optional[str] = None) -> str:
+        """Add a forward proxy token asynchronously.
+        
+        Args:
+            token: Token string, auto-generated if not provided
+            
+        Returns:
+            The token string (generated or provided)
+        """
+        return await asyncio.to_thread(self._raw.AddForwardToken, token or "")
+
+    def add_reverse_token(
+        self,
+        *,
+        token: Optional[str] = None,
+        port: Optional[int] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        allow_manage_connector: Optional[bool] = None,
+    ) -> ReverseTokenResult:
+        """Add a reverse proxy token.
+        
+        Args:
+            token: Token string, auto-generated if not provided
+            port: SOCKS5 server port, auto-assigned if not provided
+            username: SOCKS5 authentication username
+            password: SOCKS5 authentication password
+            allow_manage_connector: Whether to allow clients to manage connector tokens
+            
+        Returns:
+            Result containing the token and assigned port
+        """
+        opts = linksocks.DefaultReverseTokenOptions()
+        if token:
+            opts.Token = token
+        if port is not None:
+            opts.Port = int(port)
+        if username is not None:
+            opts.Username = username
+        if password is not None:
+            opts.Password = password
+        if allow_manage_connector is not None:
+            opts.AllowManageConnector = bool(allow_manage_connector)
+        result = self._raw.AddReverseToken(opts)
+        return ReverseTokenResult(token=result.Token, port=result.Port)
+
+    async def async_add_reverse_token(
+        self,
+        *,
+        token: Optional[str] = None,
+        port: Optional[int] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        allow_manage_connector: Optional[bool] = None,
+    ) -> ReverseTokenResult:
+        """Add a reverse proxy token asynchronously.
+        
+        Args:
+            token: Token string, auto-generated if not provided
+            port: SOCKS5 server port, auto-assigned if not provided
+            username: SOCKS5 authentication username
+            password: SOCKS5 authentication password
+            allow_manage_connector: Whether to allow clients to manage connector tokens
+            
+        Returns:
+            Result containing the token and assigned port
+        """
+        opts = linksocks.DefaultReverseTokenOptions()
+        if token:
+            opts.Token = token
+        if port is not None:
+            opts.Port = int(port)
+        if username is not None:
+            opts.Username = username
+        if password is not None:
+            opts.Password = password
+        if allow_manage_connector is not None:
+            opts.AllowManageConnector = bool(allow_manage_connector)
+        result = await asyncio.to_thread(self._raw.AddReverseToken, opts)
+        return ReverseTokenResult(token=result.Token, port=result.Port)
+
+    def add_connector_token(self, connector_token: Optional[str], reverse_token: str) -> str:
+        """Add a connector token for reverse proxy.
+        
+        Args:
+            connector_token: Connector token string, auto-generated if not provided
+            reverse_token: Associated reverse proxy token
+            
+        Returns:
+            The connector token string (generated or provided)
+        """
+        return self._raw.AddConnectorToken(connector_token or "", reverse_token)
+
+    async def async_add_connector_token(self, connector_token: Optional[str], reverse_token: str) -> str:
+        """Add a connector token for reverse proxy asynchronously.
+        
+        Args:
+            connector_token: Connector token string, auto-generated if not provided
+            reverse_token: Associated reverse proxy token
+            
+        Returns:
+            The connector token string (generated or provided)
+        """
+        return await asyncio.to_thread(self._raw.AddConnectorToken, connector_token or "", reverse_token)
+
+    def remove_token(self, token: str) -> bool:
+        """Remove a token from the server.
+        
+        Args:
+            token: Token to remove
+            
+        Returns:
+            True if token was removed, False if not found
+        """
+        return self._raw.RemoveToken(token)
+
+    async def async_remove_token(self, token: str) -> bool:
+        """Remove a token from the server asynchronously.
+        
+        Args:
+            token: Token to remove
+            
+        Returns:
+            True if token was removed, False if not found
+        """
+        return await asyncio.to_thread(self._raw.RemoveToken, token)
+
+    def wait_ready(self, timeout: Optional[DurationLike] = None) -> None:
+        """Wait for the server to be ready.
+        
+        Args:
+            timeout: Maximum time to wait, no timeout if None
+        """
+        if not self._ctx:
+            self._ctx = linksocks.NewContextWithCancel()
+        timeout = _to_duration(timeout) if timeout is not None else 0
+        self._raw.WaitReady(ctx=self._ctx.Context(), timeout=timeout)
+
+    async def async_wait_ready(self, timeout: Optional[DurationLike] = None) -> None:
+        """Wait for the server to be ready asynchronously.
+        
+        Args:
+            timeout: Maximum time to wait, no timeout if None
+        """
+        if not self._ctx:
+            self._ctx = linksocks.NewContextWithCancel()
+        timeout = _to_duration(timeout) if timeout is not None else 0
+        try:
+            return await asyncio.to_thread(self._raw.WaitReady, ctx=self._ctx.Context(), timeout=timeout)
+        except asyncio.CancelledError:
+            # Ensure the underlying Go server stops when startup wait is cancelled
+            try:
+                try:
+                    self._ctx.Cancel()
+                except Exception:
+                    pass
+                await asyncio.shield(asyncio.to_thread(self._raw.Close))
+                if hasattr(self, '_managed_logger') and self._managed_logger:
+                    try:
+                        self._managed_logger.cleanup()
+                    except Exception:
+                        pass
+            finally:
+                raise
+
+    def close(self) -> None:
+        """Close the server and clean up resources."""
+        # Close server
+        if hasattr(self, '_raw') and self._raw:
+            self._raw.Close()
+        # Clean up managed logger
+        if hasattr(self, '_managed_logger') and self._managed_logger:
+            try:
+                self._managed_logger.cleanup()
+            except:
+                # Ignore cleanup errors
+                pass
+        # Close context
+        if hasattr(self, '_ctx') and self._ctx:
+            try:
+                self._ctx.Cancel()
+            except Exception:
+                # Ignore errors during context close
+                pass
+
+    async def async_close(self) -> None:
+        """Close the server and clean up resources asynchronously."""
+        # Close server
+        if hasattr(self, '_raw') and self._raw:
+            await asyncio.to_thread(self._raw.Close)
+        # Clean up managed logger
+        if hasattr(self, '_managed_logger') and self._managed_logger:
+            try:
+                self._managed_logger.cleanup()
+            except:
+                # Ignore cleanup errors
+                pass
+        # Close context
+        if hasattr(self, '_ctx') and self._ctx:
+            try:
+                self._ctx.Cancel()
+            except Exception:
+                # Ignore errors during context close
+                pass
+    
+    def __enter__(self) -> "Server":
+        """Context manager entry."""
+        self.wait_ready()
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        """Context manager exit."""
+        self.close()
+        
+    async def __aenter__(self) -> "Server":
+        """Async context manager entry."""
+        await self.async_wait_ready()
+        return self
+        
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        """Async context manager exit."""
+        await self.async_close()
+        
+    def __del__(self):
+        """Destructor - clean up resources."""
+        try:
+            self.close()
+        except Exception:
+            # Ignore errors during cleanup
+            pass

linksocks/_utils.py

@@ -0,0 +1,144 @@
+"""
+Utility functions for linksocks CLI.
+
+This module provides common utility functions used across the linksocks
+command-line interface, including SOCKS proxy URL parsing and environment
+variable handling.
+"""
+
+import os
+import urllib.parse
+from typing import Optional, Tuple
+
+
+def parse_socks_proxy(proxy_url: str) -> Tuple[str, Optional[str], Optional[str]]:
+    """
+    Parse a SOCKS5 proxy URL and extract connection details.
+    
+    This function parses SOCKS5 proxy URLs in the standard format and extracts
+    the host address, username, and password for authentication. It supports
+    URLs with optional authentication credentials and port numbers.
+    
+    Args:
+        proxy_url: URL string in format socks5://[user:pass@]host[:port]
+                  Examples:
+                  - "socks5://127.0.0.1:9870"
+                  - "socks5://user:pass@proxy.example.com:9870"
+                  - "socks5://proxy.example.com" (uses default port 9870)
+        
+    Returns:
+        A tuple containing:
+        - address (str): Host and port in "host:port" format
+        - username (str | None): Username for authentication, if provided
+        - password (str | None): Password for authentication, if provided
+        
+    Raises:
+        ValueError: If the URL is invalid, malformed, or uses an unsupported scheme
+        
+    Examples:
+        >>> parse_socks_proxy("socks5://127.0.0.1:9870")
+        ("127.0.0.1:9870", None, None)
+        
+        >>> parse_socks_proxy("socks5://user:pass@proxy.example.com:9870")
+        ("proxy.example.com:9870", "user", "pass")
+        
+        >>> parse_socks_proxy("socks5://proxy.example.com")
+        ("proxy.example.com:9870", None, None)
+    """
+    if not proxy_url:
+        return "", None, None
+
+    try:
+        parsed_url = urllib.parse.urlparse(proxy_url)
+    except Exception as e:
+        raise ValueError(f"Invalid proxy URL: {e}") from e
+
+    # Validate the URL scheme
+    if parsed_url.scheme != "socks5":
+        raise ValueError(f"Unsupported proxy scheme: {parsed_url.scheme}. Only 'socks5' is supported.")
+
+    # Extract authentication credentials if present
+    username = None
+    password = None
+    if parsed_url.username:
+        username = urllib.parse.unquote(parsed_url.username)
+    if parsed_url.password:
+        password = urllib.parse.unquote(parsed_url.password)
+
+    # Build the address with default port if not specified
+    host = parsed_url.hostname
+    if not host:
+        raise ValueError("Proxy URL must specify a hostname")
+        
+    port = parsed_url.port or 9870  # Default SOCKS5 port
+    address = f"{host}:{port}"
+
+    return address, username, password
+
+
+def get_env_or_flag(flag_value: Optional[str], env_var: str) -> Optional[str]:
+    """
+    Get a configuration value from command-line flag or environment variable.
+    
+    This utility function implements the common pattern of checking for a value
+    in a command-line flag first, and falling back to an environment variable
+    if the flag is not provided. This allows users to configure the application
+    either through command-line arguments or environment variables.
+    
+    Args:
+        flag_value: The value provided via command-line flag (may be None)
+        env_var: The name of the environment variable to check as fallback
+        
+    Returns:
+        The configuration value from the flag if provided, otherwise from
+        the environment variable, or None if neither is set
+        
+    Examples:
+        >>> os.environ['LINKSOCKS_TOKEN'] = 'env-token'
+        >>> get_env_or_flag(None, 'LINKSOCKS_TOKEN')
+        'env-token'
+        
+        >>> get_env_or_flag('flag-token', 'LINKSOCKS_TOKEN')
+        'flag-token'
+        
+        >>> get_env_or_flag(None, 'NONEXISTENT_VAR')
+        None
+        
+    Note:
+        Command-line flags always take precedence over environment variables.
+        This allows users to override environment settings on a per-invocation basis.
+    """
+    if flag_value:
+        return flag_value
+    return os.getenv(env_var)
+
+
+def validate_required_token(token: Optional[str], env_var: str = "LINKSOCKS_TOKEN") -> str:
+    """
+    Validate that a required authentication token is provided.
+    
+    This function checks that an authentication token is available either from
+    a command-line flag or environment variable, and raises a descriptive error
+    if neither is provided.
+    
+    Args:
+        token: The token value from command-line flag (may be None)
+        env_var: The environment variable name to check (default: "LINKSOCKS_TOKEN")
+        
+    Returns:
+        The validated token string
+        
+    Raises:
+        ValueError: If no token is provided via flag or environment variable
+        
+    Example:
+        >>> validate_required_token("my-token")
+        "my-token"
+        
+        >>> validate_required_token(None)  # Will check LINKSOCKS_TOKEN env var
+        ValueError: Token is required. Provide via --token or LINKSOCKS_TOKEN environment variable.
+    """
+    actual_token = get_env_or_flag(token, env_var)
+    if not actual_token:
+        raise ValueError(f"Token is required. Provide via --token or {env_var} environment variable.")
+    return actual_token