pymux: run multiple services in one port (connection multiplexer)

mux.py

import select
import socket
import socketserver

__author__ = 'loncat <[email protected]>'

services = [
    {
        'name': 'ssh',
        'address': ('127.0.0.1', 22),
        'validator': lambda b: b.startswith(b'SSH')
    },
    {
        'name': 'aria2 rpc',
        'address': ('127.0.0.1', 6800),
        'validator': lambda b: b'/jsonrpc' in b
    },
    {
        'name': 'apache',
        'address': ('127.0.0.1', 80),
        'validator': lambda b: b'HTTP/' in b
    },
    {  # if nothing matches, forward to socks5 proxy
        'name': 'socks5',
        'address': ('127.0.0.1', 7222),
        'validator': lambda b: True
    }
]


class ThreadedTCPRequestHandler(socketserver.BaseRequestHandler):
    def handle(self):
        src_sock = self.request
        dst_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        tmp_buff = b''
        if select.select([src_sock], [], [], 3)[0]:
            tmp_buff = src_sock.recv(1024)
        dst = None
        for svc in services:
            if svc['validator'](tmp_buff):
                dst = svc
                break
        if not dst:
            return src_sock.close()

        try:
            dst_sock.connect(dst['address'])
        except Exception as e:
            print(f'unable to connect to {dst["name"]} service: {e}')
            return src_sock.close()
        dst_sock.sendall(tmp_buff)

        conn_string = ' '.join([
            ':'.join(str(s) for s in src_sock.getpeername()), '<->',
            ':'.join(str(s) for s in dst_sock.getpeername()),
            f'({dst["name"]})'
        ])
        print(f'{conn_string} connection established')

        while src_sock and dst_sock:
            r1, w1, _ = select.select([dst_sock], [src_sock], [])
            r2, w2, _ = select.select([src_sock], [dst_sock], [])
            if r1 and w2:
                data = dst_sock.recv(1024)
                if not data:
                    break
                src_sock.sendall(data)
            if r2 and w2:
                data = src_sock.recv(1024)
                if not data:
                    break
                dst_sock.sendall(data)
        print(f'{conn_string} connection terminated')


class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True


if __name__ == '__main__':
    host, port = '127.0.0.1', 8000
    server = ThreadedTCPServer((host, port), ThreadedTCPRequestHandler)
    print('pymux \u2014 connection multiplexer')
    print(f'server started at {host}:{port}')
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        server.shutdown()