socket_proxy.py

#!/usr/bin/env python3
"""
JoySword Enterprise Application-Level TCP Proxy / Port Relay.
Provides zero-admin fallback routing for internal process coordination (TCP 9100)
when elevated network interface modification is unavailable or skipped.
Implements TCP_NODELAY, SO_KEEPALIVE, Signal Handling, and Graceful Teardown.
"""

from __future__ import annotations

import argparse
import logging
import select
import signal
import socket
import sys
import threading
import time

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] [Proxy] %(message)s",
    datefmt="%H:%M:%S",
)

BUFFER_SIZE = 65536
SHUTDOWN_EVENT = threading.Event()


def set_socket_options(sock: socket.socket) -> None:
    try:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
    except OSError:
        pass
    try:
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    except OSError:
        pass


def handle_client(client_sock: socket.socket, client_addr: tuple[str, int], target_host: str, target_port: int) -> None:
    set_socket_options(client_sock)
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    set_socket_options(server_sock)

    try:
        server_sock.settimeout(5.0)
        server_sock.connect((target_host, target_port))
        server_sock.settimeout(None)
    except Exception as exc:
        logging.warning(f"Connection from {client_addr[0]}:{client_addr[1]} failed to reach target {target_host}:{target_port}: {exc}")
        client_sock.close()
        return

    sockets = [client_sock, server_sock]
    try:
        while not SHUTDOWN_EVENT.is_set():
            readable, _, exceptional = select.select(sockets, [], sockets, 1.0)
            if exceptional:
                break
            for sock in readable:
                data = sock.recv(BUFFER_SIZE)
                if not data:
                    return
                other_sock = server_sock if sock is client_sock else client_sock
                other_sock.sendall(data)
    except Exception:
        pass
    finally:
        try:
            client_sock.close()
        except OSError:
            pass
        try:
            server_sock.close()
        except OSError:
            pass


def start_proxy(listen_host: str = "0.0.0.0", listen_port: int = 9100, target_host: str = "127.0.0.1", target_port: int = 9100) -> None:
    logging.info(f"Initializing TCP proxy bridge: {listen_host}:{listen_port} -> {target_host}:{target_port}")
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        server.bind((listen_host, listen_port))
        server.listen(128)
        server.settimeout(1.0)
    except Exception as exc:
        logging.warning(f"Socket bind failed on {listen_host}:{listen_port} (port may already be in use natively): {exc}")
        return

    def signal_handler(signum, frame):
        logging.info("Received termination signal. Shutting down proxy daemon...")
        SHUTDOWN_EVENT.set()

    try:
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)
    except (AttributeError, ValueError):
        pass

    logging.info(f"TCP Proxy Daemon active and listening on {listen_host}:{listen_port}.")

    active_threads: list[threading.Thread] = []
    while not SHUTDOWN_EVENT.is_set():
        try:
            client_sock, client_addr = server.accept()
            thread = threading.Thread(
                target=handle_client,
                args=(client_sock, client_addr, target_host, target_port),
                daemon=True,
            )
            thread.start()
            active_threads.append(thread)
            active_threads = [t for t in active_threads if t.is_alive()]
        except socket.timeout:
            continue
        except Exception as exc:
            if not SHUTDOWN_EVENT.is_set():
                logging.error(f"Error accepting client connection: {exc}")
            break

    try:
        server.close()
    except OSError:
        pass
    logging.info("TCP Proxy Daemon stopped cleanly.")


def main() -> None:
    parser = argparse.ArgumentParser(description="JoySword Enterprise Application TCP Proxy Relay")
    parser.add_argument("--listen-host", default="0.0.0.0", help="Listen host address (default: 0.0.0.0)")
    parser.add_argument("--listen-port", type=int, default=9100, help="Listen port (default: 9100)")
    parser.add_argument("--target-host", default="127.0.0.1", help="Target host address (default: 127.0.0.1)")
    parser.add_argument("--target-port", type=int, default=9100, help="Target port (default: 9100)")
    args = parser.parse_args()

    start_proxy(args.listen_host, args.listen_port, args.target_host, args.target_port)


if __name__ == "__main__":
    main()