premiere-libtorrent/test/socks.py

#!/usr/bin/python

"""Minimal non-feature complete socks proxy"""

import logging
from logging import error, info, debug
import random
import socket
from SocketServer import StreamRequestHandler, ThreadingTCPServer
from struct import pack, unpack
import threading
import sys

class MyTCPServer(ThreadingTCPServer):
    allow_reuse_address = True

    def handle_timeout(self):
        raise Exception('timeout')

CLOSE = object()

logging.basicConfig(stream=sys.stderr, level=logging.INFO)

VERSION = '\x05'
NOAUTH = '\x00'
USERPASS = '\x02'
CONNECT = '\x01'
IPV4 = '\x01'
IPV6 = '\x04'
DOMAIN_NAME = '\x03'
SUCCESS = '\x00'

password = None
username = None
allow_v4 = False

def send(dest, msg):
    if msg == CLOSE:
        try: dest.shutdown(socket.SHUT_WR)
        except: pass
        dest.close()
        return 0
    else:
        return dest.sendall(msg)

def recv(source, buffer):
    data = source.recv(buffer)
    if data == '':
        return CLOSE
    else:
        return data

def forward(source, dest, name):
    while True:
        data = recv(source, 4000)
        if data == CLOSE:
            send(dest, CLOSE)
            info('%s hung up' % name)
            return
        debug('Sending (%d) %r' % (len(data), data))
        send(dest, data)

def spawn_forwarder(source, dest, name):
    t = threading.Thread(target=forward, args=(source, dest, name))
    t.daemon = True
    t.start()

