Raphaël G. Git Repositories - youtubedl/blob - youtube_dl/socks.py

   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(IOError):
  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) and '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 IOError('{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:
 197             self.close()
 198             raise Socks5Error(method)
 199
 200         if method == Socks5Auth.AUTH_USER_PASS:
 201             username = self._proxy.username.encode('utf-8')
 202             password = self._proxy.password.encode('utf-8')
 203             packet = compat_struct_pack('!B', SOCKS5_USER_AUTH_VERSION)
 204             packet += self._len_and_data(username) + self._len_and_data(password)
 205             self.sendall(packet)
 206
 207             version, status = self._recv_bytes(2)
 208
 209             self._check_response_version(SOCKS5_USER_AUTH_VERSION, version)
 210
 211             if status != SOCKS5_USER_AUTH_SUCCESS:
 212                 self.close()
 213                 raise Socks5Error(Socks5Error.ERR_GENERAL_FAILURE)
 214
 215     def _setup_socks5(self, address):
 216         destaddr, port = address
 217
 218         ipaddr = self._resolve_address(destaddr, None, use_remote_dns=True)
 219
 220         self._socks5_auth()
 221
 222         reserved = 0
 223         packet = compat_struct_pack('!BBB', SOCKS5_VERSION, Socks5Command.CMD_CONNECT, reserved)
 224         if ipaddr is None:
 225             destaddr = destaddr.encode('utf-8')
 226             packet += compat_struct_pack('!B', Socks5AddressType.ATYP_DOMAINNAME)
 227             packet += self._len_and_data(destaddr)
 228         else:
 229             packet += compat_struct_pack('!B', Socks5AddressType.ATYP_IPV4) + ipaddr
 230         packet += compat_struct_pack('!H', port)
 231
 232         self.sendall(packet)
 233
 234         version, status, reserved, atype = self._recv_bytes(4)
 235
 236         self._check_response_version(SOCKS5_VERSION, version)
 237
 238         if status != Socks5Error.ERR_SUCCESS:
 239             self.close()
 240             raise Socks5Error(status)
 241
 242         if atype == Socks5AddressType.ATYP_IPV4:
 243             destaddr = self.recvall(4)
 244         elif atype == Socks5AddressType.ATYP_DOMAINNAME:
 245             alen = compat_ord(self.recv(1))
 246             destaddr = self.recvall(alen)
 247         elif atype == Socks5AddressType.ATYP_IPV6:
 248             destaddr = self.recvall(16)
 249         destport = compat_struct_unpack('!H', self.recvall(2))[0]
 250
 251         return (destaddr, destport)
 252
 253     def _make_proxy(self, connect_func, address):
 254         if not self._proxy:
 255             return connect_func(self, address)
 256
 257         result = connect_func(self, (self._proxy.host, self._proxy.port))
 258         if result != 0 and result is not None:
 259             return result
 260         setup_funcs = {
 261             ProxyType.SOCKS4: self._setup_socks4,
 262             ProxyType.SOCKS4A: self._setup_socks4a,
 263             ProxyType.SOCKS5: self._setup_socks5,
 264         }
 265         setup_funcs[self._proxy.type](address)
 266         return result
 267
 268     def connect(self, address):
 269         self._make_proxy(socket.socket.connect, address)
 270
 271     def connect_ex(self, address):
 272         return self._make_proxy(socket.socket.connect_ex, address)