yt_dlp/socks.py

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