Merge pull request #246 from skoranda/saml_co_frontend_04

Add SAMLVirtualCoFrontend IdP entityID to state

Author: c00kiemon5ter

src/satosa/frontends/saml2.py

@@ -700,6 +700,7 @@ class SAMLVirtualCoFrontend(SAMLFrontend):
     """
     KEY_CO = 'collaborative_organizations'
     KEY_CO_NAME = 'co_name'
+    KEY_CO_ENTITY_ID = 'co_entity_id'
     KEY_CO_ATTRIBUTES = 'co_static_saml_attributes'
     KEY_CONTACT_PERSON = 'contact_person'
     KEY_ENCODEABLE_NAME = 'encodeable_name'
@@ -770,6 +771,8 @@ def _create_state_data(self, context, resp_args, relay_state):
         """
         state = super()._create_state_data(context, resp_args, relay_state)
         state[self.KEY_CO_NAME] = context.get_decoration(self.KEY_CO_NAME)
+        state[self.KEY_CO_ENTITY_ID] = context.get_decoration(
+                                                         self.KEY_CO_ENTITY_ID)
 
         return state
 
@@ -869,19 +872,22 @@ def _add_endpoints_to_config(self, config, co_name, backend_name):
 
         return config
 
-    def _add_entity_id(self, config, co_name):
+    def _add_entity_id(self, context, config, co_name):
         """
         Use the CO name to construct the entity ID for the virtual IdP
-        for the CO.
+        for the CO and add it to the config. Also add it to the
+        context.
 
         The entity ID has the form
 
         {base_entity_id}/{co_name}
 
+        :type context: The current context
         :type config: satosa.satosa_config.SATOSAConfig
         :type co_name: str
         :rtype: satosa.satosa_config.SATOSAConfig
 
+        :param context:
         :param config: satosa proxy config
         :param co_name: CO name
 
@@ -890,6 +896,7 @@ def _add_entity_id(self, config, co_name):
         base_entity_id = config['entityid']
         co_entity_id = "{}/{}".format(base_entity_id, quote_plus(co_name))
         config['entityid'] = co_entity_id
+        context.decorate(self.KEY_CO_ENTITY_ID, co_entity_id)
 
         return config
 
@@ -910,7 +917,8 @@ def _overlay_for_saml_metadata(self, config, co_name):
         """
         all_co_configs = self.config[self.KEY_CO]
         co_config = next(
-            item for item in all_co_configs if item[self.KEY_ENCODEABLE_NAME] == co_name
+            item for item in all_co_configs
+            if item[self.KEY_ENCODEABLE_NAME] == co_name
         )
 
         key = self.KEY_ORGANIZATION
@@ -974,7 +982,7 @@ def _create_co_virtual_idp(self, context):
         idp_config = self._add_endpoints_to_config(idp_config,
                                                    co_name,
                                                    backend_name)
-        idp_config = self._add_entity_id(idp_config, co_name)
+        idp_config = self._add_entity_id(context, idp_config, co_name)
 
         # Use the overwritten IdP config to generate a pysaml2 config object
         # and from it a server object.

tests/satosa/frontends/test_saml2.py

@@ -483,11 +483,14 @@ def context(self, context):
 
         return context
 
-    def test_create_state_data(self, frontend, context):
-        context.decorate(frontend.KEY_CO_NAME, self.CO)
+    def test_create_state_data(self, frontend, context, idp_conf):
+        frontend._create_co_virtual_idp(context)
         state = frontend._create_state_data(context, {}, "")
         assert state[frontend.KEY_CO_NAME] == self.CO
 
+        expected_entityid = "{}/{}".format(idp_conf['entityid'], self.CO)
+        assert state[frontend.KEY_CO_ENTITY_ID] == expected_entityid
+
     def test_get_co_name(self, frontend, context):
         co_name = frontend._get_co_name(context)
         assert co_name == self.CO
@@ -534,7 +537,7 @@ def test_co_static_attributes(self, frontend, context, internal_response,
         backend_name = context.target_backend
         idp_conf = frontend._add_endpoints_to_config(idp_conf, co_name,
                                                      backend_name)
-        idp_conf = frontend._add_entity_id(idp_conf, co_name)
+        idp_conf = frontend._add_entity_id(context, idp_conf, co_name)
 
         # Use a utility function to serialize the idp_conf IdP configuration
         # fixture to a string and then dynamically update the sp_conf