From e7b39401781e89f9dadb0f8ed2e9dea43f113c80 Mon Sep 17 00:00:00 2001
From: Wolf Vollprecht <w.vollprecht@gmail.com>
Date: Tue, 30 Jan 2024 17:01:43 +0100
Subject: [PATCH] feat: add google authentication proxy middleware plugin
 (#683)

---
 environment.yml                               |   2 +-
 .../quetz_googleiap/__init__.py               |   0
 .../quetz_googleiap/middleware.py             | 123 ++++++++++++++++++
 plugins/quetz_googleiap/setup.py              |  14 ++
 plugins/quetz_googleiap/tests/conftest.py     |  15 +++
 .../quetz_googleiap/tests/test_googleiap.py   |  23 ++++
 quetz/main.py                                 |  20 +--
 setup.cfg                                     |   2 +-
 8 files changed, 189 insertions(+), 10 deletions(-)
 create mode 100644 plugins/quetz_googleiap/quetz_googleiap/__init__.py
 create mode 100644 plugins/quetz_googleiap/quetz_googleiap/middleware.py
 create mode 100644 plugins/quetz_googleiap/setup.py
 create mode 100644 plugins/quetz_googleiap/tests/conftest.py
 create mode 100644 plugins/quetz_googleiap/tests/test_googleiap.py

diff --git a/environment.yml b/environment.yml
index 6a498249..7c2e78c7 100644
--- a/environment.yml
+++ b/environment.yml
@@ -42,7 +42,7 @@ dependencies:
   - adlfs
   - importlib_metadata
   - pre-commit
-  - pytest
+  - pytest 7.*
   - pytest-mock
   - rq
   - libcflib
diff --git a/plugins/quetz_googleiap/quetz_googleiap/__init__.py b/plugins/quetz_googleiap/quetz_googleiap/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/plugins/quetz_googleiap/quetz_googleiap/middleware.py b/plugins/quetz_googleiap/quetz_googleiap/middleware.py
new file mode 100644
index 00000000..eaedd3d0
--- /dev/null
+++ b/plugins/quetz_googleiap/quetz_googleiap/middleware.py
@@ -0,0 +1,123 @@
+import logging
+import uuid
+
+from starlette.middleware.base import BaseHTTPMiddleware
+
+import quetz.authentication.base as auth_base
+from quetz import rest_models
+from quetz.config import Config, ConfigEntry, ConfigSection
+from quetz.dao import Dao
+from quetz.deps import get_config, get_db
+
+logger = logging.getLogger("quetz.googleiam")
+
+
+def email_to_channel_name(email):
+    name = email.split("@")[0]
+    name = name.replace(".", "-")
+    name = name.replace("_", "-")
+    return name
+
+
+class GoogleIAMMiddleware(BaseHTTPMiddleware):
+    """
+    Handles Google IAM headers and authorizes users based on the
+    Google IAM headers.
+    """
+
+    def __init__(self, app, config: Config):
+        if config is not None:
+            self.configure(config)
+        else:
+            self.configured = False
+
+        super().__init__(app)
+
+    def configure(self, config: Config):
+        config.register(
+            [
+                ConfigSection(
+                    "googleiam",
+                    [
+                        ConfigEntry("server_admin_emails", list, default=[]),
+                    ],
+                )
+            ]
+        )
+
+        # load configuration values
+        if config.configured_section("googleiam"):
+            self.server_admin_emails = config.googleiam_server_admin_emails
+            logger.info("Google IAM successfully configured")
+            logger.info(f"Google IAM server admin emails: {self.server_admin_emails}")
+            self.configured = True
+        else:
+            self.configured = False
+
+    async def dispatch(self, request, call_next):
+        # ignore middleware if it is not configured
+        if not self.configured or request.url.path.startswith("/health"):
+            response = await call_next(request)
+            return response
+
+        user_id = request.headers.get("x-goog-authenticated-user-id")
+        email = request.headers.get("x-goog-authenticated-user-email")
+
+        if user_id and email:
+            db = next(get_db(get_config()))
+            dao = Dao(db)
+
+            _, email = email.split(":", 1)
+            _, user_id = user_id.split(":", 1)
+
+            user = dao.get_user_by_username(email)
+            if not user:
+                email_data: auth_base.Email = {
+                    "email": email,
+                    "verified": True,
+                    "primary": True,
+                }
+                user = dao.create_user_with_profile(
+                    email, "google", user_id, email, "", None, True, [email_data]
+                )
+            user_channel = email_to_channel_name(email)
+
+            if dao.get_channel(email_to_channel_name(user_channel)) is None:
+                logger.info(f"Creating channel for user: {user_channel}")
+                channel = rest_models.Channel(
+                    name=user_channel,
+                    private=False,
+                    description="Channel for user: " + email,
+                )
+                dao.create_channel(channel, user.id, "owner")
+
+            self.google_role_for_user(user_id, email, dao)
+            user_id = uuid.UUID(bytes=user.id)
+            # drop the db and dao to remove the connection
+            del db, dao
+            # we also need to find the role of the user
+            request.session['identity_provider'] = "dummy"
+            request.session["user_id"] = str(user_id)
+        else:
+            request.session["user_id"] = None
+            request.session["identity_provider"] = None
+
+        response = await call_next(request)
+        return response
+
+    def google_role_for_user(self, user_id, username, dao):
+        if not user_id or not username:
+            return
+
+        if username in self.server_admin_emails:
+            logger.info(f"User '{username}' with user id '{user_id}' is server admin")
+            dao.set_user_role(user_id, "owner")
+        else:
+            logger.info(
+                f"User '{username}' with user id '{user_id}' is not a server admin"
+            )
+            dao.set_user_role(user_id, "member")
+
+
+def middleware():
+    return GoogleIAMMiddleware
diff --git a/plugins/quetz_googleiap/setup.py b/plugins/quetz_googleiap/setup.py
new file mode 100644
index 00000000..bfb664f0
--- /dev/null
+++ b/plugins/quetz_googleiap/setup.py
@@ -0,0 +1,14 @@
+from setuptools import setup
+
+plugin_name = "quetz-googleiap"
+
+setup(
+    name=plugin_name,
+    install_requires=[],
+    entry_points={
+        "quetz.middlewares": [f"{plugin_name} = quetz_googleiap.middleware"],
+    },
+    packages=[
+        "quetz_googleiap",
+    ],
+)
diff --git a/plugins/quetz_googleiap/tests/conftest.py b/plugins/quetz_googleiap/tests/conftest.py
new file mode 100644
index 00000000..6c2854bf
--- /dev/null
+++ b/plugins/quetz_googleiap/tests/conftest.py
@@ -0,0 +1,15 @@
+import pytest
+
+pytest_plugins = "quetz.testing.fixtures"
+
+
+@pytest.fixture
+def plugins():
+    # defines plugins to enable for testing
+    return ['quetz-googleiap']
+
+
+@pytest.fixture
+def sqlite_in_memory():
+    # use sqlite on disk so that we can modify it in a different process
+    return False
diff --git a/plugins/quetz_googleiap/tests/test_googleiap.py b/plugins/quetz_googleiap/tests/test_googleiap.py
new file mode 100644
index 00000000..3e9cd386
--- /dev/null
+++ b/plugins/quetz_googleiap/tests/test_googleiap.py
@@ -0,0 +1,23 @@
+import pytest
+
+
+@pytest.mark.parametrize(
+    "config_extra", ['[googleiam]\nserver_admin_emails=["test@tester.com"]']
+)
+def test_authentication(client, db):
+    response = client.get("/api/me")
+    assert response.status_code == 401
+
+    # add headers
+    headers = {
+        'X-Goog-Authenticated-User-Email': 'accounts.google.com:someone@tester.com',
+        'X-Goog-Authenticated-User-Id': 'accounts.google.com:someone@tester.com',
+    }
+
+    response = client.get("/api/me", headers=headers)
+    assert response.status_code == 200
+
+    # # check if channel was created
+    # response = client.get("/api/channels", headers=headers)
+    # assert response.status_code == 200
+    # assert response.json()['channels'][0]['name'] == 'someone'
diff --git a/quetz/main.py b/quetz/main.py
index 9bd35574..0a653063 100644
--- a/quetz/main.py
+++ b/quetz/main.py
@@ -104,12 +104,6 @@
 
 logger = logging.getLogger("quetz")
 
-app.add_middleware(
-    SessionMiddleware,
-    secret_key=config.session_secret,
-    https_only=config.session_https_only,
-)
-
 if config.general_redirect_http_to_https:
     logger.info("Configuring http to https redirect ")
     app.add_middleware(HTTPSRedirectMiddleware)
@@ -163,6 +157,18 @@ async def dispatch(self, request, call_next):
 
 app.add_middleware(CondaTokenMiddleware)
 
+plugin_middlewares: List[Type[BaseHTTPMiddleware]] = [
+    ep.load() for ep in entry_points().select(group='quetz.middlewares')
+]
+
+for middleware in plugin_middlewares:
+    app.add_middleware(middleware.middleware(), config=config)
+
+app.add_middleware(
+    SessionMiddleware,
+    secret_key=config.session_secret,
+    https_only=config.session_https_only,
+)
 
 if config.configured_section("profiling") and config.profiling_enable_sampling:
     from pyinstrument.profiler import Profiler
@@ -187,8 +193,6 @@ async def profile_request(
 pkgstore = config.get_package_store()
 
 # authenticators
-
-
 builtin_authenticators: List[Type[BaseAuthenticator]] = [
     authenticator
     for authenticator in [
diff --git a/setup.cfg b/setup.cfg
index 81c3b02d..9ac47e75 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -76,7 +76,7 @@ dev =
   flake8
   isort
   pre-commit
-  pytest
+  pytest >=7,<8
   pytest-asyncio
   pytest-mock
   pytest-cov