From 6209714f8748934e60ae7d80c3974441f8155eb6 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Thu, 9 Dec 2021 09:37:41 +0100
Subject: [PATCH] policies/expression: add ak_call_policy

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/policies/expression/evaluator.py | 13 +++++++++++++
 authentik/policies/process.py              | 18 +++++++++++-------
 authentik/root/celery.py                   |  1 -
 website/docs/policies/expression.mdx       | 19 ++++++++++++++++++-
 4 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/authentik/policies/expression/evaluator.py b/authentik/policies/expression/evaluator.py
index 6d7fa59c6..04df20d55 100644
--- a/authentik/policies/expression/evaluator.py
+++ b/authentik/policies/expression/evaluator.py
@@ -11,6 +11,8 @@ from authentik.flows.planner import PLAN_CONTEXT_SSO
 from authentik.lib.expression.evaluator import BaseEvaluator
 from authentik.lib.utils.http import get_client_ip
 from authentik.policies.exceptions import PolicyException
+from authentik.policies.models import Policy, PolicyBinding
+from authentik.policies.process import PolicyProcess
 from authentik.policies.types import PolicyRequest, PolicyResult
 
 LOGGER = get_logger()
@@ -31,6 +33,7 @@ class PolicyEvaluator(BaseEvaluator):
         self._context["ak_logger"] = get_logger(policy_name)
         self._context["ak_message"] = self.expr_func_message
         self._context["ak_user_has_authenticator"] = self.expr_func_user_has_authenticator
+        self._context["ak_call_policy"] = self.expr_func_call_policy
         self._context["ip_address"] = ip_address
         self._context["ip_network"] = ip_network
         self._filename = policy_name or "PolicyEvaluator"
@@ -39,6 +42,16 @@ class PolicyEvaluator(BaseEvaluator):
         """Wrapper to append to messages list, which is returned with PolicyResult"""
         self._messages.append(message)
 
+    def expr_func_call_policy(self, name: str, **kwargs) -> PolicyResult:
+        """Call policy by name, with current request"""
+        policy = Policy.objects.filter(name=name).select_subclasses().first()
+        if not policy:
+            raise ValueError(f"Policy '{name}' not found.")
+        req: PolicyRequest = self._context["request"]
+        req.context.update(kwargs)
+        proc = PolicyProcess(PolicyBinding(policy=policy), request=req, connection=None)
+        return proc.profiling_wrapper()
+
     def expr_func_user_has_authenticator(
         self, user: User, device_type: Optional[str] = None
     ) -> bool:
diff --git a/authentik/policies/process.py b/authentik/policies/process.py
index e8856297f..f5e07b8d5 100644
--- a/authentik/policies/process.py
+++ b/authentik/policies/process.py
@@ -127,8 +127,8 @@ class PolicyProcess(PROCESS_CLASS):
         )
         return policy_result
 
-    def run(self):  # pragma: no cover
-        """Task wrapper to run policy checking"""
+    def profiling_wrapper(self):
+        """Run with profiling enabled"""
         with Hub.current.start_span(
             op="policy.process.execute",
         ) as span, HIST_POLICIES_EXECUTION_TIME.labels(
@@ -142,8 +142,12 @@ class PolicyProcess(PROCESS_CLASS):
             span: Span
             span.set_data("policy", self.binding.policy)
             span.set_data("request", self.request)
-            try:
-                self.connection.send(self.execute())
-            except Exception as exc:  # pylint: disable=broad-except
-                LOGGER.warning(str(exc))
-                self.connection.send(PolicyResult(False, str(exc)))
+            return self.execute()
+
+    def run(self):  # pragma: no cover
+        """Task wrapper to run policy checking"""
+        try:
+            self.connection.send(self.profiling_wrapper())
+        except Exception as exc:  # pylint: disable=broad-except
+            LOGGER.warning(str(exc))
+            self.connection.send(PolicyResult(False, str(exc)))
diff --git a/authentik/root/celery.py b/authentik/root/celery.py
index 7267ae59a..17e9a863f 100644
--- a/authentik/root/celery.py
+++ b/authentik/root/celery.py
@@ -1,7 +1,6 @@
 """authentik core celery"""
 import os
 from logging.config import dictConfig
-from uuid import uuid4
 
 from celery import Celery
 from celery.signals import (
diff --git a/website/docs/policies/expression.mdx b/website/docs/policies/expression.mdx
index 5beca9b8c..bb102c201 100644
--- a/website/docs/policies/expression.mdx
+++ b/website/docs/policies/expression.mdx
@@ -25,7 +25,7 @@ ak_message("Access denied")
 return False
 ```
 
-### `ak_user_has_authenticator(user: User, device_type: Optional[str] = None)` (2021.9+)
+### `ak_user_has_authenticator(user: User, device_type: Optional[str] = None) -> bool` (2021.9+)
 
 Check if a user has any authenticator devices. Only fully validated devices are counted.
 
@@ -42,6 +42,23 @@ Example:
 return ak_user_has_authenticator(request.user)
 ```
 
+### `ak_call_policy(name: str, **kwargs) -> PolicyResult` (2021.12+)
+
+Call another policy with the name *name*. Current request is passed to policy. Key-word arguments
+can be used to modify the request's context.
+
+Example:
+
+```python
+result = ak_call_policy("test-policy")
+# result is a PolicyResult object, so you can access `.passing` and `.messages`.
+return result.passing
+
+result = ak_call_policy("test-policy-2", foo="bar")
+# Inside the `test-policy-2` you can then use `request.context["foo"]`
+return result.passing
+```
+
 import Functions from '../expressions/_functions.md'
 
 <Functions />