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

Generated by cgit