Initial commit: castle personal software platform
Add three submodules (central-context, mboxer, notification-bridge), devbox-connect as tracked files, and top-level project docs.
This commit is contained in:
3
devbox-connect/src/devbox_connect/__init__.py
Normal file
3
devbox-connect/src/devbox_connect/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""SSH tunnel manager for connecting to devbox ports."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
235
devbox-connect/src/devbox_connect/cli.py
Normal file
235
devbox-connect/src/devbox_connect/cli.py
Normal file
@@ -0,0 +1,235 @@
|
||||
"""Command-line interface for devbox-connect."""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from .config import ConfigError, load_config
|
||||
from .tunnel import TunnelManager, TunnelState
|
||||
|
||||
# ANSI color codes
|
||||
GREEN = "\033[92m"
|
||||
RED = "\033[91m"
|
||||
YELLOW = "\033[93m"
|
||||
CYAN = "\033[96m"
|
||||
RESET = "\033[0m"
|
||||
BOLD = "\033[1m"
|
||||
|
||||
# State to color mapping
|
||||
STATE_COLORS = {
|
||||
TunnelState.CONNECTED: GREEN,
|
||||
TunnelState.CONNECTING: YELLOW,
|
||||
TunnelState.DISCONNECTED: RED,
|
||||
TunnelState.ERROR: RED,
|
||||
}
|
||||
|
||||
STATE_SYMBOLS = {
|
||||
TunnelState.CONNECTED: "[OK]",
|
||||
TunnelState.CONNECTING: "[..]",
|
||||
TunnelState.DISCONNECTED: "[--]",
|
||||
TunnelState.ERROR: "[!!]",
|
||||
}
|
||||
|
||||
|
||||
def setup_logging(verbose: bool) -> None:
|
||||
"""Configure logging."""
|
||||
level = logging.DEBUG if verbose else logging.INFO
|
||||
format_str = "%(asctime)s - %(levelname)s - %(message)s" if verbose else "%(message)s"
|
||||
logging.basicConfig(level=level, format=format_str)
|
||||
|
||||
|
||||
def print_banner() -> None:
|
||||
"""Print application banner."""
|
||||
print(f"{BOLD}{CYAN}devbox-connect{RESET} - SSH Tunnel Manager")
|
||||
print()
|
||||
|
||||
|
||||
def print_status(manager: TunnelManager) -> None:
|
||||
"""Print current tunnel status."""
|
||||
status = manager.get_status()
|
||||
|
||||
for host_key, tunnels in status.items():
|
||||
print(f"{BOLD}{host_key}{RESET}")
|
||||
for name, tunnel_status in tunnels.items():
|
||||
state = tunnel_status.state
|
||||
color = STATE_COLORS[state]
|
||||
symbol = STATE_SYMBOLS[state]
|
||||
local_port = tunnel_status.config.local_port or tunnel_status.config.remote_port
|
||||
remote = f"{tunnel_status.config.remote_host}:{tunnel_status.config.remote_port}"
|
||||
|
||||
line = f" {color}{symbol}{RESET} {name}: localhost:{local_port} -> {remote}"
|
||||
|
||||
if tunnel_status.connections > 0:
|
||||
line += f" ({tunnel_status.connections} active)"
|
||||
|
||||
if tunnel_status.error:
|
||||
line += f" {RED}({tunnel_status.error}){RESET}"
|
||||
|
||||
print(line)
|
||||
print()
|
||||
|
||||
|
||||
def on_status_change(host: str, tunnel: str, state: TunnelState) -> None:
|
||||
"""Callback for tunnel status changes."""
|
||||
color = STATE_COLORS[state]
|
||||
symbol = STATE_SYMBOLS[state]
|
||||
logging.info(f"{color}{symbol}{RESET} {host}/{tunnel}: {state.value}")
|
||||
|
||||
|
||||
def run_manager(config_path: Path, verbose: bool) -> int:
|
||||
"""Run the tunnel manager."""
|
||||
setup_logging(verbose)
|
||||
print_banner()
|
||||
|
||||
# Load configuration
|
||||
try:
|
||||
config = load_config(config_path)
|
||||
except ConfigError as e:
|
||||
logging.error(f"Configuration error: {e}")
|
||||
return 1
|
||||
|
||||
if not config.hosts:
|
||||
logging.error("No tunnels configured")
|
||||
return 1
|
||||
|
||||
# Count tunnels
|
||||
total_tunnels = sum(len(h.tunnels) for h in config.hosts)
|
||||
logging.info(f"Loaded {total_tunnels} tunnel(s) for {len(config.hosts)} host(s)")
|
||||
print()
|
||||
|
||||
# Create and start manager
|
||||
manager = TunnelManager(config, on_status_change=on_status_change)
|
||||
|
||||
# Handle shutdown signals
|
||||
shutdown_requested = False
|
||||
|
||||
def signal_handler(signum: int, frame: object) -> None:
|
||||
nonlocal shutdown_requested
|
||||
if shutdown_requested:
|
||||
# Force exit on second signal
|
||||
sys.exit(1)
|
||||
shutdown_requested = True
|
||||
print(f"\n{YELLOW}Shutting down...{RESET}")
|
||||
manager.stop()
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
|
||||
# Start tunnels
|
||||
manager.start()
|
||||
|
||||
# Wait a moment for connections
|
||||
time.sleep(2)
|
||||
|
||||
# Print initial status
|
||||
print_status(manager)
|
||||
print(f"{CYAN}Press Ctrl+C to stop{RESET}")
|
||||
print()
|
||||
|
||||
# Keep running until shutdown
|
||||
try:
|
||||
while not shutdown_requested:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
manager.stop()
|
||||
print(f"{GREEN}Stopped{RESET}")
|
||||
return 0
|
||||
|
||||
|
||||
def show_status(config_path: Path) -> int:
|
||||
"""Show status of configured tunnels (quick check)."""
|
||||
setup_logging(verbose=False)
|
||||
|
||||
try:
|
||||
config = load_config(config_path)
|
||||
except ConfigError as e:
|
||||
logging.error(f"Configuration error: {e}")
|
||||
return 1
|
||||
|
||||
print_banner()
|
||||
print(f"Config: {config_path}")
|
||||
print()
|
||||
|
||||
for host in config.hosts:
|
||||
print(f"{BOLD}{host.user}@{host.host}:{host.port}{RESET}")
|
||||
for tunnel in host.tunnels:
|
||||
local_port = tunnel.local_port or tunnel.remote_port
|
||||
remote = f"{tunnel.remote_host}:{tunnel.remote_port}"
|
||||
print(f" - {tunnel.name}: localhost:{local_port} -> {remote}")
|
||||
print()
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def validate_config(config_path: Path) -> int:
|
||||
"""Validate configuration file."""
|
||||
setup_logging(verbose=False)
|
||||
|
||||
try:
|
||||
config = load_config(config_path)
|
||||
total_tunnels = sum(len(h.tunnels) for h in config.hosts)
|
||||
print(f"{GREEN}Configuration valid{RESET}")
|
||||
print(f" Hosts: {len(config.hosts)}")
|
||||
print(f" Tunnels: {total_tunnels}")
|
||||
return 0
|
||||
except ConfigError as e:
|
||||
print(f"{RED}Configuration invalid: {e}{RESET}")
|
||||
return 1
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Main entry point."""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="devbox-connect",
|
||||
description="SSH tunnel manager for connecting to devbox ports",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-c",
|
||||
"--config",
|
||||
type=Path,
|
||||
default=Path("tunnels.yaml"),
|
||||
help="Path to configuration file (default: tunnels.yaml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-v",
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
help="Enable verbose output",
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest="command", help="Commands")
|
||||
|
||||
# start command (default)
|
||||
start_parser = subparsers.add_parser("start", help="Start tunnels (default)")
|
||||
start_parser.add_argument("-v", "--verbose", action="store_true")
|
||||
|
||||
# status command
|
||||
subparsers.add_parser("status", help="Show configured tunnels")
|
||||
|
||||
# validate command
|
||||
subparsers.add_parser("validate", help="Validate configuration file")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Handle default command
|
||||
command = args.command or "start"
|
||||
|
||||
if command == "start":
|
||||
verbose = getattr(args, "verbose", False)
|
||||
return run_manager(args.config, verbose)
|
||||
elif command == "status":
|
||||
return show_status(args.config)
|
||||
elif command == "validate":
|
||||
return validate_config(args.config)
|
||||
else:
|
||||
parser.print_help()
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
162
devbox-connect/src/devbox_connect/config.py
Normal file
162
devbox-connect/src/devbox_connect/config.py
Normal file
@@ -0,0 +1,162 @@
|
||||
"""Configuration loading and validation."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
@dataclass
|
||||
class TunnelConfig:
|
||||
"""Configuration for a single SSH tunnel."""
|
||||
|
||||
name: str
|
||||
host: str
|
||||
remote_port: int
|
||||
local_port: int | None = None # Defaults to remote_port if not specified
|
||||
user: str | None = None # Defaults to global user
|
||||
key_file: str | None = None # Defaults to global key_file
|
||||
remote_host: str = "localhost" # The host on the remote side to connect to
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.local_port is None:
|
||||
self.local_port = self.remote_port
|
||||
|
||||
|
||||
@dataclass
|
||||
class HostConfig:
|
||||
"""Configuration for a host with multiple tunnels."""
|
||||
|
||||
host: str
|
||||
user: str
|
||||
key_file: str | None = None
|
||||
port: int = 22
|
||||
tunnels: list[TunnelConfig] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""Root configuration."""
|
||||
|
||||
hosts: list[HostConfig] = field(default_factory=list)
|
||||
reconnect_delay: int = 5 # Seconds between reconnection attempts
|
||||
max_reconnect_delay: int = 60 # Max delay with exponential backoff
|
||||
|
||||
|
||||
class ConfigError(Exception):
|
||||
"""Configuration error."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def load_config(path: Path) -> Config:
|
||||
"""Load configuration from a YAML file."""
|
||||
if not path.exists():
|
||||
raise ConfigError(f"Config file not found: {path}")
|
||||
|
||||
with open(path) as f:
|
||||
data = yaml.safe_load(f)
|
||||
|
||||
if not data:
|
||||
raise ConfigError("Empty configuration file")
|
||||
|
||||
return _parse_config(data)
|
||||
|
||||
|
||||
def _parse_config(data: dict) -> Config:
|
||||
"""Parse configuration dictionary into Config object."""
|
||||
config = Config(
|
||||
reconnect_delay=data.get("reconnect_delay", 5),
|
||||
max_reconnect_delay=data.get("max_reconnect_delay", 60),
|
||||
)
|
||||
|
||||
# Handle simple format: list of tunnels with host info per tunnel
|
||||
if "tunnels" in data:
|
||||
config.hosts = _parse_simple_format(data["tunnels"], data)
|
||||
# Handle grouped format: hosts with nested tunnels
|
||||
elif "hosts" in data:
|
||||
config.hosts = _parse_grouped_format(data["hosts"])
|
||||
else:
|
||||
raise ConfigError("Config must contain either 'tunnels' or 'hosts' key")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _parse_simple_format(tunnels: list[dict], global_config: dict) -> list[HostConfig]:
|
||||
"""Parse simple format where each tunnel specifies its host."""
|
||||
# Group tunnels by host
|
||||
hosts_map: dict[str, HostConfig] = {}
|
||||
|
||||
global_user = global_config.get("user")
|
||||
global_key_file = global_config.get("key_file")
|
||||
|
||||
for t in tunnels:
|
||||
if "host" not in t:
|
||||
raise ConfigError(f"Tunnel '{t.get('name', 'unnamed')}' missing 'host'")
|
||||
if "remote_port" not in t:
|
||||
raise ConfigError(f"Tunnel '{t.get('name', 'unnamed')}' missing 'remote_port'")
|
||||
|
||||
host = t["host"]
|
||||
user = t.get("user", global_user)
|
||||
if not user:
|
||||
raise ConfigError(f"Tunnel '{t.get('name', 'unnamed')}' missing 'user'")
|
||||
|
||||
key_file = t.get("key_file", global_key_file)
|
||||
ssh_port = t.get("ssh_port", 22)
|
||||
|
||||
# Create host key for grouping
|
||||
host_key = f"{user}@{host}:{ssh_port}"
|
||||
|
||||
if host_key not in hosts_map:
|
||||
hosts_map[host_key] = HostConfig(
|
||||
host=host,
|
||||
user=user,
|
||||
key_file=key_file,
|
||||
port=ssh_port,
|
||||
)
|
||||
|
||||
tunnel = TunnelConfig(
|
||||
name=t.get("name", f"{host}:{t['remote_port']}"),
|
||||
host=host,
|
||||
remote_port=t["remote_port"],
|
||||
local_port=t.get("local_port"),
|
||||
remote_host=t.get("remote_host", "localhost"),
|
||||
)
|
||||
hosts_map[host_key].tunnels.append(tunnel)
|
||||
|
||||
return list(hosts_map.values())
|
||||
|
||||
|
||||
def _parse_grouped_format(hosts: list[dict]) -> list[HostConfig]:
|
||||
"""Parse grouped format with hosts containing nested tunnels."""
|
||||
result = []
|
||||
|
||||
for h in hosts:
|
||||
if "host" not in h:
|
||||
raise ConfigError("Host entry missing 'host' field")
|
||||
if "user" not in h:
|
||||
raise ConfigError(f"Host '{h['host']}' missing 'user' field")
|
||||
|
||||
host_config = HostConfig(
|
||||
host=h["host"],
|
||||
user=h["user"],
|
||||
key_file=h.get("key_file"),
|
||||
port=h.get("port", 22),
|
||||
)
|
||||
|
||||
for t in h.get("tunnels", []):
|
||||
if "remote_port" not in t:
|
||||
raise ConfigError(f"Tunnel in host '{h['host']}' missing 'remote_port'")
|
||||
|
||||
tunnel = TunnelConfig(
|
||||
name=t.get("name", f"{h['host']}:{t['remote_port']}"),
|
||||
host=h["host"],
|
||||
remote_port=t["remote_port"],
|
||||
local_port=t.get("local_port"),
|
||||
remote_host=t.get("remote_host", "localhost"),
|
||||
)
|
||||
host_config.tunnels.append(tunnel)
|
||||
|
||||
result.append(host_config)
|
||||
|
||||
return result
|
||||
355
devbox-connect/src/devbox_connect/tunnel.py
Normal file
355
devbox-connect/src/devbox_connect/tunnel.py
Normal file
@@ -0,0 +1,355 @@
|
||||
"""SSH tunnel manager using paramiko."""
|
||||
|
||||
import logging
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import paramiko
|
||||
|
||||
from .config import Config, HostConfig, TunnelConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TunnelState(Enum):
|
||||
"""State of a tunnel."""
|
||||
|
||||
DISCONNECTED = "disconnected"
|
||||
CONNECTING = "connecting"
|
||||
CONNECTED = "connected"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TunnelStatus:
|
||||
"""Status of a single tunnel."""
|
||||
|
||||
config: TunnelConfig
|
||||
state: TunnelState = TunnelState.DISCONNECTED
|
||||
error: str | None = None
|
||||
connections: int = 0 # Active forwarded connections
|
||||
|
||||
|
||||
@dataclass
|
||||
class HostConnection:
|
||||
"""Manages SSH connection and tunnels for a single host."""
|
||||
|
||||
config: HostConfig
|
||||
tunnels: dict[str, TunnelStatus] = field(default_factory=dict)
|
||||
client: paramiko.SSHClient | None = None
|
||||
transport: paramiko.Transport | None = None
|
||||
_forward_threads: list[threading.Thread] = field(default_factory=list)
|
||||
_forward_servers: list[socket.socket] = field(default_factory=list)
|
||||
_stop_event: threading.Event = field(default_factory=threading.Event)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
for tunnel in self.config.tunnels:
|
||||
self.tunnels[tunnel.name] = TunnelStatus(config=tunnel)
|
||||
|
||||
|
||||
class TunnelManager:
|
||||
"""Manages SSH tunnels based on configuration."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Config,
|
||||
on_status_change: Callable[[str, str, TunnelState], None] | None = None,
|
||||
):
|
||||
self.config = config
|
||||
self.on_status_change = on_status_change
|
||||
self.hosts: dict[str, HostConnection] = {}
|
||||
self._stop_event = threading.Event()
|
||||
self._reconnect_threads: list[threading.Thread] = []
|
||||
|
||||
# Initialize host connections
|
||||
for host_config in config.hosts:
|
||||
key = f"{host_config.user}@{host_config.host}"
|
||||
self.hosts[key] = HostConnection(config=host_config)
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start all tunnels."""
|
||||
self._stop_event.clear()
|
||||
for host_key, host_conn in self.hosts.items():
|
||||
self._connect_host(host_key, host_conn)
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop all tunnels."""
|
||||
logger.info("Stopping all tunnels...")
|
||||
self._stop_event.set()
|
||||
|
||||
for host_conn in self.hosts.values():
|
||||
self._disconnect_host(host_conn)
|
||||
|
||||
# Wait for reconnect threads to finish
|
||||
for thread in self._reconnect_threads:
|
||||
thread.join(timeout=2)
|
||||
self._reconnect_threads.clear()
|
||||
|
||||
def get_status(self) -> dict[str, dict[str, TunnelStatus]]:
|
||||
"""Get status of all tunnels."""
|
||||
return {host_key: host.tunnels for host_key, host in self.hosts.items()}
|
||||
|
||||
def _connect_host(self, host_key: str, host_conn: HostConnection) -> bool:
|
||||
"""Connect to a host and establish all tunnels."""
|
||||
config = host_conn.config
|
||||
|
||||
# Update all tunnel states to connecting
|
||||
for tunnel_status in host_conn.tunnels.values():
|
||||
self._update_state(host_key, tunnel_status, TunnelState.CONNECTING)
|
||||
|
||||
try:
|
||||
# Create SSH client
|
||||
client = paramiko.SSHClient()
|
||||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
|
||||
# Prepare connection kwargs
|
||||
connect_kwargs: dict = {
|
||||
"hostname": config.host,
|
||||
"port": config.port,
|
||||
"username": config.user,
|
||||
}
|
||||
|
||||
# Add key file if specified
|
||||
if config.key_file:
|
||||
key_path = Path(config.key_file).expanduser()
|
||||
if not key_path.exists():
|
||||
raise FileNotFoundError(f"SSH key file not found: {key_path}")
|
||||
connect_kwargs["key_filename"] = str(key_path)
|
||||
|
||||
logger.info(f"Connecting to {config.user}@{config.host}:{config.port}...")
|
||||
client.connect(**connect_kwargs)
|
||||
|
||||
host_conn.client = client
|
||||
host_conn.transport = client.get_transport()
|
||||
|
||||
if host_conn.transport is None:
|
||||
raise ConnectionError("Failed to get transport")
|
||||
|
||||
# Start port forwarding for each tunnel
|
||||
host_conn._stop_event.clear()
|
||||
for tunnel_config in config.tunnels:
|
||||
self._start_tunnel(host_key, host_conn, tunnel_config)
|
||||
|
||||
logger.info(f"Connected to {config.host}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect to {config.host}: {e}")
|
||||
for tunnel_status in host_conn.tunnels.values():
|
||||
tunnel_status.error = str(e)
|
||||
self._update_state(host_key, tunnel_status, TunnelState.ERROR)
|
||||
self._schedule_reconnect(host_key, host_conn)
|
||||
return False
|
||||
|
||||
def _disconnect_host(self, host_conn: HostConnection) -> None:
|
||||
"""Disconnect from a host."""
|
||||
host_conn._stop_event.set()
|
||||
|
||||
# Close forward servers
|
||||
for server in host_conn._forward_servers:
|
||||
try:
|
||||
server.close()
|
||||
except Exception:
|
||||
pass
|
||||
host_conn._forward_servers.clear()
|
||||
|
||||
# Wait for forward threads
|
||||
for thread in host_conn._forward_threads:
|
||||
thread.join(timeout=1)
|
||||
host_conn._forward_threads.clear()
|
||||
|
||||
# Close SSH connection
|
||||
if host_conn.client:
|
||||
try:
|
||||
host_conn.client.close()
|
||||
except Exception:
|
||||
pass
|
||||
host_conn.client = None
|
||||
host_conn.transport = None
|
||||
|
||||
def _start_tunnel(
|
||||
self, host_key: str, host_conn: HostConnection, tunnel_config: TunnelConfig
|
||||
) -> None:
|
||||
"""Start a single port forward tunnel."""
|
||||
tunnel_status = host_conn.tunnels[tunnel_config.name]
|
||||
|
||||
try:
|
||||
# Create local listening socket
|
||||
local_port = tunnel_config.local_port or tunnel_config.remote_port
|
||||
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
server.bind(("127.0.0.1", local_port))
|
||||
server.listen(5)
|
||||
server.settimeout(1.0) # Allow checking stop event
|
||||
|
||||
host_conn._forward_servers.append(server)
|
||||
|
||||
# Start accept thread
|
||||
thread = threading.Thread(
|
||||
target=self._accept_loop,
|
||||
args=(host_key, host_conn, tunnel_config, server),
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
host_conn._forward_threads.append(thread)
|
||||
|
||||
self._update_state(host_key, tunnel_status, TunnelState.CONNECTED)
|
||||
logger.info(
|
||||
f"Tunnel '{tunnel_config.name}' open: "
|
||||
f"localhost:{local_port} -> {tunnel_config.remote_host}:{tunnel_config.remote_port}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start tunnel '{tunnel_config.name}': {e}")
|
||||
tunnel_status.error = str(e)
|
||||
self._update_state(host_key, tunnel_status, TunnelState.ERROR)
|
||||
|
||||
def _accept_loop(
|
||||
self,
|
||||
host_key: str,
|
||||
host_conn: HostConnection,
|
||||
tunnel_config: TunnelConfig,
|
||||
server: socket.socket,
|
||||
) -> None:
|
||||
"""Accept loop for forwarding connections."""
|
||||
tunnel_status = host_conn.tunnels[tunnel_config.name]
|
||||
|
||||
while not host_conn._stop_event.is_set() and not self._stop_event.is_set():
|
||||
try:
|
||||
client_socket, addr = server.accept()
|
||||
except socket.timeout:
|
||||
continue
|
||||
except OSError:
|
||||
break # Socket was closed
|
||||
|
||||
# Check if transport is still active
|
||||
if host_conn.transport is None or not host_conn.transport.is_active():
|
||||
client_socket.close()
|
||||
self._update_state(host_key, tunnel_status, TunnelState.ERROR)
|
||||
tunnel_status.error = "SSH connection lost"
|
||||
break
|
||||
|
||||
# Open channel to remote
|
||||
try:
|
||||
channel = host_conn.transport.open_channel(
|
||||
"direct-tcpip",
|
||||
(tunnel_config.remote_host, tunnel_config.remote_port),
|
||||
client_socket.getpeername(),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to open channel for {tunnel_config.name}: {e}")
|
||||
client_socket.close()
|
||||
continue
|
||||
|
||||
if channel is None:
|
||||
logger.error(f"Failed to open channel for {tunnel_config.name}")
|
||||
client_socket.close()
|
||||
continue
|
||||
|
||||
# Start forwarding thread
|
||||
tunnel_status.connections += 1
|
||||
thread = threading.Thread(
|
||||
target=self._forward_data,
|
||||
args=(client_socket, channel, tunnel_status),
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
|
||||
# Connection lost, trigger reconnect
|
||||
if not self._stop_event.is_set():
|
||||
self._schedule_reconnect(host_key, host_conn)
|
||||
|
||||
def _forward_data(
|
||||
self,
|
||||
client_socket: socket.socket,
|
||||
channel: paramiko.Channel,
|
||||
tunnel_status: TunnelStatus,
|
||||
) -> None:
|
||||
"""Forward data between client socket and SSH channel."""
|
||||
try:
|
||||
while True:
|
||||
# Check both directions for data
|
||||
r_ready = []
|
||||
|
||||
# Use select for multiplexing
|
||||
import select
|
||||
|
||||
try:
|
||||
r_ready, _, _ = select.select([client_socket, channel], [], [], 1.0)
|
||||
except Exception:
|
||||
break
|
||||
|
||||
if client_socket in r_ready:
|
||||
data = client_socket.recv(4096)
|
||||
if not data:
|
||||
break
|
||||
channel.send(data)
|
||||
|
||||
if channel in r_ready:
|
||||
data = channel.recv(4096)
|
||||
if not data:
|
||||
break
|
||||
client_socket.send(data)
|
||||
|
||||
except Exception as e:
|
||||
logger.debug(f"Forward connection closed: {e}")
|
||||
finally:
|
||||
tunnel_status.connections -= 1
|
||||
try:
|
||||
client_socket.close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
channel.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _schedule_reconnect(self, host_key: str, host_conn: HostConnection) -> None:
|
||||
"""Schedule a reconnection attempt."""
|
||||
if self._stop_event.is_set():
|
||||
return
|
||||
|
||||
thread = threading.Thread(
|
||||
target=self._reconnect_loop,
|
||||
args=(host_key, host_conn),
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
self._reconnect_threads.append(thread)
|
||||
|
||||
def _reconnect_loop(self, host_key: str, host_conn: HostConnection) -> None:
|
||||
"""Reconnection loop with exponential backoff."""
|
||||
delay = self.config.reconnect_delay
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
logger.info(f"Reconnecting to {host_conn.config.host} in {delay}s...")
|
||||
|
||||
# Wait with stop check
|
||||
for _ in range(delay):
|
||||
if self._stop_event.is_set():
|
||||
return
|
||||
time.sleep(1)
|
||||
|
||||
# Disconnect first
|
||||
self._disconnect_host(host_conn)
|
||||
|
||||
# Try to connect
|
||||
if self._connect_host(host_key, host_conn):
|
||||
return # Success
|
||||
|
||||
# Exponential backoff
|
||||
delay = min(delay * 2, self.config.max_reconnect_delay)
|
||||
|
||||
def _update_state(self, host_key: str, tunnel_status: TunnelStatus, state: TunnelState) -> None:
|
||||
"""Update tunnel state and notify callback."""
|
||||
tunnel_status.state = state
|
||||
if state != TunnelState.ERROR:
|
||||
tunnel_status.error = None
|
||||
|
||||
if self.on_status_change:
|
||||
self.on_status_change(host_key, tunnel_status.config.name, state)
|
||||
Reference in New Issue
Block a user