youtube-dl

Another place where youtube-dl lives on
git clone git://git.oshgnacknak.de/youtube-dl.git
Log | Files | Refs | README | LICENSE

socks.py (8861B)


      1 # Public Domain SOCKS proxy protocol implementation
      2 # Adapted from https://gist.github.com/bluec0re/cafd3764412967417fd3
      3 
      4 from __future__ import unicode_literals
      5 
      6 # References:
      7 # SOCKS4 protocol http://www.openssh.com/txt/socks4.protocol
      8 # SOCKS4A protocol http://www.openssh.com/txt/socks4a.protocol
      9 # SOCKS5 protocol https://tools.ietf.org/html/rfc1928
     10 # SOCKS5 username/password authentication https://tools.ietf.org/html/rfc1929
     11 
     12 import collections
     13 import socket
     14 
     15 from .compat import (
     16     compat_ord,
     17     compat_struct_pack,
     18     compat_struct_unpack,
     19 )
     20 
     21 __author__ = 'Timo Schmid <coding@timoschmid.de>'
     22 
     23 SOCKS4_VERSION = 4
     24 SOCKS4_REPLY_VERSION = 0x00
     25 # Excerpt from SOCKS4A protocol:
     26 # if the client cannot resolve the destination host's domain name to find its
     27 # IP address, it should set the first three bytes of DSTIP to NULL and the last
     28 # byte to a non-zero value.
     29 SOCKS4_DEFAULT_DSTIP = compat_struct_pack('!BBBB', 0, 0, 0, 0xFF)
     30 
     31 SOCKS5_VERSION = 5
     32 SOCKS5_USER_AUTH_VERSION = 0x01
     33 SOCKS5_USER_AUTH_SUCCESS = 0x00
     34 
     35 
     36 class Socks4Command(object):
     37     CMD_CONNECT = 0x01
     38     CMD_BIND = 0x02
     39 
     40 
     41 class Socks5Command(Socks4Command):
     42     CMD_UDP_ASSOCIATE = 0x03
     43 
     44 
     45 class Socks5Auth(object):
     46     AUTH_NONE = 0x00
     47     AUTH_GSSAPI = 0x01
     48     AUTH_USER_PASS = 0x02
     49     AUTH_NO_ACCEPTABLE = 0xFF  # For server response
     50 
     51 
     52 class Socks5AddressType(object):
     53     ATYP_IPV4 = 0x01
     54     ATYP_DOMAINNAME = 0x03
     55     ATYP_IPV6 = 0x04
     56 
     57 
     58 class ProxyError(socket.error):
     59     ERR_SUCCESS = 0x00
     60 
     61     def __init__(self, code=None, msg=None):
     62         if code is not None and msg is None:
     63             msg = self.CODES.get(code) or 'unknown error'
     64         super(ProxyError, self).__init__(code, msg)
     65 
     66 
     67 class InvalidVersionError(ProxyError):
     68     def __init__(self, expected_version, got_version):
     69         msg = ('Invalid response version from server. Expected {0:02x} got '
     70                '{1:02x}'.format(expected_version, got_version))
     71         super(InvalidVersionError, self).__init__(0, msg)
     72 
     73 
     74 class Socks4Error(ProxyError):
     75     ERR_SUCCESS = 90
     76 
     77     CODES = {
     78         91: 'request rejected or failed',
     79         92: 'request rejected because SOCKS server cannot connect to identd on the client',
     80         93: 'request rejected because the client program and identd report different user-ids'
     81     }
     82 
     83 
     84 class Socks5Error(ProxyError):
     85     ERR_GENERAL_FAILURE = 0x01
     86 
     87     CODES = {
     88         0x01: 'general SOCKS server failure',
     89         0x02: 'connection not allowed by ruleset',
     90         0x03: 'Network unreachable',
     91         0x04: 'Host unreachable',
     92         0x05: 'Connection refused',
     93         0x06: 'TTL expired',
     94         0x07: 'Command not supported',
     95         0x08: 'Address type not supported',
     96         0xFE: 'unknown username or invalid password',
     97         0xFF: 'all offered authentication methods were rejected'
     98     }
     99 
    100 
    101 class ProxyType(object):
    102     SOCKS4 = 0
    103     SOCKS4A = 1
    104     SOCKS5 = 2
    105 
    106 
    107 Proxy = collections.namedtuple('Proxy', (
    108     'type', 'host', 'port', 'username', 'password', 'remote_dns'))
    109 
    110 
    111 class sockssocket(socket.socket):
    112     def __init__(self, *args, **kwargs):
    113         self._proxy = None
    114         super(sockssocket, self).__init__(*args, **kwargs)
    115 
    116     def setproxy(self, proxytype, addr, port, rdns=True, username=None, password=None):
    117         assert proxytype in (ProxyType.SOCKS4, ProxyType.SOCKS4A, ProxyType.SOCKS5)
    118 
    119         self._proxy = Proxy(proxytype, addr, port, username, password, rdns)
    120 
    121     def recvall(self, cnt):
    122         data = b''
    123         while len(data) < cnt:
    124             cur = self.recv(cnt - len(data))
    125             if not cur:
    126                 raise EOFError('{0} bytes missing'.format(cnt - len(data)))
    127             data += cur
    128         return data
    129 
    130     def _recv_bytes(self, cnt):
    131         data = self.recvall(cnt)
    132         return compat_struct_unpack('!{0}B'.format(cnt), data)
    133 
    134     @staticmethod
    135     def _len_and_data(data):
    136         return compat_struct_pack('!B', len(data)) + data
    137 
    138     def _check_response_version(self, expected_version, got_version):
    139         if got_version != expected_version:
    140             self.close()
    141             raise InvalidVersionError(expected_version, got_version)
    142 
    143     def _resolve_address(self, destaddr, default, use_remote_dns):
    144         try:
    145             return socket.inet_aton(destaddr)
    146         except socket.error:
    147             if use_remote_dns and self._proxy.remote_dns:
    148                 return default
    149             else:
    150                 return socket.inet_aton(socket.gethostbyname(destaddr))
    151 
    152     def _setup_socks4(self, address, is_4a=False):
    153         destaddr, port = address
    154 
    155         ipaddr = self._resolve_address(destaddr, SOCKS4_DEFAULT_DSTIP, use_remote_dns=is_4a)
    156 
    157         packet = compat_struct_pack('!BBH', SOCKS4_VERSION, Socks4Command.CMD_CONNECT, port) + ipaddr
    158 
    159         username = (self._proxy.username or '').encode('utf-8')
    160         packet += username + b'\x00'
    161 
    162         if is_4a and self._proxy.remote_dns:
    163             packet += destaddr.encode('utf-8') + b'\x00'
    164 
    165         self.sendall(packet)
    166 
    167         version, resp_code, dstport, dsthost = compat_struct_unpack('!BBHI', self.recvall(8))
    168 
    169         self._check_response_version(SOCKS4_REPLY_VERSION, version)
    170 
    171         if resp_code != Socks4Error.ERR_SUCCESS:
    172             self.close()
    173             raise Socks4Error(resp_code)
    174 
    175         return (dsthost, dstport)
    176 
    177     def _setup_socks4a(self, address):
    178         self._setup_socks4(address, is_4a=True)
    179 
    180     def _socks5_auth(self):
    181         packet = compat_struct_pack('!B', SOCKS5_VERSION)
    182 
    183         auth_methods = [Socks5Auth.AUTH_NONE]
    184         if self._proxy.username and self._proxy.password:
    185             auth_methods.append(Socks5Auth.AUTH_USER_PASS)
    186 
    187         packet += compat_struct_pack('!B', len(auth_methods))
    188         packet += compat_struct_pack('!{0}B'.format(len(auth_methods)), *auth_methods)
    189 
    190         self.sendall(packet)
    191 
    192         version, method = self._recv_bytes(2)
    193 
    194         self._check_response_version(SOCKS5_VERSION, version)
    195 
    196         if method == Socks5Auth.AUTH_NO_ACCEPTABLE or (
    197                 method == Socks5Auth.AUTH_USER_PASS and (not self._proxy.username or not self._proxy.password)):
    198             self.close()
    199             raise Socks5Error(Socks5Auth.AUTH_NO_ACCEPTABLE)
    200 
    201         if method == Socks5Auth.AUTH_USER_PASS:
    202             username = self._proxy.username.encode('utf-8')
    203             password = self._proxy.password.encode('utf-8')
    204             packet = compat_struct_pack('!B', SOCKS5_USER_AUTH_VERSION)
    205             packet += self._len_and_data(username) + self._len_and_data(password)
    206             self.sendall(packet)
    207 
    208             version, status = self._recv_bytes(2)
    209 
    210             self._check_response_version(SOCKS5_USER_AUTH_VERSION, version)
    211 
    212             if status != SOCKS5_USER_AUTH_SUCCESS:
    213                 self.close()
    214                 raise Socks5Error(Socks5Error.ERR_GENERAL_FAILURE)
    215 
    216     def _setup_socks5(self, address):
    217         destaddr, port = address
    218 
    219         ipaddr = self._resolve_address(destaddr, None, use_remote_dns=True)
    220 
    221         self._socks5_auth()
    222 
    223         reserved = 0
    224         packet = compat_struct_pack('!BBB', SOCKS5_VERSION, Socks5Command.CMD_CONNECT, reserved)
    225         if ipaddr is None:
    226             destaddr = destaddr.encode('utf-8')
    227             packet += compat_struct_pack('!B', Socks5AddressType.ATYP_DOMAINNAME)
    228             packet += self._len_and_data(destaddr)
    229         else:
    230             packet += compat_struct_pack('!B', Socks5AddressType.ATYP_IPV4) + ipaddr
    231         packet += compat_struct_pack('!H', port)
    232 
    233         self.sendall(packet)
    234 
    235         version, status, reserved, atype = self._recv_bytes(4)
    236 
    237         self._check_response_version(SOCKS5_VERSION, version)
    238 
    239         if status != Socks5Error.ERR_SUCCESS:
    240             self.close()
    241             raise Socks5Error(status)
    242 
    243         if atype == Socks5AddressType.ATYP_IPV4:
    244             destaddr = self.recvall(4)
    245         elif atype == Socks5AddressType.ATYP_DOMAINNAME:
    246             alen = compat_ord(self.recv(1))
    247             destaddr = self.recvall(alen)
    248         elif atype == Socks5AddressType.ATYP_IPV6:
    249             destaddr = self.recvall(16)
    250         destport = compat_struct_unpack('!H', self.recvall(2))[0]
    251 
    252         return (destaddr, destport)
    253 
    254     def _make_proxy(self, connect_func, address):
    255         if not self._proxy:
    256             return connect_func(self, address)
    257 
    258         result = connect_func(self, (self._proxy.host, self._proxy.port))
    259         if result != 0 and result is not None:
    260             return result
    261         setup_funcs = {
    262             ProxyType.SOCKS4: self._setup_socks4,
    263             ProxyType.SOCKS4A: self._setup_socks4a,
    264             ProxyType.SOCKS5: self._setup_socks5,
    265         }
    266         setup_funcs[self._proxy.type](address)
    267         return result
    268 
    269     def connect(self, address):
    270         self._make_proxy(socket.socket.connect, address)
    271 
    272     def connect_ex(self, address):
    273         return self._make_proxy(socket.socket.connect_ex, address)