Port REMOTE_GROUPS feature from upstream.

See mozilla/redash#311 and mozilla/redash#37.

pyproject.toml

@@ -57,4 +57,5 @@ datasource_link = "redash_stmo.data_sources.link:extension"
 datasource_validator = "redash_stmo.data_sources.validator:extension"
 datasource_version = "redash_stmo.data_sources.version:extension"
 handler_queryresults = "redash_stmo.handlers.query_results:extension"
+handler_remote_user_auth = "redash_stmo.handlers.authentication.remote_user_auth:extension"
 queryrunner_presto = "redash_stmo.query_runner.presto:extension"

redash_stmo/handlers/authentication/remote_user_auth.py

@@ -0,0 +1,53 @@
+from flask import redirect, request, url_for
+from redash import settings
+from redash.authentication import get_next_path
+from redash.authentication.org_resolving import current_org
+from redash.authentication.remote_user_auth import logger
+
+from redash_stmo import settings as extension_settings
+
+
+def check_remote_groups():
+    """Check if there is a header of user groups and if yes
+    check it against a list of allowed user groups from the settings"""
+    # Quick shortcut out if remote user auth or remote groups aren't enabled
+    if (
+        not settings.REMOTE_USER_LOGIN_ENABLED
+        or not extension_settings.REMOTE_GROUPS_ENABLED
+    ):
+        return
+
+    # Generate the URL to the remote auth login endpoint
+    if settings.MULTI_ORG:
+        org = current_org._get_current_object()
+        remote_auth_path = url_for("remote_user_auth.login", org_slug=org.slug)
+    else:
+        remote_auth_path = url_for("remote_user_auth.login")
+
+    # Then only act if the request is for the remote user auth view
+    if request.path.startswith(remote_auth_path):
+        remote_groups = settings.set_from_string(
+            request.headers.get(extension_settings.REMOTE_GROUPS_HEADER) or ""
+        )
+        # Finally check if the remote groups found in the request header
+        # intersect with the allowed remote groups
+        if not extension_settings.REMOTE_GROUPS_ALLOWED.intersection(remote_groups):
+            logger.error(
+                "User groups provided in the %s header are not "
+                "matching the allowed groups.",
+                extension_settings.REMOTE_GROUPS_HEADER,
+            )
+            # Otherwise redirect back to the frontpage
+            unsafe_next_path = request.args.get("next")
+            next_path = get_next_path(unsafe_next_path)
+            if settings.MULTI_ORG:
+                org = current_org._get_current_object()
+                index_url = url_for("redash.index", org_slug=org.slug, next=next_path)
+            else:
+                index_url = url_for("redash.index", next=next_path)
+            return redirect(index_url)
+
+
+def extension(app):
+    """An extension that checks the REMOTE_GROUPS_HEADER."""
+    app.before_request(check_remote_groups)

redash_stmo/settings.py

@@ -1,4 +1,21 @@
 import os
 
+from redash.settings.helpers import parse_boolean, set_from_string
+
+
 # Frequency of health query runs in minutes (12 hours by default)
-HEALTH_QUERIES_REFRESH_SCHEDULE = int(os.environ.get("REDASH_HEALTH_QUERIES_REFRESH_SCHEDULE", 720))
+HEALTH_QUERIES_REFRESH_SCHEDULE = int(
+    os.environ.get("REDASH_HEALTH_QUERIES_REFRESH_SCHEDULE", 720)
+)
+
+# When enabled this will match the given remote groups request header with a
+# configured list of allowed user groups.
+REMOTE_GROUPS_ENABLED = parse_boolean(
+    os.environ.get("REDASH_REMOTE_GROUPS_ENABLED", "false")
+)
+REMOTE_GROUPS_HEADER = os.environ.get(
+    "REDASH_REMOTE_GROUPS_HEADER", "X-Forwarded-Remote-Groups"
+)
+REMOTE_GROUPS_ALLOWED = set_from_string(
+    os.environ.get("REDASH_REMOTE_GROUPS_ALLOWED", "")
+)

tests/handlers/test_remote_user_auth.py

@@ -0,0 +1,40 @@
+import mock
+from flask import url_for
+from redash import settings
+
+from redash_stmo import settings as extension_settings
+from tests import BaseTestCase
+
+
+class TestRemoteAuthGroups(BaseTestCase):
+    @mock.patch.object(settings, "REMOTE_USER_LOGIN_ENABLED", return_value=True)
+    @mock.patch.object(extension_settings, "REMOTE_GROUPS_ENABLED", return_value=True)
+    @mock.patch.object(
+        extension_settings,
+        "REMOTE_GROUPS_ALLOWED",
+        return_value=set(["admins", "managers"]),
+    )
+    @mock.patch("redash.authentication.remote_user_auth.logger.error")
+    def test_redirect_if_disabled(self, mock_logger, *args, **kwargs):
+        """Test to make sure requests to /login are directed to the
+        remote auth URL"""
+        next_path = "/"
+        if settings.MULTI_ORG:
+            test_url = url_for(
+                "remote_user_auth.login", org_slug="default", next=next_path
+            )
+        else:
+            test_url = url_for("remote_user_auth.login", next=next_path)
+
+        with self.app.test_request_context(test_url):
+            # make sure to call the before_request callback used
+            self.app.preprocess_request()
+            response = self.get_request(
+                test_url, headers={"X-Forwarded-Remote-Groups": "staff,contributors"}
+            )
+            self.assertTrue(mock_logger.called)
+            self.assertEqual(response.status_code, 302)
+            index_url = url_for(
+                "redash.index", org_slug="default", next=next_path, _external=True
+            )
+            self.assertEqual(response.location, index_url)