Fix greatuk 1595 check x forward for header

conf/env.py

@@ -28,6 +28,8 @@ class BaseSettings(PydanticBaseSettings):
     sentry_enable_tracing: bool = False
     sentry_traces_sample_rate: float = 1.0
 
+    allowed_ips: str = ''
+
     feature_enforce_staff_sso_enabled: bool = False
 
     staff_sso_authbroker_url: str

conf/settings.py

@@ -29,6 +29,8 @@
 # PaaS, we can open ALLOWED_HOSTS
 ALLOWED_HOSTS = ['*']
 
+ALLOWED_IPS = [host.strip() for host in env.allowed_ips.split(',')]
+
 
 INSTALLED_APPS = [
     'django.contrib.auth',
@@ -71,6 +73,7 @@
     'django.middleware.cache.UpdateCacheMiddleware',
     'directory_components.middleware.MaintenanceModeMiddleware',
     'core.middleware.SSODisplayLoggedInCookieMiddleware',
+    'core.middleware.XForwardForCheckMiddleware',
     'django.middleware.security.SecurityMiddleware',
     'whitenoise.middleware.WhiteNoiseMiddleware',
     'conf.signature.SignatureCheckMiddleware',

core/middleware.py

@@ -1,3 +1,4 @@
+from dbt_copilot_python.utility import is_copilot
 from django.conf import settings
 from django.http import HttpResponse
 from django.utils.deprecation import MiddlewareMixin
@@ -37,3 +38,18 @@ def process_view(self, request, view_func, view_args, view_kwarg):
             if self.is_admin_name_space(request) or request.path_info.startswith('/admin/login'):
                 if not request.user.is_staff:
                     return HttpResponse(self.SSO_UNAUTHORISED_ACCESS_MESSAGE, status=401)
+
+
+class XForwardForCheckMiddleware(MiddlewareMixin):
+    CLIENT_IP_ERROR_MESSAGE = 'X Forward For checks failed'
+
+    def process_request(self, request):
+        if is_copilot():
+            # 200 response if client IP from x-forwarded-for header in ALLOWED_IPS, else 401.
+            try:
+                client_ips = request.META['HTTP_X_FORWARDED_FOR'].split(',')
+                for ip in client_ips:
+                    if ip.strip() not in settings.ALLOWED_IPS:
+                        return HttpResponse(self.CLIENT_IP_ERROR_MESSAGE, status=401)
+            except KeyError:
+                pass

core/tests/test_middleware.py

@@ -1,4 +1,7 @@
+import os
+
 import pytest
+from dbt_copilot_python.utility import is_copilot
 from django.http import HttpResponse
 from django.test.client import Client
 from django.urls import reverse
@@ -80,3 +83,43 @@ def test_admin_permission_middleware_authorised_with_staff(client, settings, adm
     response = client.get(reverse('admin:login'))
 
     assert response.status_code == 302
+
+
+@pytest.mark.django_db
+def test_x_forward_for_middleware_with_expected_ip(client, settings):
+    os.environ["COPILOT_ENVIRONMENT_NAME"] = "dev"
+    settings.ALLOWED_IPS = ['1.2.3.4', '123.123.123.123']
+    reload_urlconf()
+
+    # Middleware is for DBT only and should only trigger is is_copilot() is true
+    assert is_copilot() is True
+
+    response = client.get(
+        reverse('pingdom'),
+        content_type='',
+        HTTP_X_FORWARDED_FOR='1.2.3.4, 123.123.123.123',
+    )
+
+    assert response.status_code == 200
+    os.environ.pop("COPILOT_ENVIRONMENT_NAME")
+
+
+@pytest.mark.django_db
+def test_x_forward_for_middleware_with_unexpected_ip(client, settings):
+    os.environ["COPILOT_ENVIRONMENT_NAME"] = "dev"
+    settings.ALLOWED_IPS = [
+        '0.0.0.0',
+    ]
+    reload_urlconf()
+
+    # Middleware is for DBT only and should only trigger is is_copilot() is true
+    assert is_copilot() is True
+
+    response = client.get(
+        reverse('pingdom'),
+        content_type='',
+        HTTP_X_FORWARDED_FOR='1.2.3.4, 123.123.123.123',
+    )
+
+    assert response.status_code == 401
+    os.environ.pop("COPILOT_ENVIRONMENT_NAME")