From 0a5de27335d8a39ec438ca2b24fd7e131f1eb431 Mon Sep 17 00:00:00 2001
From: Arthur Hanson <worldnomad@gmail.com>
Date: Mon, 30 Sep 2024 11:46:35 -0700
Subject: [PATCH] review feedback

---
 netbox/core/data_backends.py  | 16 +++++++++++-----
 netbox/utilities/constants.py |  4 ++++
 netbox/utilities/socks.py     | 27 +++++++++++++++++++++------
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/netbox/core/data_backends.py b/netbox/core/data_backends.py
index 72c7bae53..f20ac017c 100644
--- a/netbox/core/data_backends.py
+++ b/netbox/core/data_backends.py
@@ -8,10 +8,12 @@ from urllib.parse import urlparse
 
 from django import forms
 from django.conf import settings
+from django.core.exceptions import ImproperlyConfigured
 from django.utils.translation import gettext as _
 
 from netbox.data_backends import DataBackend
 from netbox.utils import register_data_backend
+from utilities.constants import DATA_SOURCE_SUPPORTED_SCHEMAS, DATA_SOURCE_SUPPORTED_SOCK_SCHEMAS
 from utilities.socks import ProxyPoolManager
 from .exceptions import SyncError
 
@@ -71,11 +73,15 @@ class GitBackend(DataBackend):
         self.use_socks = False
 
         # Apply HTTP proxy (if configured)
-        if settings.HTTP_PROXIES and self.url_scheme in ('http', 'https'):
-            if proxy := settings.HTTP_PROXIES.get(self.url_scheme):
-                config.set("http", "proxy", proxy)
-                if urlparse(proxy).scheme in ['socks4', 'socks4a', 'socks4h', 'socks5', 'socks5a', 'socks5h']:
-                    self.use_socks = True
+        if settings.HTTP_PROXIES:
+            if proxy := settings.HTTP_PROXIES.get(self.url_scheme, None):
+                if urlparse(proxy).scheme not in DATA_SOURCE_SUPPORTED_SCHEMAS:
+                    raise ImproperlyConfigured(f"Unsupported Git DataSource proxy scheme: {urlparse(proxy).scheme}")
+
+                if self.url_scheme in ('http', 'https'):
+                    config.set("http", "proxy", proxy)
+                    if urlparse(proxy).scheme in DATA_SOURCE_SUPPORTED_SOCK_SCHEMAS:
+                        self.use_socks = True
 
         return config
 
diff --git a/netbox/utilities/constants.py b/netbox/utilities/constants.py
index c7c26f6b3..3ec6fd2dd 100644
--- a/netbox/utilities/constants.py
+++ b/netbox/utilities/constants.py
@@ -93,3 +93,7 @@ HTML_ALLOWED_ATTRIBUTES = {
     "td": {"align"},
     "th": {"align"},
 }
+
+DATA_SOURCE_SUPPORTED_SOCK_SCHEMAS = ['socks4', 'socks4a', 'socks4h', 'socks5', 'socks5a', 'socks5h']
+DATA_SOURCE_SOCK_RDNS_SCHEMAS = ['socks4h', 'socks4a', 'socks5h', 'socks5a']
+DATA_SOURCE_SUPPORTED_SCHEMAS = ['http', 'https', 'socks4', 'socks4a', 'socks4h', 'socks5', 'socks5a', 'socks5h']
diff --git a/netbox/utilities/socks.py b/netbox/utilities/socks.py
index 0b5864e44..8df82fb7a 100644
--- a/netbox/utilities/socks.py
+++ b/netbox/utilities/socks.py
@@ -1,10 +1,15 @@
 from urllib.parse import urlparse
 from urllib3 import PoolManager, HTTPConnectionPool, HTTPSConnectionPool
 from urllib3.connection import HTTPConnection, HTTPSConnection
+from .constants import DATA_SOURCE_SOCK_RDNS_SCHEMAS
 
 
-# These Proxy Methods are for handling SOCKS connection proxy
 class ProxyHTTPConnection(HTTPConnection):
+    """
+    A Proxy connection class that uses a SOCK proxy - used to create
+    a urllib3 PoolManager that routes connections via the proxy.
+    This is for an HTTP (not HTTPS) connection
+    """
     use_rdns = False
 
     def __init__(self, *args, **kwargs):
@@ -23,14 +28,25 @@ class ProxyHTTPConnection(HTTPConnection):
 
 
 class ProxyHTTPSConnection(ProxyHTTPConnection, HTTPSConnection):
+    """
+    A Proxy connection class for an HTTPS (not HTTP) connection.
+    """
     pass
 
 
 class RdnsProxyHTTPConnection(ProxyHTTPConnection):
+    """
+    A Proxy connection class for an HTTP remote-dns connection.
+    I.E. socks4a, socks4h, socks5a, socks5h
+    """
     use_rdns = True
 
 
 class RdnsProxyHTTPSConnection(ProxyHTTPSConnection):
+    """
+    A Proxy connection class for an HTTPS remote-dns connection.
+    I.E. socks4a, socks4h, socks5a, socks5h
+    """
     use_rdns = True
 
 
@@ -54,12 +70,11 @@ 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:')
+        if use_rdns := urlparse(proxy_url).scheme in DATA_SOURCE_SOCK_RDNS_SCHEMAS:
+            proxy_url = proxy_url.replace('socks5h:', 'socks5:').replace('socks5a:', 'socks5:')
+            proxy_url = proxy_url.replace('socks4h:', 'socks4:').replace('socks4a:', 'socks4:')
 
-        connection_pool_kw['_socks_options'] = {'proxy_url': cleaned_proxy_url}
+        connection_pool_kw['_socks_options'] = {'proxy_url': proxy_url}
         connection_pool_kw['timeout'] = timeout
 
         super().__init__(num_pools, headers, **connection_pool_kw)