websocket-tunnel 0.1.0__py3-none-any.whl

websocket_tunnel/__init__.py

@@ -0,0 +1,3 @@
+"""websocket-tunnel: a frp-like intranet penetration tool over WebSocket."""
+
+__version__ = "0.1.0"

websocket_tunnel/cli.py

@@ -0,0 +1,81 @@
+"""Command line interface for wtunnel."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import sys
+from pathlib import Path
+
+from . import __version__
+from .client import TunnelClient
+from .config import ConfigError, load_client_config, load_server_config
+from .server import TunnelServer
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="wtunnel",
+        description="WebSocket-based intranet penetration tool (frp-like).",
+    )
+    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    server = subparsers.add_parser("server", help="run the tunnel server")
+    server.add_argument("-c", "--config", type=Path, required=True, help="path to server TOML config")
+    server.add_argument("--listen", metavar="HOST:PORT", help="override listen address")
+    server.add_argument("--token", help="override shared auth token")
+    server.add_argument("--tls-cert", metavar="PATH", help="override TLS certificate path")
+    server.add_argument("--tls-key", metavar="PATH", help="override TLS private key path")
+    server.add_argument("-v", "--verbose", action="count", default=0, help="increase log verbosity (repeatable)")
+
+    client = subparsers.add_parser("client", help="run the tunnel client")
+    client.add_argument("-c", "--config", type=Path, required=True, help="path to client TOML config")
+    client.add_argument("--server", metavar="HOST:PORT", help="override tunnel server address")
+    client.add_argument("--token", help="override shared auth token")
+    client.add_argument("--tls", action="store_true", default=None, help="override: use wss transport")
+    client.add_argument(
+        "--tls-skip-verify",
+        action="store_true",
+        default=None,
+        help="override: skip TLS certificate verification",
+    )
+    client.add_argument("-v", "--verbose", action="count", default=0, help="increase log verbosity (repeatable)")
+    return parser
+
+
+def _setup_logging(verbosity: int) -> None:
+    level = logging.DEBUG if verbosity else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s %(levelname)-7s %(name)s: %(message)s",
+    )
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _build_parser().parse_args(argv)
+    _setup_logging(args.verbose)
+    try:
+        if args.command == "server":
+            config = load_server_config(
+                args.config,
+                listen=args.listen,
+                token=args.token,
+                tls_cert=args.tls_cert,
+                tls_key=args.tls_key,
+            )
+            return asyncio.run(TunnelServer(config).run())
+        config = load_client_config(
+            args.config,
+            server=args.server,
+            token=args.token,
+            tls=args.tls,
+            tls_skip_verify=args.tls_skip_verify,
+        )
+        return asyncio.run(TunnelClient(config).run())
+    except ConfigError as exc:
+        print(f"wtunnel: configuration error: {exc}", file=sys.stderr)
+        return 2
+    except KeyboardInterrupt:
+        return 0

websocket_tunnel/client.py

@@ -0,0 +1,115 @@
+"""Tunnel client: connects to a tunnel server and reconnects on failure."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import ssl
+
+import websockets
+from websockets.asyncio.client import connect
+
+from .config import ClientConfig, split_host_port
+from .protocol import MAX_CONTROL_SIZE
+from .session import Session
+
+CONNECT_TIMEOUT = 15.0
+RECONNECT_BASE = 1.0
+RECONNECT_MAX = 30.0
+RECONNECT_AFTER_SESSION = 1.0
+
+
+def _ws_url(host: str, port: int, secure: bool) -> str:
+    if ":" in host and not host.startswith("["):
+        host = f"[{host}]"
+    scheme = "wss" if secure else "ws"
+    return f"{scheme}://{host}:{port}/"
+
+
+class TunnelClient:
+    def __init__(self, config: ClientConfig) -> None:
+        self._config = config
+        self._log = logging.getLogger("wtunnel.client")
+        self._session: Session | None = None
+        self._stopping = False
+        self._stop_wait = asyncio.Event()
+        self.ready_event = asyncio.Event()
+
+    async def run(self) -> None:
+        backoff = RECONNECT_BASE
+        while not self._stopping:
+            try:
+                ws = await asyncio.wait_for(self._connect(), CONNECT_TIMEOUT)
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                self._log.warning(
+                    "connect to %s failed: %s; retrying in %.1fs", self._config.server, exc, backoff
+                )
+                await self._sleep_until_stop_or(backoff)
+                backoff = min(backoff * 2, RECONNECT_MAX)
+                continue
+            backoff = RECONNECT_BASE
+            self.ready_event = asyncio.Event()
+            session = Session(
+                ws,
+                role="client",
+                own_name="client",
+                own_proxies=self._config.proxies,
+                token=self._config.token,
+                ready_event=self.ready_event,
+                logger=self._log,
+            )
+            self._session = session
+            try:
+                await session.run()
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                self._log.warning("session failed: %s", exc)
+            finally:
+                self._session = None
+                try:
+                    await ws.close()
+                except Exception:
+                    pass
+            if self._stopping:
+                break
+            self._log.info("session ended; reconnecting in %.1fs", RECONNECT_AFTER_SESSION)
+            await self._sleep_until_stop_or(RECONNECT_AFTER_SESSION)
+
+    async def _connect(self) -> websockets.asyncio.client.ClientConnection:
+        ssl_context = self._build_ssl()
+        host, port = split_host_port(self._config.server)
+        uri = _ws_url(host, port, ssl_context is not None)
+        return await connect(
+            uri,
+            ssl=ssl_context,
+            max_size=MAX_CONTROL_SIZE,
+            ping_interval=20,
+            ping_timeout=20,
+            close_timeout=1,
+            compression=None,
+        )
+
+    def _build_ssl(self) -> ssl.SSLContext | None:
+        if not self._config.tls:
+            return None
+        context = ssl.create_default_context()
+        if self._config.tls_skip_verify:
+            context.check_hostname = False
+            context.verify_mode = ssl.CERT_NONE
+        return context
+
+    async def _sleep_until_stop_or(self, seconds: float) -> None:
+        try:
+            await asyncio.wait_for(self._stop_wait.wait(), timeout=seconds)
+        except asyncio.TimeoutError:
+            pass
+
+    async def stop(self) -> None:
+        self._stopping = True
+        self._stop_wait.set()
+        session = self._session
+        if session is not None:
+            await session.close()

