_ Git - youtube-dl/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 import collections
   7 import socket
   8
   9 from .compat import (
  10     struct_pack,
  11     struct_unpack,
  12 )
  13
  14 __author__ = 'Timo Schmid <coding@timoschmid.de>'
  15
  16
  17 class ProxyError(IOError):
  18     pass
  19
  20
  21 class Socks4Error(ProxyError):
  22     CODES = {
  23         0x5B: 'request rejected or failed',
  24         0x5C: 'request rejected becasue SOCKS server cannot connect to identd on the client',
  25         0x5D: 'request rejected because the client program and identd report different user-ids'
  26     }
  27
  28     def __init__(self, code=None, msg=None):
  29         if code is not None and msg is None:
  30             msg = self.CODES.get(code)
  31             if msg is None:
  32                 msg = 'unknown error'
  33         super(Socks4Error, self).__init__(code, msg)
  34
  35
  36 class Socks5Error(Socks4Error):
  37     CODES = {
  38         0x01: 'general SOCKS server failure',
  39         0x02: 'connection not allowed by ruleset',
  40         0x03: 'Network unreachable',
  41         0x04: 'Host unreachable',
  42         0x05: 'Connection refused',
  43         0x06: 'TTL expired',
  44         0x07: 'Command not supported',
  45         0x08: 'Address type not supported',
  46         0xFE: 'unknown username or invalid password',
  47         0xFF: 'all offered authentication methods were rejected'
  48     }
  49
  50
  51 class ProxyType(object):
  52     SOCKS4 = 0
  53     SOCKS4A = 1
  54     SOCKS5 = 2
  55
  56 Proxy = collections.namedtuple('Proxy', ('type', 'host', 'port', 'username', 'password', 'remote_dns'))
  57
  58
  59 class sockssocket(socket.socket):
  60     @property
  61     def _proxy(self):
  62         return self.__proxy
  63
  64     @property
  65     def _proxy_port(self):
  66         if self._proxy:
  67             if self._proxy.port:
  68                 return self._proxy.port
  69             return 1080
  70         return None
  71
  72     def setproxy(self, proxytype=None, addr=None, port=None, rdns=True, username=None, password=None):
  73         if proxytype is None:
  74             self.__proxy = None
  75         else:
  76             self.__proxy = Proxy(proxytype, addr, port, username, password, rdns)
  77
  78     def recvall(self, cnt):
  79         data = b''
  80         while len(data) < cnt:
  81             cur = self.recv(cnt - len(data))
  82             if not cur:
  83                 raise IOError('{0} bytes missing'.format(cnt - len(data)))
  84             data += cur
  85         return data
  86
  87     def _setup_socks4(self, address, is_4a=False):
  88         destaddr, port = address
  89
  90         try:
  91             ipaddr = socket.inet_aton(destaddr)
  92         except socket.error:
  93             if is_4a and self._proxy.remote_dns:
  94                 ipaddr = struct_pack('!BBBB', 0, 0, 0, 0xFF)
  95             else:
  96                 ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
  97
  98         packet = struct_pack('!BBH', 0x4, 0x1, port) + ipaddr
  99         if self._proxy.username:
 100             username = self._proxy.username
 101             if hasattr(username, 'encode'):
 102                 username = username.encode()
 103             packet += struct_pack('!{0}s'.format(len(username) + 1), username)
 104         else:
 105             packet += b'\x00'
 106
 107         if is_4a and self._proxy.remote_dns:
 108             if hasattr(destaddr, 'encode'):
 109                 destaddr = destaddr.encode()
 110             packet += struct_pack('!{0}s'.format(len(destaddr) + 1), destaddr)
 111
 112         self.sendall(packet)
 113
 114         packet = self.recvall(8)
 115         nbyte, resp_code, dstport, dsthost = struct_unpack('!BBHI', packet)
 116
 117         # check valid response
 118         if nbyte != 0x00:
 119             self.close()
 120             raise ProxyError(
 121                 0, 'Invalid response from server. Expected {0:02x} got {1:02x}'.format(0, nbyte))
 122
 123         # access granted
 124         if resp_code != 0x5a:
 125             self.close()
 126             raise Socks4Error(resp_code)
 127
 128         return (dsthost, dstport)
 129
 130     def _setup_socks5(self, address):
 131         destaddr, port = address
 132
 133         try:
 134             ipaddr = socket.inet_aton(destaddr)
 135         except socket.error:
 136             if self._proxy.remote_dns:
 137                 ipaddr = None
 138             else:
 139                 ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
 140
 141         auth_methods = 1
 142         if self._proxy.username and self._proxy.password:
 143             # two auth methods available
 144             auth_methods = 2
 145         packet = struct_pack('!BBB', 0x5, auth_methods, 0x00)  # no auth
 146         if self._proxy.username and self._proxy.password:
 147             packet += struct_pack('!B', 0x02)  # user/pass auth
 148
 149         self.sendall(packet)
 150
 151         packet = self.recvall(2)
 152         version, method = struct_unpack('!BB', packet)
 153
 154         # check valid response
 155         if version != 0x05:
 156             self.close()
 157             raise ProxyError(
 158                 0, 'Invalid response from server. Expected {0:02x} got {1:02x}'.format(5, version))
 159
 160         # no auth methods
 161         if method == 0xFF:
 162             self.close()
 163             raise Socks5Error(method)
 164
 165         # user/pass auth
 166         if method == 0x01:
 167             username = self._proxy.username
 168             if hasattr(username, 'encode'):
 169                 username = username.encode()
 170             password = self._proxy.password
 171             if hasattr(password, 'encode'):
 172                 password = password.encode()
 173             packet = struct_pack('!BB', 1, len(username)) + username
 174             packet += struct_pack('!B', len(password)) + password
 175             self.sendall(packet)
 176
 177             packet = self.recvall(2)
 178             version, status = struct_unpack('!BB', packet)
 179
 180             if version != 0x01:
 181                 self.close()
 182                 raise ProxyError(
 183                     0, 'Invalid response from server. Expected {0:02x} got {1:02x}'.format(1, version))
 184
 185             if status != 0x00:
 186                 self.close()
 187                 raise Socks5Error(1)
 188         elif method == 0x00:  # no auth
 189             pass
 190
 191         packet = struct_pack('!BBB', 5, 1, 0)
 192         if ipaddr is None:
 193             if hasattr(destaddr, 'encode'):
 194                 destaddr = destaddr.encode()
 195             packet += struct_pack('!BB', 3, len(destaddr)) + destaddr
 196         else:
 197             packet += struct_pack('!B', 1) + ipaddr
 198         packet += struct_pack('!H', port)
 199
 200         self.sendall(packet)
 201
 202         packet = self.recvall(4)
 203         version, status, _, atype = struct_unpack('!BBBB', packet)
 204
 205         if version != 0x05:
 206             self.close()
 207             raise ProxyError(
 208                 0, 'Invalid response from server. Expected {0:02x} got {1:02x}'.format(5, version))
 209
 210         if status != 0x00:
 211             self.close()
 212             raise Socks5Error(status)
 213
 214         if atype == 0x01:
 215             destaddr = self.recvall(4)
 216         elif atype == 0x03:
 217             alen = struct_unpack('!B', self.recv(1))[0]
 218             destaddr = self.recvall(alen)
 219         elif atype == 0x04:
 220             destaddr = self.recvall(16)
 221         destport = struct_unpack('!H', self.recvall(2))[0]
 222
 223         return (destaddr, destport)
 224
 225     def _make_proxy(self, connect_func, address):
 226         if self._proxy.type == ProxyType.SOCKS4:
 227             result = connect_func(self, (self._proxy.host, self._proxy_port))
 228             if result != 0 and result is not None:
 229                 return result
 230             self._setup_socks4(address)
 231         elif self._proxy.type == ProxyType.SOCKS4A:
 232             result = connect_func(self, (self._proxy.host, self._proxy_port))
 233             if result != 0 and result is not None:
 234                 return result
 235             self._setup_socks4(address, is_4a=True)
 236         elif self._proxy.type == ProxyType.SOCKS5:
 237             result = connect_func(self, (self._proxy.host, self._proxy_port))
 238             if result != 0 and result is not None:
 239                 return result
 240             self._setup_socks5(address)
 241         else:
 242             return connect_func(self, address)
 243
 244     def connect(self, address):
 245         self._make_proxy(socket.socket.connect, address)
 246
 247     def connect_ex(self, address):
 248         return self._make_proxy(socket.socket.connect_ex, address)