Coverage for utilities/socks.py: 53%

1import logging 

2from urllib.parse import urlparse 

3 

4from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, PoolManager 

5from urllib3.connection import HTTPConnection, HTTPSConnection 

6 

7from .constants import HTTP_PROXY_SOCK_RDNS_SCHEMAS 

8 

9logger = logging.getLogger('netbox.utilities') 

10 

11 

12class ProxyHTTPConnection(HTTPConnection): 

13 """ 

14 A Proxy connection class that uses a SOCK proxy - used to create 

15 a urllib3 PoolManager that routes connections via the proxy. 

16 This is for an HTTP (not HTTPS) connection 

17 """ 

18 use_rdns = False 

19 

20 def __init__(self, *args, **kwargs): 

21 socks_options = kwargs.pop('_socks_options') 

22 self._proxy_url = socks_options['proxy_url'] 

23 super().__init__(*args, **kwargs) 

24 

25 def _new_conn(self): 

26 try: 

27 from python_socks.sync import Proxy 

28 except ModuleNotFoundError as e: 

29 logger.info( 

30 "Configuring an HTTP proxy using SOCKS requires the python_socks library. Check that it has been " 

31 "installed." 

32 ) 

33 raise e 

34 

35 proxy = Proxy.from_url(self._proxy_url, rdns=self.use_rdns) 

36 return proxy.connect( 

37 dest_host=self.host, 

38 dest_port=self.port, 

39 timeout=self.timeout 

40 ) 

41 

42 

43class ProxyHTTPSConnection(ProxyHTTPConnection, HTTPSConnection): 

44 """ 

45 A Proxy connection class for an HTTPS (not HTTP) connection. 

46 """ 

47 pass 

48 

49 

50class RdnsProxyHTTPConnection(ProxyHTTPConnection): 

51 """ 

52 A Proxy connection class for an HTTP remote-dns connection. 

53 I.E. socks4a, socks4h, socks5a, socks5h 

54 """ 

55 use_rdns = True 

56 

57 

58class RdnsProxyHTTPSConnection(ProxyHTTPSConnection): 

59 """ 

60 A Proxy connection class for an HTTPS remote-dns connection. 

61 I.E. socks4a, socks4h, socks5a, socks5h 

62 """ 

63 use_rdns = True 

64 

65 

66class ProxyHTTPConnectionPool(HTTPConnectionPool): 

67 ConnectionCls = ProxyHTTPConnection 

68 

69 

70class ProxyHTTPSConnectionPool(HTTPSConnectionPool): 

71 ConnectionCls = ProxyHTTPSConnection 

72 

73 

74class RdnsProxyHTTPConnectionPool(HTTPConnectionPool): 

75 ConnectionCls = RdnsProxyHTTPConnection 

76 

77 

78class RdnsProxyHTTPSConnectionPool(HTTPSConnectionPool): 

79 ConnectionCls = RdnsProxyHTTPSConnection 

80 

81 

82class ProxyPoolManager(PoolManager): 

83 def __init__(self, proxy_url, timeout=5, num_pools=10, headers=None, **connection_pool_kw): 

84 # python_socks uses rdns param to denote remote DNS parsing and 

85 # doesn't accept the 'h' or 'a' in the proxy URL 

86 if use_rdns := urlparse(proxy_url).scheme in HTTP_PROXY_SOCK_RDNS_SCHEMAS: 

87 proxy_url = proxy_url.replace('socks5h:', 'socks5:').replace('socks5a:', 'socks5:') 

88 proxy_url = proxy_url.replace('socks4h:', 'socks4:').replace('socks4a:', 'socks4:') 

89 

90 connection_pool_kw['_socks_options'] = {'proxy_url': proxy_url} 

91 connection_pool_kw['timeout'] = timeout 

92 

93 super().__init__(num_pools, headers, **connection_pool_kw) 

94 

95 if use_rdns: 

96 self.pool_classes_by_scheme = { 

97 'http': RdnsProxyHTTPConnectionPool, 

98 'https': RdnsProxyHTTPSConnectionPool, 

99 } 

100 else: 

101 self.pool_classes_by_scheme = { 

102 'http': ProxyHTTPConnectionPool, 

103 'https': ProxyHTTPSConnectionPool, 

104 }