websocket_tunnel/config.py

@@ -0,0 +1,173 @@
+"""TOML configuration parsing and validation for both tunnel nodes."""
+
+from __future__ import annotations
+
+import re
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+_HOST_PORT_RE = re.compile(r"^(\[[^\[\]]+\]|[^:]+):([0-9]+)$")
+
+
+class ConfigError(Exception):
+    """Raised when a configuration file or value is invalid."""
+
+
+@dataclass(frozen=True)
+class ProxyConfig:
+    name: str
+    listen: str
+    listen_side: str  # "local" | "peer"
+    backend: str
+    backend_side: str  # "local" | "peer"
+
+
+@dataclass(frozen=True)
+class ServerConfig:
+    listen: str = "0.0.0.0:7000"
+    token: str | None = None
+    tls_cert: str | None = None
+    tls_key: str | None = None
+    proxies: tuple[ProxyConfig, ...] = ()
+
+
+@dataclass(frozen=True)
+class ClientConfig:
+    server: str = "127.0.0.1:7000"
+    token: str | None = None
+    tls: bool = False
+    tls_skip_verify: bool = False
+    proxies: tuple[ProxyConfig, ...] = ()
+
+
+def split_host_port(value: str) -> tuple[str, int]:
+    """Split ``host:port`` (IPv6 literals must use ``[::1]:port`` form)."""
+    if not isinstance(value, str):
+        raise ConfigError(f"expected 'host:port', got {value!r}")
+    match = _HOST_PORT_RE.match(value.strip())
+    if match is None:
+        raise ConfigError(f"invalid 'host:port' address: {value!r}")
+    host = match.group(1)
+    if host.startswith("[") and host.endswith("]"):
+        host = host[1:-1]
+    port = int(match.group(2))
+    if not 1 <= port <= 65535:
+        raise ConfigError(f"port out of range in {value!r}")
+    if not host:
+        raise ConfigError(f"empty host in {value!r}")
+    return host, port
+
+
+def _check_side(value: Any, key: str) -> str:
+    if value not in ("local", "peer"):
+        raise ConfigError(f"{key} must be 'local' or 'peer', got {value!r}")
+    return value
+
+
+def proxy_from_dict(raw: Any) -> ProxyConfig:
+    """Build and validate a proxy from a mapping (config table or wire payload)."""
+    if not isinstance(raw, dict):
+        raise ConfigError("proxy entry must be a table/object")
+    missing = [key for key in ("name", "listen", "listen_side", "backend", "backend_side") if key not in raw]
+    if missing:
+        raise ConfigError(f"proxy entry missing keys: {', '.join(missing)}")
+    name = raw["name"]
+    if not isinstance(name, str) or not name.strip():
+        raise ConfigError("proxy name must be a non-empty string")
+    listen = raw["listen"]
+    backend = raw["backend"]
+    split_host_port(listen)
+    split_host_port(backend)
+    return ProxyConfig(
+        name=name.strip(),
+        listen=str(listen),
+        listen_side=_check_side(raw["listen_side"], "listen_side"),
+        backend=str(backend),
+        backend_side=_check_side(raw["backend_side"], "backend_side"),
+    )
+
+
+def _parse_proxies(raw: Any) -> tuple[ProxyConfig, ...]:
+    if raw is None:
+        return ()
+    if not isinstance(raw, list):
+        raise ConfigError("'proxies' must be a list of tables")
+    proxies = tuple(proxy_from_dict(item) for item in raw)
+    seen: set[str] = set()
+    for proxy in proxies:
+        if proxy.name in seen:
+            raise ConfigError(f"duplicate proxy name: {proxy.name!r}")
+        seen.add(proxy.name)
+    return proxies
+
+
+def _load_toml(path: Path) -> dict[str, Any]:
+    path = Path(path)
+    try:
+        data = tomllib.loads(path.read_text(encoding="utf-8"))
+    except FileNotFoundError as exc:
+        raise ConfigError(f"config file not found: {path}") from exc
+    except (OSError, tomllib.TOMLDecodeError) as exc:
+        raise ConfigError(f"failed to parse config {path}: {exc}") from exc
+    if not isinstance(data, dict):
+        raise ConfigError(f"config {path} must be a TOML table")
+    return data
+
+
+def load_server_config(
+    path: Path,
+    *,
+    listen: str | None = None,
+    token: str | None = None,
+    tls_cert: str | None = None,
+    tls_key: str | None = None,
+) -> ServerConfig:
+    data = _load_toml(path)
+    listen = listen if listen is not None else data.get("listen", "0.0.0.0:7000")
+    split_host_port(listen)
+    token = token if token is not None else data.get("token")
+    if token is not None and not isinstance(token, str):
+        raise ConfigError("'token' must be a string")
+    tls_raw = data.get("tls")
+    if tls_raw is None:
+        tls_raw = {}
+    if not isinstance(tls_raw, dict):
+        raise ConfigError("'tls' must be a table with 'cert' and 'key'")
+    tls_cert = tls_cert if tls_cert is not None else tls_raw.get("cert")
+    tls_key = tls_key if tls_key is not None else tls_raw.get("key")
+    if bool(tls_cert) != bool(tls_key):
+        raise ConfigError("'tls.cert' and 'tls.key' must be set together")
+    return ServerConfig(
+        listen=listen,
+        token=token,
+        tls_cert=tls_cert,
+        tls_key=tls_key,
+        proxies=_parse_proxies(data.get("proxies")),
+    )
+
+
+def load_client_config(
+    path: Path,
+    *,
+    server: str | None = None,
+    token: str | None = None,
+    tls: bool | None = None,
+    tls_skip_verify: bool | None = None,
+) -> ClientConfig:
+    data = _load_toml(path)
+    server = server if server is not None else data.get("server", "127.0.0.1:7000")
+    split_host_port(server)
+    token = token if token is not None else data.get("token")
+    if token is not None and not isinstance(token, str):
+        raise ConfigError("'token' must be a string")
+    tls = tls if tls is not None else bool(data.get("tls", False))
+    tls_skip_verify = tls_skip_verify if tls_skip_verify is not None else bool(data.get("tls_skip_verify", False))
+    return ClientConfig(
+        server=server,
+        token=token,
+        tls=tls,
+        tls_skip_verify=tls_skip_verify,
+        proxies=_parse_proxies(data.get("proxies")),
+    )

