From c4915ae521c7a167f6c23a0c2338df938af2335f Mon Sep 17 00:00:00 2001 From: Arthur Hanson Date: Mon, 30 Sep 2024 07:28:31 -0700 Subject: [PATCH] refactor Socks to utilities --- netbox/core/data_backends.py | 76 +----------------------------------- netbox/utilities/socks.py | 75 +++++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 75 deletions(-) create mode 100644 netbox/utilities/socks.py diff --git a/netbox/core/data_backends.py b/netbox/core/data_backends.py index 935cb3aa9..72c7bae53 100644 --- a/netbox/core/data_backends.py +++ b/netbox/core/data_backends.py @@ -4,8 +4,6 @@ import re import tempfile from contextlib import contextmanager from pathlib import Path -from urllib3 import PoolManager, HTTPConnectionPool, HTTPSConnectionPool -from urllib3.connection import HTTPConnection, HTTPSConnection from urllib.parse import urlparse from django import forms @@ -14,6 +12,7 @@ from django.utils.translation import gettext as _ from netbox.data_backends import DataBackend from netbox.utils import register_data_backend +from utilities.socks import ProxyPoolManager from .exceptions import SyncError __all__ = ( @@ -25,79 +24,6 @@ __all__ = ( logger = logging.getLogger('netbox.data_backends') -# These Proxy Methods are for handling SOCKS connection proxy -class ProxyHTTPConnection(HTTPConnection): - use_rdns = False - - def __init__(self, *args, **kwargs): - socks_options = kwargs.pop('_socks_options') - self._proxy_url = socks_options['proxy_url'] - super().__init__(*args, **kwargs) - - def _new_conn(self): - from python_socks.sync import Proxy - proxy = Proxy.from_url(self._proxy_url, rdns=self.use_rdns) - return proxy.connect( - dest_host=self.host, - dest_port=self.port, - timeout=self.timeout - ) - - -class ProxyHTTPSConnection(ProxyHTTPConnection, HTTPSConnection): - pass - - -class rdnsProxyHTTPConnection(ProxyHTTPConnection): - use_rdns = True - - -class rdnsProxyHTTPSConnection(ProxyHTTPSConnection): - use_rdns = True - - -class ProxyHTTPConnectionPool(HTTPConnectionPool): - ConnectionCls = ProxyHTTPConnection - - -class ProxyHTTPSConnectionPool(HTTPSConnectionPool): - ConnectionCls = ProxyHTTPSConnection - - -class rdnsProxyHTTPConnectionPool(HTTPConnectionPool): - ConnectionCls = rdnsProxyHTTPConnection - - -class rdnsProxyHTTPSConnectionPool(HTTPSConnectionPool): - ConnectionCls = rdnsProxyHTTPSConnection - - -class ProxyPoolManager(PoolManager): - def __init__(self, proxy_url, timeout=5, num_pools=10, headers=None, **connection_pool_kw): - # python_socks uses rdns param to denote remote DNS parsing and - # doesn't accept the 'h' or 'a' in the proxy URL - cleaned_proxy_url = proxy_url - if use_rdns := urlparse(cleaned_proxy_url).scheme in ['socks4h', 'socks4a' 'socks5h', 'socks5a']: - cleaned_proxy_url = cleaned_proxy_url.replace('socks5h:', 'socks5:').replace('socks5a:', 'socks5:') - cleaned_proxy_url = cleaned_proxy_url.replace('socks4h:', 'socks4:').replace('socks4a:', 'socks4:') - - connection_pool_kw['_socks_options'] = {'proxy_url': cleaned_proxy_url} - connection_pool_kw['timeout'] = timeout - - super().__init__(num_pools, headers, **connection_pool_kw) - - if use_rdns: - self.pool_classes_by_scheme = { - 'http': rdnsProxyHTTPConnectionPool, - 'https': rdnsProxyHTTPSConnectionPool, - } - else: - self.pool_classes_by_scheme = { - 'http': ProxyHTTPConnectionPool, - 'https': ProxyHTTPSConnectionPool, - } - - @register_data_backend() class LocalBackend(DataBackend): name = 'local' diff --git a/netbox/utilities/socks.py b/netbox/utilities/socks.py new file mode 100644 index 000000000..30cc782bc --- /dev/null +++ b/netbox/utilities/socks.py @@ -0,0 +1,75 @@ +from urllib3 import PoolManager, HTTPConnectionPool, HTTPSConnectionPool +from urllib3.connection import HTTPConnection, HTTPSConnection + + +# These Proxy Methods are for handling SOCKS connection proxy +class ProxyHTTPConnection(HTTPConnection): + use_rdns = False + + def __init__(self, *args, **kwargs): + socks_options = kwargs.pop('_socks_options') + self._proxy_url = socks_options['proxy_url'] + super().__init__(*args, **kwargs) + + def _new_conn(self): + from python_socks.sync import Proxy + proxy = Proxy.from_url(self._proxy_url, rdns=self.use_rdns) + return proxy.connect( + dest_host=self.host, + dest_port=self.port, + timeout=self.timeout + ) + + +class ProxyHTTPSConnection(ProxyHTTPConnection, HTTPSConnection): + pass + + +class RdnsProxyHTTPConnection(ProxyHTTPConnection): + use_rdns = True + + +class RdnsProxyHTTPSConnection(ProxyHTTPSConnection): + use_rdns = True + + +class ProxyHTTPConnectionPool(HTTPConnectionPool): + ConnectionCls = ProxyHTTPConnection + + +class ProxyHTTPSConnectionPool(HTTPSConnectionPool): + ConnectionCls = ProxyHTTPSConnection + + +class RdnsProxyHTTPConnectionPool(HTTPConnectionPool): + ConnectionCls = RdnsProxyHTTPConnection + + +class RdnsProxyHTTPSConnectionPool(HTTPSConnectionPool): + ConnectionCls = RdnsProxyHTTPSConnection + + +class ProxyPoolManager(PoolManager): + def __init__(self, proxy_url, timeout=5, num_pools=10, headers=None, **connection_pool_kw): + # python_socks uses rdns param to denote remote DNS parsing and + # doesn't accept the 'h' or 'a' in the proxy URL + cleaned_proxy_url = proxy_url + if use_rdns := urlparse(cleaned_proxy_url).scheme in ['socks4h', 'socks4a' 'socks5h', 'socks5a']: + cleaned_proxy_url = cleaned_proxy_url.replace('socks5h:', 'socks5:').replace('socks5a:', 'socks5:') + cleaned_proxy_url = cleaned_proxy_url.replace('socks4h:', 'socks4:').replace('socks4a:', 'socks4:') + + connection_pool_kw['_socks_options'] = {'proxy_url': cleaned_proxy_url} + connection_pool_kw['timeout'] = timeout + + super().__init__(num_pools, headers, **connection_pool_kw) + + if use_rdns: + self.pool_classes_by_scheme = { + 'http': RdnsProxyHTTPConnectionPool, + 'https': RdnsProxyHTTPSConnectionPool, + } + else: + self.pool_classes_by_scheme = { + 'http': ProxyHTTPConnectionPool, + 'https': ProxyHTTPSConnectionPool, + }