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(socket.error):
  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):
 138         try:
 139             return socket.inet_aton(destaddr)
 140         except OSError:
 141             if use_remote_dns and self._proxy.remote_dns:
 142                 return default
 143             else:
 144                 return socket.inet_aton(socket.gethostbyname(destaddr))
 145
 146     def _setup_socks4(self, address, is_4a=False):
 147         destaddr, port = address
 148
 149         ipaddr = self._resolve_address(destaddr, SOCKS4_DEFAULT_DSTIP, use_remote_dns=is_4a)
 150
 151         packet = struct.pack('!BBH', SOCKS4_VERSION, Socks4Command.CMD_CONNECT, port) + ipaddr
 152
 153         username = (self._proxy.username or '').encode()
 154         packet += username + b'\x00'
 155
 156         if is_4a and self._proxy.remote_dns:
 157             packet += destaddr.encode() + b'\x00'
 158
 159         self.sendall(packet)
 160
 161         version, resp_code, dstport, dsthost = struct.unpack('!BBHI', self.recvall(8))
 162
 163         self._check_response_version(SOCKS4_REPLY_VERSION, version)
 164
 165         if resp_code != Socks4Error.ERR_SUCCESS:
 166             self.close()
 167             raise Socks4Error(resp_code)
 168
 169         return (dsthost, dstport)
 170
 171     def _setup_socks4a(self, address):
 172         self._setup_socks4(address, is_4a=True)
 173
 174     def _socks5_auth(self):
 175         packet = struct.pack('!B', SOCKS5_VERSION)
 176
 177         auth_methods = [Socks5Auth.AUTH_NONE]
 178         if self._proxy.username and self._proxy.password:
 179             auth_methods.append(Socks5Auth.AUTH_USER_PASS)
 180
 181         packet += struct.pack('!B', len(auth_methods))
 182         packet += struct.pack(f'!{len(auth_methods)}B', *auth_methods)
 183
 184         self.sendall(packet)
 185
 186         version, method = self._recv_bytes(2)
 187
 188         self._check_response_version(SOCKS5_VERSION, version)
 189
 190         if method == Socks5Auth.AUTH_NO_ACCEPTABLE or (
 191                 method == Socks5Auth.AUTH_USER_PASS and (not self._proxy.username or not self._proxy.password)):
 192             self.close()
 193             raise Socks5Error(Socks5Auth.AUTH_NO_ACCEPTABLE)
 194
 195         if method == Socks5Auth.AUTH_USER_PASS:
 196             username = self._proxy.username.encode()
 197             password = self._proxy.password.encode()
 198             packet = struct.pack('!B', SOCKS5_USER_AUTH_VERSION)
 199             packet += self._len_and_data(username) + self._len_and_data(password)
 200             self.sendall(packet)
 201
 202             version, status = self._recv_bytes(2)
 203
 204             self._check_response_version(SOCKS5_USER_AUTH_VERSION, version)
 205
 206             if status != SOCKS5_USER_AUTH_SUCCESS:
 207                 self.close()
 208                 raise Socks5Error(Socks5Error.ERR_GENERAL_FAILURE)
 209
 210     def _setup_socks5(self, address):
 211         destaddr, port = address
 212
 213         ipaddr = self._resolve_address(destaddr, None, use_remote_dns=True)
 214
 215         self._socks5_auth()
 216
 217         reserved = 0
 218         packet = struct.pack('!BBB', SOCKS5_VERSION, Socks5Command.CMD_CONNECT, reserved)
 219         if ipaddr is None:
 220             destaddr = destaddr.encode()
 221             packet += struct.pack('!B', Socks5AddressType.ATYP_DOMAINNAME)
 222             packet += self._len_and_data(destaddr)
 223         else:
 224             packet += struct.pack('!B', Socks5AddressType.ATYP_IPV4) + ipaddr
 225         packet += struct.pack('!H', port)
 226
 227         self.sendall(packet)
 228
 229         version, status, reserved, atype = self._recv_bytes(4)
 230
 231         self._check_response_version(SOCKS5_VERSION, version)
 232
 233         if status != Socks5Error.ERR_SUCCESS:
 234             self.close()
 235             raise Socks5Error(status)
 236
 237         if atype == Socks5AddressType.ATYP_IPV4:
 238             destaddr = self.recvall(4)
 239         elif atype == Socks5AddressType.ATYP_DOMAINNAME:
 240             alen = compat_ord(self.recv(1))
 241             destaddr = self.recvall(alen)
 242         elif atype == Socks5AddressType.ATYP_IPV6:
 243             destaddr = self.recvall(16)
 244         destport = struct.unpack('!H', self.recvall(2))[0]
 245
 246         return (destaddr, destport)
 247
 248     def _make_proxy(self, connect_func, address):
 249         if not self._proxy:
 250             return connect_func(self, address)
 251
 252         result = connect_func(self, (self._proxy.host, self._proxy.port))
 253         if result != 0 and result is not None:
 254             return result
 255         setup_funcs = {
 256             ProxyType.SOCKS4: self._setup_socks4,
 257             ProxyType.SOCKS4A: self._setup_socks4a,
 258             ProxyType.SOCKS5: self._setup_socks5,
 259         }
 260         setup_funcs[self._proxy.type](address)
 261         return result
 262
 263     def connect(self, address):
 264         self._make_proxy(socket.socket.connect, address)
 265
 266     def connect_ex(self, address):
 267         return self._make_proxy(socket.socket.connect_ex, address)