cli.py

import socket
import struct
import threading
import argparse
import logging
import time
import sys

# Import core protocol
try:
    from mux_core import SessionManager, CMD_OPEN, CMD_CLOSE, CMD_LISTEN
except ImportError:
    sys.stderr.write("Error: mux_core.py not found.\n")
    sys.exit(1)

# Import transports
try:
    from transport.proc import SubprocessTransport
    from transport.sock import SocketTransport
    from transport.ws import WebSocketTransport
    from transport.ssh import SSHTransport
except ImportError:
    sys.stderr.write("Error: transport package not found.\n")
    sys.exit(1)

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger(__name__)

# --- Client Session Manager ---


class ClientSessionManager(SessionManager):
    def __init__(self, reader, writer):
        super().__init__(reader, writer)
        self.r_rules = {}  # {rport: (lhost, lport)}

    def handle_open(self, cid, data):
        """Handle incoming connection request (R Mode)."""
        # Payload: Port(2) - The remote port that received a connection
        try:
            rport = struct.unpack("!H", data)[0]
            if rport in self.r_rules:
                lhost, lport = self.r_rules[rport]
                logger.info(
                    f"[R] Remote accepted on {rport}, connecting to {lhost}:{lport} (CID={cid})"
                )

                def _opener():
                    return socket.create_connection((lhost, lport), timeout=10)

                self.open_session_async(cid, _opener)
            else:
                logger.error(f"[R] Unknown remote port {rport} for CID={cid}")
                self.mux.send_packet(cid, CMD_CLOSE)
        except Exception as e:
            logger.error(f"[R] Failed to connect local target: {e}")
            self.mux.send_packet(cid, CMD_CLOSE)

    def start_local_listener(self, lport, rhost, rport):
        """L Mode: Listen locally, forward to remote."""
        try:
            server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            server.bind(("127.0.0.1", lport))
            server.listen(100)
            logger.info(f"[L] Listening on 127.0.0.1:{lport} -> {rhost}:{rport}")

            while self.running:
                conn, addr = server.accept()

                # Allocate CID (Use even numbers for client initiated)
                cid = self.allocate_cid(parity=0)

                logger.info(f"[L] Accepted {addr}, assigning CID={cid}")

                # Send OPEN: HostLen(2) | Host | Port(2)
                host_bytes = rhost.encode()
                payload = (
                    struct.pack("!H", len(host_bytes))
                    + host_bytes
                    + struct.pack("!H", rport)
                )
                self.mux.send_packet(cid, CMD_OPEN, payload)

                self.add_session(cid, conn)
        except Exception as e:
            logger.error(f"[L] Listener failed: {e}")

    def request_remote_listener(self, rport, lhost, lport):
        """R Mode: Tell remote to listen."""
        self.r_rules[rport] = (lhost, lport)
        payload = struct.pack("!H", rport)
        self.mux.send_packet(0, CMD_LISTEN, payload)
        logger.info(f"[R] Requested remote listen on port {rport} -> {lhost}:{lport}")


def run_client(transport, args):
    try:
        reader, writer = transport.connect()
    except Exception as e:
        logger.error(f"Transport connection failed: {e}")
        return

    manager = ClientSessionManager(reader, writer)

    # Start Protocol Loop in a thread so we can start listeners
    t_loop = threading.Thread(target=manager.run_loop, daemon=True)
    t_loop.start()

    # Configure Forwarding Rules
    if args.L:
        try:
            lport, rhost, rport = args.L.split(":")
            lport, rport = int(lport), int(rport)
            threading.Thread(
                target=manager.start_local_listener,
                args=(lport, rhost, rport),
                daemon=True,
            ).start()
        except Exception as e:
            logger.error(f"Invalid -L argument: {e}")
            manager.stop()
            return

    if args.R:
        try:
            rport, lhost, lport = args.R.split(":")
            rport, lport = int(rport), int(lport)
            # Give the loop a moment to establish
            time.sleep(0.5)
            manager.request_remote_listener(rport, lhost, lport)
        except Exception as e:
            logger.error(f"Invalid -R argument: {e}")
            manager.stop()
            return

    # Keep main thread alive
    try:
        while manager.running:
            time.sleep(1)
    except KeyboardInterrupt:
        pass
    finally:
        transport.close()
        manager.stop()


def main():
    parser = argparse.ArgumentParser(description="stdioOverTCP CLI (Decoupled)")
    parser.add_argument(
        "--cmd", help="Command transport (e.g. 'ssh user@host python srv.py')"
    )
    parser.add_argument("--socket", help="Socket transport (e.g. '127.0.0.1:8888')")
    parser.add_argument("--ws", help="WebSocket transport (e.g. 'ws://127.0.0.1:8765')")
    parser.add_argument(
        "--ssh",
        help="Use Legacy SSH Transport (configured in transport/ssh.py)",
    )
    parser.add_argument("-L", help="local_port:remote_host:remote_port")
    parser.add_argument("-R", help="remote_port:local_host:local_port")
    args = parser.parse_args()

    transport = None
    if args.cmd:
        transport = SubprocessTransport(args.cmd)
    elif args.socket:
        try:
            host, port = args.socket.split(":")
            transport = SocketTransport(host, int(port))
        except:
            print("Invalid socket format. Use host:port")
            return
    elif args.ws:
        transport = WebSocketTransport(args.ws)
    elif args.ssh:
        # Default command for SSH transport is python srv.py, or user can override if we add arg
        # For now, we use the default hardcoded in transport/ssh.py if we want, or pass it here
        # But wait, SSHTransport init takes 'command'.
        # Let's use args.cmd if provided, otherwise default.
        # But args.cmd is mutually exclusive in my if/else block above.
        # Let's change the logic.
        #cmd = "python srv.py"  # Default
        transport = SSHTransport(args.ssh)
    else:
        print("Error: You must specify a transport (--cmd, --socket, --ws, or --ssh)")
        return

    run_client(transport, args)


if __name__ == "__main__":
    main()