From d29399f2c9c3de5162481fdde7ed7648dc731445 Mon Sep 17 00:00:00 2001
From: Gijs Vermeulen <gverm@inuits.eu>
Date: Mon, 23 Oct 2023 13:21:21 +0200
Subject: [PATCH] Add multi-tenant policy

The idea is to add generic policies for elody here.
---
 src/elody/loader.py                           | 11 ++-
 src/elody/policies/__init__.py                |  0
 src/elody/policies/authentication/__init__.py |  0
 .../authentication/multi_tenant_policy.py     | 68 +++++++++++++++++++
 4 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 src/elody/policies/__init__.py
 create mode 100644 src/elody/policies/authentication/__init__.py
 create mode 100644 src/elody/policies/authentication/multi_tenant_policy.py

diff --git a/src/elody/loader.py b/src/elody/loader.py
index 1efc7c6..9e9e4c4 100644
--- a/src/elody/loader.py
+++ b/src/elody/loader.py
@@ -89,7 +89,16 @@ def __instantiate_authentication_policy(policy_module_name, policy, logger):
         return policy(
             token_schema, os.getenv("ROLE_SCOPE_MAPPING", "role_scope_mapping.json")
         )
-
+    if policy_module_name == "elody.policies.authentication.multi_tenant_policy":
+        tenant_defining_types = os.getenv("TENANT_DEFINING_TYPES")
+        tenant_defining_types = (
+            tenant_defining_types.split(",") if tenant_defining_types else []
+        )
+        return policy(
+            os.getenv("TENANT_DEFINING_HEADER", "X-tenant-id"),
+            tenant_defining_types,
+            os.getenv("AUTO_CREATE_TENANTS"),
+        )
     return policy()
 
 
diff --git a/src/elody/policies/__init__.py b/src/elody/policies/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/elody/policies/authentication/__init__.py b/src/elody/policies/authentication/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/elody/policies/authentication/multi_tenant_policy.py b/src/elody/policies/authentication/multi_tenant_policy.py
new file mode 100644
index 0000000..053838e
--- /dev/null
+++ b/src/elody/policies/authentication/multi_tenant_policy.py
@@ -0,0 +1,68 @@
+from inuits_policy_based_auth import BaseAuthenticationPolicy, RequestContext
+from inuits_policy_based_auth.contexts import UserContext
+from storage.storagemanager import StorageManager
+from werkzeug.exceptions import Forbidden
+
+
+class MultiTenantPolicy(BaseAuthenticationPolicy):
+    """
+    An authentication policy that gets or creates (when applicable) a
+    tenant, using the details from the configured tenant defining
+    header (default 'X-tenant-id').
+
+    Parameters:
+    -----------
+    defining_header : str
+        Which header receives an identifier to get or create a
+        tenant. (default 'X-tenant-id')
+    defining_types : list
+        List of types that defines a tenant.
+    auto_create_tenants : bool
+        Configures the auto creation of tenants.
+    """
+
+    def __init__(self, defining_header, defining_types, auto_create_tenants):
+        self._defining_header = defining_header
+        self._defining_types = defining_types
+        self._auto_create_tenants = auto_create_tenants
+
+    def authenticate(self, user_context: UserContext, request_context: RequestContext):
+        """
+        Get tenant from tenant defining header and set x_tenant accordingly.
+
+        Parameters:
+        -----------
+        user_context : UserContext
+            The context of the user requesting authentication.
+        request_context : RequestContext
+            The context of the request.
+
+        Returns:
+        --------
+        UserContext
+            The user context with x_tenant set.
+
+        Raises:
+        -------
+        Unauthorized
+            If the authentication fails.
+        """
+
+        if self._defining_types:
+            return user_context
+        auth_header = self._defining_header
+        if not (tenant_id := request_context.http_request.headers.get(auth_header)):
+            raise Forbidden(description=f"{auth_header} header not present")
+        storage = StorageManager().get_db_engine()
+        tenant = storage.get_item_from_collection_by_id("entities", tenant_id)
+        if not tenant:
+            if not self._auto_create_tenants:
+                raise Forbidden(
+                    description=f"Tenant with identifier {tenant_id} not found"
+                )
+            tenant = storage.save_item_to_collection(
+                "entities", {"type": "tenant", "identifiers": [tenant_id]}
+            )
+        user_context.x_tenant.id = tenant["_id"]
+        user_context.x_tenant.raw = tenant
+        return user_context