class SocksHandler(StreamRequestHandler):
    """Highly feature incomplete SOCKS 5 implementation"""

    def close_request(self):
        self.server.close_request(self.request)

    def read(self, n):
        data = ''
        while len(data) < n:
            extra = self.rfile.read(n)
            if extra == '':
                raise Exception('Connection closed')
            data += extra
        return data

    def handle(self):
        # IMRPOVEMENT: Report who requests are from in logging
        # IMPROVEMENT: Timeout on client
        info('Connection - authenticating')
        version = self.read(1)

        if allow_v4 and version == '\x04':
            cmd = self.read(1)
            if cmd != CONNECT:
                error('Only supports connect method not (%r) closing' % cmd)
                self.close_request()
                return

            raw_dest_port = self.read(2)
            dest_port, = unpack('>H', raw_dest_port)

            raw_dest_address = self.read(4)
            dest_address = '.'.join(map(str, unpack('>4B', raw_dest_address)))

            user_id = ''
            c = self.read(1)
            while c != '\0':
                user_id += c
                c = self.read(1)

            outbound_sock = socket.socket(socket.AF_INET)
            out_address = socket.getaddrinfo(dest_address,dest_port)[0][4]
            debug("Creating forwarder connection to %r", out_address)
            outbound_sock.connect(out_address)

            self.send_reply_v4(outbound_sock.getsockname())

            spawn_forwarder(outbound_sock, self.request, 'destination')
            forward(self.request, outbound_sock, 'client')
            return

        if version != '\x05':
            error('Wrong version number (%r) closing...' % version)
            self.close_request()
            return

        nmethods = ord(self.read(1))
        method_list = self.read(nmethods)

        global password
        global username

        if password == None and NOAUTH in method_list:
            self.send_no_auth_method()
            info('Authenticated (no-auth)')
        elif USERPASS in method_list:
            self.send_user_pass_auth_method()
            auth_version = self.read(1)
            if auth_version != '\x01':
                error('Wrong sub-negotiation version number (%r) closing...' % version)
                self.close_request()
                return
            usr_len = ord(self.read(1))
            usr_name = self.read(usr_len)
            pwd_len = ord(self.read(1))
            pwd = self.read(pwd_len)

            if usr_name != username or pwd != password:
                error('Invalid username or password')
                self.close_request()
                return
            info('Authenticated (user/password)')
            self.send_authenticated()
        else:
            error('Server only supports NOAUTH and user/pass')
            self.send_no_method()
            return

        # If we were authenticating it would go here
        version, cmd, zero, address_type = self.read(4)
        if version != '\x05':
            error('Wrong version number (%r) closing...' % version)
            self.close_request()
        elif cmd != CONNECT:
            error('Only supports connect method not (%r) closing' % cmd)
            self.close_request()
        elif zero != '\x00':
            error('Mangled request. Reserved field (%r) is not null' % zero)
            self.close_request()

        if address_type == IPV4:
            raw_dest_address = self.read(4)
            dest_address = '.'.join(map(str, unpack('>4B', raw_dest_address)))
        elif address_type == IPV6:
            raw_dest_address = self.read(16)
            dest_address = ":".join(map(lambda x: hex(x)[2:],unpack('>8H',raw_dest_address)))
        elif address_type == DOMAIN_NAME:
            dns_length = ord(self.read(1))
            dns_name = self.read(dns_length)
            dest_address = dns_name
        else:
            error('Unknown addressing (%r)' % address_type)
            self.close_request()

        raw_dest_port = self.read(2)
        dest_port, = unpack('>H', raw_dest_port)

        if address_type == IPV6:
            outbound_sock = socket.socket(socket.AF_INET6)
        else:
            outbound_sock = socket.socket(socket.AF_INET)
        try:
            out_address = socket.getaddrinfo(dest_address,dest_port)[0][4]
        except Exception, e:
            print e
            return
        debug("Creating forwarder connection to %r", out_address)
        outbound_sock.connect(out_address)

        if address_type == IPV6:
            self.send_reply6(outbound_sock.getsockname())
        else:
            self.send_reply(outbound_sock.getsockname())

        spawn_forwarder(outbound_sock, self.request, 'destination')
        forward(self.request, outbound_sock, 'client')

    def send_reply_v4(self, (bind_addr, bind_port)):
        self.wfile.write('\0\x5a\0\0\0\0\0\0')
        self.wfile.flush()

    def send_reply(self, (bind_addr, bind_port)):
        bind_tuple = tuple(map(int, bind_addr.split('.')))
        full_address = bind_tuple + (bind_port,)
        info('Setting up forwarding port %r' % (full_address,))
        msg = pack('>cccc4BH', VERSION, SUCCESS, '\x00', IPV4, *full_address)
        self.wfile.write(msg)

    def send_reply6(self, (bind_addr, bind_port, unused1, unused2)):
        bind_tuple = tuple(map(lambda x: int(x,16), bind_addr.split(':')))
        full_address = bind_tuple + (bind_port,)
        info('Setting up forwarding port %r' % (full_address,))
        msg = pack('>cccc8HH', VERSION, SUCCESS, '\x00', IPV6, *full_address)
        self.wfile.write(msg)

    def send_no_method(self):
        self.wfile.write('\x05\xff')
        self.close_request()

    def send_no_auth_method(self):
        self.wfile.write('\x05\x00')
        self.wfile.flush()

    def send_user_pass_auth_method(self):
        self.wfile.write('\x05\x02')
        self.wfile.flush()

    def send_authenticated(self):
        self.wfile.write('\x01\x00')
        self.wfile.flush()

if __name__ == '__main__':

    listen_port = 8002
    i = 1
    while i < len(sys.argv):
        if sys.argv[i] == '--username':
            username = sys.argv[i+1]
            i += 1
        elif sys.argv[i] == '--password':
            password = sys.argv[i+1]
            i += 1
        elif sys.argv[i] == '--port':
            listen_port = int(sys.argv[i+1])
            i += 1
        elif sys.argv[i] == '--allow-v4':
            allow_v4 = True
        else:
            if sys.argv[i] != '--help': info('unknown option "%s"' % sys.argv[i])
            print('usage: socks.py [--username <user> --password <password>] [--port <listen-port>]')
            sys.exit(1)
        i += 1

    info('Listening on port %d...' % listen_port)
    server = MyTCPServer(('localhost', listen_port), SocksHandler)
    server.timeout = 120
    while True:
        server.handle_request()