websocket_tunnel/protocol.py

@@ -0,0 +1,78 @@
+"""Binary wire protocol shared by tunnel nodes.
+
+Every WebSocket message is one binary frame:
+
+* control messages: 1 type byte + UTF-8 JSON object
+* stream data:      1 type byte + 4-byte big-endian stream id + raw chunk
+* ping/pong:        1 type byte
+"""
+
+from __future__ import annotations
+
+import json
+from enum import IntEnum
+from typing import Any
+
+
+class MessageType(IntEnum):
+    HELLO = 1
+    HELLO_OK = 2
+    HELLO_ERROR = 3
+    PROXY_REGISTER = 4
+    PROXY_OK = 5
+    PROXY_ERROR = 6
+    PROXY_UNREGISTER = 7
+    STREAM_OPEN = 8
+    STREAM_OK = 9
+    STREAM_ERROR = 10
+    STREAM_DATA = 11
+    STREAM_CLOSE = 12
+    PING = 13
+    PONG = 14
+
+
+PROTOCOL_VERSION = 1
+DATA_CHUNK_SIZE = 32 * 1024
+MAX_CONTROL_SIZE = 1 << 20
+STREAM_ID_SIZE = 4
+
+
+class InvalidFrameError(ValueError):
+    """Raised when a frame cannot be decoded."""
+
+
+def encode_control(message_type: MessageType, payload: dict[str, Any]) -> bytes:
+    return bytes([message_type]) + json.dumps(payload, separators=(",", ":")).encode("utf-8")
+
+
+def encode_data(stream_id: int, chunk: bytes) -> bytes:
+    return bytes([MessageType.STREAM_DATA]) + stream_id.to_bytes(STREAM_ID_SIZE, "big") + chunk
+
+
+def encode_ping() -> bytes:
+    return bytes([MessageType.PING])
+
+
+def encode_pong() -> bytes:
+    return bytes([MessageType.PONG])
+
+
+def decode(frame: bytes) -> tuple[MessageType, Any]:
+    if not frame:
+        raise InvalidFrameError("empty frame")
+    try:
+        message_type = MessageType(frame[0])
+    except ValueError as exc:
+        raise InvalidFrameError(f"unknown message type {frame[0]}") from exc
+    payload = frame[1:]
+    if message_type is MessageType.STREAM_DATA:
+        if len(payload) < STREAM_ID_SIZE:
+            raise InvalidFrameError("stream data frame too short")
+        stream_id = int.from_bytes(payload[:STREAM_ID_SIZE], "big")
+        return message_type, (stream_id, payload[STREAM_ID_SIZE:])
+    if message_type in (MessageType.PING, MessageType.PONG):
+        return message_type, None
+    try:
+        return message_type, json.loads(payload.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise InvalidFrameError(f"invalid JSON control payload: {exc}") from exc

websocket_tunnel/server.py

@@ -0,0 +1,81 @@
+"""Tunnel server: accepts control connections from tunnel clients."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import ssl
+
+import websockets
+from websockets.asyncio.server import ServerConnection
+
+from . import __version__
+from .config import ConfigError, ServerConfig, split_host_port
+from .protocol import MAX_CONTROL_SIZE
+from .session import Session
+
+
+class TunnelServer:
+    def __init__(self, config: ServerConfig) -> None:
+        self._config = config
+        self._log = logging.getLogger("wtunnel.server")
+        self._sessions: set[Session] = set()
+        self._stop_event = asyncio.Event()
+        self.ready_event = asyncio.Event()
+
+    async def run(self) -> None:
+        ssl_context = self._build_ssl()
+        host, port = split_host_port(self._config.listen)
+        async with websockets.serve(
+            self._handle_connection,
+            host,
+            port,
+            ssl=ssl_context,
+            max_size=MAX_CONTROL_SIZE,
+            ping_interval=20,
+            ping_timeout=20,
+            close_timeout=1,
+            compression=None,
+        ) as _:
+            scheme = "wss" if ssl_context else "ws"
+            self._log.info("wtunnel %s listening on %s:%s (%s)", __version__, host, port, scheme)
+            try:
+                await self._stop_event.wait()
+            finally:
+                await self._close_sessions()
+
+    def _build_ssl(self) -> ssl.SSLContext | None:
+        if not self._config.tls_cert:
+            return None
+        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        try:
+            context.load_cert_chain(self._config.tls_cert, self._config.tls_key)
+        except (OSError, ssl.SSLError) as exc:
+            raise ConfigError(f"failed to load TLS certificate: {exc}") from exc
+        return context
+
+    async def _handle_connection(self, ws: ServerConnection) -> None:
+        session = Session(
+            ws,
+            role="server",
+            own_name="server",
+            own_proxies=self._config.proxies,
+            token=self._config.token,
+            ready_event=self.ready_event,
+            logger=self._log,
+        )
+        self._sessions.add(session)
+        try:
+            await session.run()
+        finally:
+            self._sessions.discard(session)
+
+    async def _close_sessions(self) -> None:
+        sessions = list(self._sessions)
+        for session in sessions:
+            await session.close()
+        self._sessions.clear()
+
+    async def stop(self) -> None:
+        self._stop_event.set()
+        await self._close_sessions()