_ Git - youtube-dl/blob - youtube_dl/socks.py

   1 # This is free and unencumbered software released into the public domain.
   2 #
   3 # Anyone is free to copy, modify, publish, use, compile, sell, or
   4 # distribute this software, either in source code form or as a compiled
   5 # binary, for any purpose, commercial or non-commercial, and by any
   6 # means.
   7 #
   8 # In jurisdictions that recognize copyright laws, the author or authors
   9 # of this software dedicate any and all copyright interest in the
  10 # software to the public domain. We make this dedication for the benefit
  11 # of the public at large and to the detriment of our heirs and
  12 # successors. We intend this dedication to be an overt act of
  13 # relinquishment in perpetuity of all present and future rights to this
  14 # software under copyright law.
  15 #
  16 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  17 # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  18 # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
  19 # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
  20 # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
  21 # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  22 # OTHER DEALINGS IN THE SOFTWARE.
  23 #
  24 # For more information, please refer to <http://unlicense.org/>
  25 #
  26 # Example:
  27 # import socks
  28 # import ftplib
  29 # import socket
  30 #
  31 # socks.patch_socket()
  32 #
  33 # f = ftplib.FTP('ftp.kernel.org')
  34 # f.login()
  35 # print f.retrlines('LIST')
  36 # f.quit()
  37 #
  38 # s = socket.create_connection(('www.google.com', 80))
  39 # s.sendall('HEAD / HTTP/1.0\r\n\r\n')
  40 # print s.recv(1024)
  41 # s.close()
  42 from __future__ import unicode_literals
  43 import os
  44 import struct
  45 import socket
  46 import time
  47
  48 __author__ = 'Timo Schmid <coding@timoschmid.de>'
  49
  50 _orig_socket = socket.socket
  51
  52 try:
  53     from collections import namedtuple
  54 except ImportError:
  55     from Collections import namedtuple
  56
  57 try:
  58     from urllib.parse import urlparse
  59 except:
  60     from urlparse import urlparse
  61
  62 try:
  63     from enum import Enum
  64 except ImportError:
  65     Enum = object
  66
  67
  68 class ProxyError(IOError): pass
  69 class Socks4Error(ProxyError):
  70     CODES = {
  71         0x5B: 'request rejected or failed',
  72         0x5C: 'request rejected becasue SOCKS server cannot connect to identd on the client',
  73         0x5D: 'request rejected because the client program and identd report different user-ids'
  74     }
  75     def __init__(self, code=None, msg=None):
  76         if code is not None and msg is None:
  77             msg = self.CODES.get(code)
  78             if msg is None:
  79                 msg = 'unknown error'
  80         super(Socks4Error, self).__init__(code, msg)
  81
  82 class Socks5Error(Socks4Error):
  83     CODES = {
  84         0x01: 'general SOCKS server failure',
  85         0x02: 'connection not allowed by ruleset',
  86         0x03: 'Network unreachable',
  87         0x04: 'Host unreachable',
  88         0x05: 'Connection refused',
  89         0x06: 'TTL expired',
  90         0x07: 'Command not supported',
  91         0x08: 'Address type not supported',
  92         0xFE: 'unknown username or invalid password',
  93         0xFF: 'all offered authentication methods were rejected'
  94     }
  95
  96 class ProxyType(Enum):
  97     SOCKS4  = 0
  98     SOCKS4A = 1
  99     SOCKS5  = 2
 100
 101 Proxy = namedtuple('Proxy', ('type', 'host', 'port', 'username', 'password', 'remote_dns'))
 102
 103 _default_proxy = None
 104
 105 def setdefaultproxy(proxytype=None, addr=None, port=None, rdns=True, username=None, password=None, allow_env_override=True):
 106     global _default_proxy
 107     if allow_env_override:
 108         all_proxy = os.environ.get('ALL_PROXY', os.environ.get('all_proxy'))
 109         if all_proxy:
 110             all_proxy = urlparse(all_proxy)
 111             if all_proxy.scheme.startswith('socks'):
 112                 if all_proxy.scheme == 'socks' or all_proxy.scheme == 'socks4':
 113                     proxytype = ProxyType.SOCKS4
 114                 elif all_proxy.scheme == 'socks4a':
 115                     proxytype = ProxyType.SOCKS4A
 116                 elif all_proxy.scheme == 'socks5':
 117                     proxytype = ProxyType.SOCKS5
 118                 addr = all_proxy.hostname
 119                 port = all_proxy.port
 120                 username = all_proxy.username
 121                 password = all_proxy.password
 122
 123     if proxytype is not None:
 124         _default_proxy = Proxy(proxytype, addr, port, username, password, rdns)
 125
 126
 127 def wrap_socket(sock):
 128     return socksocket(_sock=sock._sock)
 129
 130 def wrap_module(module):
 131     if hasattr(module, 'socket'):
 132         sock = module.socket
 133         if isinstance(sock, socket.socket):
 134             module.socket = sockssocket
 135         elif hasattr(socket, 'socket'):
 136             socket.socket = sockssocket
 137
 138 def patch_socket():
 139     import sys
 140     if 'socket' not in sys.modules:
 141         import socket
 142     sys.modules['socket'].socket = sockssocket
 143
 144
 145 class sockssocket(socket.socket):
 146     def __init__(self, *args, **kwargs):
 147         self.__proxy = None
 148         if 'proxy' in kwargs:
 149             self.__proxy = kwargs['proxy']
 150             del kwargs['proxy']
 151         super(sockssocket, self).__init__(*args, **kwargs)
 152
 153     @property
 154     def _proxy(self):
 155         if self.__proxy:
 156             return self.__proxy
 157         return _default_proxy
 158
 159     @property
 160     def _proxy_port(self):
 161         if self._proxy:
 162             if self._proxy.port:
 163                 return self._proxy.port
 164             return 1080
 165         return None
 166
 167     def setproxy(self, proxytype=None, addr=None, port=None, rdns=True, username=None, password=None):
 168         if proxytype is None:
 169             self.__proxy = None
 170         else:
 171             self.__proxy = Proxy(proxytype, addr, port, username, password, rdns)
 172
 173     def recvall(self, cnt):
 174         data = b''
 175         while len(data) < cnt:
 176             cur = self.recv(cnt - len(data))
 177             if not cur:
 178                 raise IOError("{0} bytes missing".format(cnt-len(data)))
 179             data += cur
 180         return data
 181
 182     def _setup_socks4(self, address, is_4a=False):
 183         destaddr, port = address
 184
 185         try:
 186             ipaddr = socket.inet_aton(destaddr)
 187         except socket.error:
 188             if is_4a and self._proxy.remote_dns:
 189                 ipaddr = struct.pack('!BBBB', 0, 0, 0, 0xFF)
 190             else:
 191                 ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
 192
 193         packet = struct.pack('!BBH', 0x4, 0x1, port) + ipaddr
 194         if self._proxy.username:
 195             username = self._proxy.username
 196             if hasattr(username, 'encode'):
 197                 username = username.encode()
 198             packet += struct.pack('!{0}s'.format(len(username)+1), username)
 199         else:
 200             packet += b'\x00'
 201
 202         if is_4a and self._proxy.remote_dns:
 203             if hasattr(destaddr, 'encode'):
 204                 destaddr = destaddr.encode()
 205             packet += struct.pack('!{0}s'.format(len(destaddr)+1), destaddr)
 206
 207         self.sendall(packet)
 208
 209         packet = self.recvall(8)
 210         nbyte, resp_code, dstport, dsthost = struct.unpack('!BBHI', packet)
 211
 212         # check valid response
 213         if nbyte != 0x00:
 214             self.close()
 215             raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(0, nbyte))
 216
 217         # access granted
 218         if resp_code != 0x5a:
 219             self.close()
 220             raise Socks4Error(resp_code)
 221
 222     def _setup_socks5(self, address):
 223         destaddr, port = address
 224
 225         try:
 226             ipaddr = socket.inet_aton(destaddr)
 227         except socket.error:
 228             if self._proxy.remote_dns:
 229                 ipaddr = None
 230             else:
 231                 ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
 232
 233         auth_methods = 1
 234         if self._proxy.username and self._proxy.password:
 235             # two auth methods available
 236             auth_methods = 2
 237         packet = struct.pack('!BBB', 0x5, auth_methods, 0x00) # no auth
 238         if self._proxy.username and self._proxy.password:
 239             packet += struct.pack('!B', 0x02) # user/pass auth
 240
 241         self.sendall(packet)
 242
 243         packet = self.recvall(2)
 244         version, method = struct.unpack('!BB', packet)
 245
 246         # check valid response
 247         if version != 0x05:
 248             self.close()
 249             raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(5, version))
 250
 251         # no auth methods
 252         if method == 0xFF:
 253             self.close()
 254             raise Socks5Error(method)
 255
 256         # user/pass auth
 257         if method == 0x01:
 258             username = self._proxy.username
 259             if hasattr(username, 'encode'):
 260                 username = username.encode()
 261             password = self._proxy.password
 262             if hasattr(password, 'encode'):
 263                 password = password.encode()
 264             packet = struct.pack('!BB', 1, len(username)) + username
 265             packet += struct.pack('!B', len(password)) + password
 266             self.sendall(packet)
 267
 268             packet = self.recvall(2)
 269             version, status = struct.unpack('!BB', packet)
 270
 271             if version != 0x01:
 272                 self.close()
 273                 raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(1, version))
 274
 275             if status != 0x00:
 276                 self.close()
 277                 raise Socks5Error(1)
 278         elif method == 0x00: # no auth
 279             pass
 280
 281
 282         packet = struct.pack('!BBB', 5, 1, 0)
 283         if ipaddr is None:
 284             if hasattr(destaddr, 'encode'):
 285                 destaddr = destaddr.encode()
 286             packet += struct.pack('!BB', 3, len(destaddr)) + destaddr
 287         else:
 288             packet += struct.pack('!B', 1) + ipaddr
 289         packet += struct.pack('!H', port)
 290
 291         self.sendall(packet)
 292
 293         packet = self.recvall(4)
 294         version, status, _, atype = struct.unpack('!BBBB', packet)
 295
 296         if version != 0x05:
 297             self.close()
 298             raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(5, version))
 299
 300         if status != 0x00:
 301             self.close()
 302             raise Socks5Error(status)
 303
 304         if atype == 0x01:
 305             destaddr = self.recvall(4)
 306         elif atype == 0x03:
 307             alen = struct.unpack('!B', self.recv(1))[0]
 308             destaddr = self.recvall(alen)
 309         elif atype == 0x04:
 310             destaddr = self.recvall(16)
 311         destport = struct.unpack('!H', self.recvall(2))[0]
 312
 313     def _make_proxy(self, connect_func, address):
 314         if self._proxy.type == ProxyType.SOCKS4:
 315             result = connect_func(self, (self._proxy.host, self._proxy_port))
 316             if result != 0 and result is not None:
 317                 return result
 318             self._setup_socks4(address)
 319         elif self._proxy.type == ProxyType.SOCKS4A:
 320             result = connect_func(self, (self._proxy.host, self._proxy_port))
 321             if result != 0 and result is not None:
 322                 return result
 323             self._setup_socks4(address, is_4a=True)
 324         elif self._proxy.type == ProxyType.SOCKS5:
 325             result = connect_func(self, (self._proxy.host, self._proxy_port))
 326             if result != 0 and result is not None:
 327                 return result
 328             self._setup_socks5(address)
 329         else:
 330             return connect_func(self, address)
 331
 332     def connect(self, address):
 333         self._make_proxy(_orig_socket.connect, address)
 334
 335     def connect_ex(self, address):
 336         return self._make_proxy(_orig_socket.connect_ex, address)