Decorate context with target entity id on mirrored frontend setup

Author: c00kiemon5ter

src/satosa/backends/saml2.py

@@ -17,6 +17,7 @@
 import satosa.util as util
 from satosa.base import SAMLBaseModule
 from satosa.base import SAMLEIDASBaseModule
+from satosa.context import Context
 from .base import BackendModule
 from ..exception import SATOSAAuthenticationError
 from ..internal_data import (InternalResponse,
@@ -92,15 +93,13 @@ def start_auth(self, context, internal_req):
         if len(idps) == 1 and "mdq" not in self.config["sp_config"]["metadata"]:
             return self.authn_request(context, idps[0])
 
-        try:
-            # find mirrored entity id
-            entity_id = context.internal_data["mirror.target_entity_id"]
-        except KeyError:
-            # redirect to discovery server
+        entity_id = context.get_decoration(
+                Context.KEY_MIRROR_TARGET_ENTITYID)
+        if None is entity_id:
             return self.disco_query()
-        else:
-            entity_id = urlsafe_b64decode(entity_id).decode("utf-8")
-            return self.authn_request(context, entity_id)
+
+        entity_id = urlsafe_b64decode(entity_id).decode("utf-8")
+        return self.authn_request(context, entity_id)
 
     def disco_query(self):
         """

src/satosa/context.py

@@ -15,6 +15,7 @@ class Context(object):
     """
     Holds information about the current request.
     """
+    KEY_MIRROR_TARGET_ENTITYID = 'mirror_target_entity_id'
 
     def __init__(self):
         self._path = None

src/satosa/frontends/saml2.py

@@ -16,6 +16,7 @@
 from saml2.server import Server
 
 from satosa.base import SAMLBaseModule
+from satosa.context import Context
 from .base import FrontendModule
 from ..internal_data import InternalRequest, UserIdHashType
 from ..logging_util import satosa_logging
@@ -503,7 +504,7 @@ def _load_idp_dynamic_endpoints(self, context):
         :return: An idp server
         """
         target_entity_id = context.path.split("/")[1]
-        context.internal_data["mirror.target_entity_id"] = target_entity_id
+        context.decorate(Context.KEY_MIRROR_TARGET_ENTITYID, target_entity_id)
         idp_conf_file = self._load_endpoints_to_config(context.target_backend, target_entity_id)
         idp_config = IdPConfig().load(idp_conf_file, metadata_construction=False)
         return Server(config=idp_config)

src/satosa/micro_services/custom_routing.py

@@ -1,6 +1,8 @@
 import logging
 from base64 import urlsafe_b64encode
 
+from satosa.context import Context
+
 from .base import RequestMicroService
 from ..exception import SATOSAConfigurationError
 from ..exception import SATOSAError
@@ -56,9 +58,9 @@ def _b64_url(self, data):
         return urlsafe_b64encode(data.encode("utf-8")).decode("utf-8")
 
     def process(self, context, data):
-        try:
-            target_entity_id = context.internal_data["mirror.target_entity_id"]
-        except KeyError:
+        target_entity_id = context.get_decoration(
+                Context.KEY_MIRROR_TARGET_ENTITYID)
+        if None is target_entity_id:
             logger.error("DecideIfRequesterIsAllowed can only be used with SAMLMirrorFrontend")
             raise SATOSAError("DecideIfRequesterIsAllowed can only be used with SAMLMirrorFrontend")
 

tests/satosa/backends/test_saml2.py

@@ -141,9 +141,12 @@ def test_full_flow(self, context, idp_conf, sp_conf):
         assert context.state[test_state_key] == "my_state"
         self.assert_authn_response(internal_resp)
 
-    def test_start_auth_redirects_directly_to_mirrored_idp(self, context, idp_conf):
-        context.internal_data["mirror.target_entity_id"] = urlsafe_b64encode(
-            idp_conf["entityid"].encode("utf-8")).decode("utf-8")
+    def test_start_auth_redirects_directly_to_mirrored_idp(
+            self, context, idp_conf):
+        entityid = idp_conf["entityid"]
+        entityid_bytes = entityid.encode("utf-8")
+        entityid_b64_str = urlsafe_b64encode(entityid_bytes).decode("utf-8")
+        context.decorate(Context.KEY_MIRROR_TARGET_ENTITYID, entityid_b64_str)
 
         resp = self.samlbackend.start_auth(context, InternalRequest(None, None))
         self.assert_redirect_to_idp(resp, idp_conf)

tests/satosa/micro_services/test_custom_routing.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from satosa.context import Context
 from satosa.exception import SATOSAError, SATOSAConfigurationError
 from satosa.internal_data import InternalRequest
 from satosa.micro_services.custom_routing import DecideIfRequesterIsAllowed
@@ -11,7 +12,9 @@
 
 @pytest.fixture
 def target_context(context):
-    context.internal_data["mirror.target_entity_id"] = urlsafe_b64encode(TARGET_ENTITY.encode("utf-8")).decode("utf-8")
+    entityid_bytes = TARGET_ENTITY.encode("utf-8")
+    entityid_b64_str = urlsafe_b64encode(entityid_bytes).decode("utf-8")
+    context.decorate(Context.KEY_MIRROR_TARGET_ENTITYID, entityid_b64_str)
     return context