RDP via SOCKS

SOCKS 経由で RDP 接続するためのローカル TCP relay. 色々と思うところはあるけど、とりあえず動く. 作り込みとかはまた考える.

from socket import socket, inet_aton, AF_INET, SOCK_STREAM
from struct import pack
from threading import Thread

relay_server = '127.0.0.1:13389'
socks_server = 'socks-server:1080'
rdp_server = '192.0.2.1:3389'

def main():
    BUFFER_SIZE = 2048

    def relay(src, dst):
        try:
            while 1:
                data = src.recv(BUFFER_SIZE)
                if not data: break
                dst.send(data)
        except:
            pass

    def parse_server(s):
        host, port = s.split(':')
        return (host, int(port))

    def socks_connect():
        def is_ipv4(host):
            try:
                t = host.split('.')
                return (len(t) == 4) and all(map(lambda x: (x >= 0) and (x <= 255), [int(x) for x in t]))
            except:
                return False
        def destination():
            (host, port) = parse_server(rdp_server)
            if is_ipv4(host):
                return '\x01' + inet_aton(host) + pack('>H', port)
            else:
                return '\x03' + chr(len(host)) + host + pack('>H', port)
        SOCKSV5_CONNECT = '\x05\x01\x00'
        dst = socket(AF_INET, SOCK_STREAM)
        dst.connect(parse_server(socks_server))
        dst.sendall(SOCKSV5_CONNECT)
        dst.recv(BUFFER_SIZE)
        SOCKSV5_REQUEST = '\x05\x01\x00'
        dst.sendall(SOCKSV5_REQUEST + destination())
        dst.recv(BUFFER_SIZE)
        return dst

    s = socket(AF_INET, SOCK_STREAM)
    s.bind(parse_server(relay_server))
    s.listen(1)
    while 1:
        (src, _) = s.accept()
        dst = socks_connect()
        Thread(target = relay, args = (dst, src)).start()
        Thread(target = relay, args = (src, dst)).start()

if __name__ == '__main__':
    main()