]> jfr.im git - yt-dlp.git/blobdiff - yt_dlp/socks.py
[ie/sproutvideo] Add extractors (#10098)
[yt-dlp.git] / yt_dlp / socks.py
index f93328f63a410797f858c1a8308a7894cf50a2c5..e553a5f1f353179e4cb79615af2d4fd21d528523 100644 (file)
@@ -49,7 +49,7 @@ class Socks5AddressType:
     ATYP_IPV6 = 0x04
 
 
-class ProxyError(socket.error):
+class ProxyError(OSError):
     ERR_SUCCESS = 0x00
 
     def __init__(self, code=None, msg=None):
@@ -60,8 +60,8 @@ def __init__(self, code=None, msg=None):
 
 class InvalidVersionError(ProxyError):
     def __init__(self, expected_version, got_version):
-        msg = ('Invalid response version from server. Expected {:02x} got '
-               '{:02x}'.format(expected_version, got_version))
+        msg = (f'Invalid response version from server. Expected {expected_version:02x} got '
+               f'{got_version:02x}')
         super().__init__(0, msg)
 
 
@@ -71,7 +71,7 @@ class Socks4Error(ProxyError):
     CODES = {
         91: 'request rejected or failed',
         92: 'request rejected because SOCKS server cannot connect to identd on the client',
-        93: 'request rejected because the client program and identd report different user-ids'
+        93: 'request rejected because the client program and identd report different user-ids',
     }
 
 
@@ -88,7 +88,7 @@ class Socks5Error(ProxyError):
         0x07: 'Command not supported',
         0x08: 'Address type not supported',
         0xFE: 'unknown username or invalid password',
-        0xFF: 'all offered authentication methods were rejected'
+        0xFF: 'all offered authentication methods were rejected',
     }
 
 
@@ -134,26 +134,31 @@ def _check_response_version(self, expected_version, got_version):
             self.close()
             raise InvalidVersionError(expected_version, got_version)
 
-    def _resolve_address(self, destaddr, default, use_remote_dns):
-        try:
-            return socket.inet_aton(destaddr)
-        except OSError:
-            if use_remote_dns and self._proxy.remote_dns:
-                return default
-            else:
-                return socket.inet_aton(socket.gethostbyname(destaddr))
+    def _resolve_address(self, destaddr, default, use_remote_dns, family=None):
+        for f in (family,) if family else (socket.AF_INET, socket.AF_INET6):
+            try:
+                return f, socket.inet_pton(f, destaddr)
+            except OSError:
+                continue
+
+        if use_remote_dns and self._proxy.remote_dns:
+            return 0, default
+        else:
+            res = socket.getaddrinfo(destaddr, None, family=family or 0)
+            f, _, _, _, ipaddr = res[0]
+            return f, socket.inet_pton(f, ipaddr[0])
 
     def _setup_socks4(self, address, is_4a=False):
         destaddr, port = address
 
-        ipaddr = self._resolve_address(destaddr, SOCKS4_DEFAULT_DSTIP, use_remote_dns=is_4a)
+        _, ipaddr = self._resolve_address(destaddr, SOCKS4_DEFAULT_DSTIP, use_remote_dns=is_4a, family=socket.AF_INET)
 
         packet = struct.pack('!BBH', SOCKS4_VERSION, Socks4Command.CMD_CONNECT, port) + ipaddr
 
         username = (self._proxy.username or '').encode()
         packet += username + b'\x00'
 
-        if is_4a and self._proxy.remote_dns:
+        if is_4a and self._proxy.remote_dns and ipaddr == SOCKS4_DEFAULT_DSTIP:
             packet += destaddr.encode() + b'\x00'
 
         self.sendall(packet)
@@ -210,7 +215,7 @@ def _socks5_auth(self):
     def _setup_socks5(self, address):
         destaddr, port = address
 
-        ipaddr = self._resolve_address(destaddr, None, use_remote_dns=True)
+        family, ipaddr = self._resolve_address(destaddr, None, use_remote_dns=True)
 
         self._socks5_auth()
 
@@ -220,8 +225,10 @@ def _setup_socks5(self, address):
             destaddr = destaddr.encode()
             packet += struct.pack('!B', Socks5AddressType.ATYP_DOMAINNAME)
             packet += self._len_and_data(destaddr)
-        else:
+        elif family == socket.AF_INET:
             packet += struct.pack('!B', Socks5AddressType.ATYP_IPV4) + ipaddr
+        elif family == socket.AF_INET6:
+            packet += struct.pack('!B', Socks5AddressType.ATYP_IPV6) + ipaddr
         packet += struct.pack('!H', port)
 
         self.sendall(packet)