Merge pull request #230 from skoranda/dynamic_saml_disovery_service

Dynamic saml discovery service

Author: c00kiemon5ter

src/satosa/backends/saml2.py

@@ -40,6 +40,8 @@ class SAMLBackend(BackendModule, SAMLBaseModule):
     A saml2 backend module (acting as a SP).
     """
     KEY_DISCO_SRV = 'disco_srv'
+    KEY_SAML_DISCOVERY_SERVICE_URL = 'saml_discovery_service_url'
+    KEY_SAML_DISCOVERY_SERVICE_POLICY = 'saml_discovery_service_policy'
     KEY_SP_CONFIG = 'sp_config'
     VALUE_ACR_COMPARISON_DEFAULT = 'exact'
 
@@ -102,9 +104,9 @@ def start_auth(self, context, internal_req):
             entity_id = idps[0]
             return self.authn_request(context, entity_id)
 
-        return self.disco_query()
+        return self.disco_query(context)
 
-    def disco_query(self):
+    def disco_query(self, context):
         """
         Makes a request to the discovery server
 
@@ -116,8 +118,21 @@ def disco_query(self):
         :param internal_req: The request
         :return: Response
         """
-        return_url = self.sp.config.getattr("endpoints", "sp")["discovery_response"][0][0]
-        loc = self.sp.create_discovery_service_request(self.discosrv, self.sp.config.entityid, **{"return": return_url})
+        endpoints = self.sp.config.getattr("endpoints", "sp")
+        return_url = endpoints["discovery_response"][0][0]
+
+        disco_url = (
+            context.get_decoration(self.KEY_SAML_DISCOVERY_SERVICE_URL) or self.discosrv
+        )
+        disco_policy = context.get_decoration(self.KEY_SAML_DISCOVERY_SERVICE_POLICY)
+
+        args = {"return": return_url}
+        if disco_policy:
+            args["policy"] = disco_policy
+
+        loc = self.sp.create_discovery_service_request(
+            disco_url, self.sp.config.entityid, **args
+        )
         return SeeOther(loc)
 
     def construct_requested_authn_context(self, entity_id):

tests/satosa/backends/test_saml2.py

@@ -47,11 +47,13 @@ def assert_redirect_to_idp(self, redirect_response, idp_conf):
         assert redirect_location == idp_conf["service"]["idp"]["endpoints"]["single_sign_on_service"][0][0]
         assert "SAMLRequest" in parse_qs(parsed.query)
 
-    def assert_redirect_to_discovery_server(self, redirect_response, sp_conf):
+    def assert_redirect_to_discovery_server(
+        self, redirect_response, sp_conf, expected_discosrv_url
+    ):
         assert redirect_response.status == "303 See Other"
         parsed = urlparse(redirect_response.message)
         redirect_location = "{parsed.scheme}://{parsed.netloc}{parsed.path}".format(parsed=parsed)
-        assert redirect_location == DISCOSRV_URL
+        assert redirect_location == expected_discosrv_url
 
         request_params = dict(parse_qsl(parsed.query))
         assert request_params["return"] == sp_conf["service"]["sp"]["endpoints"]["discovery_response"][0][0]
@@ -99,7 +101,15 @@ def get_path_from_url(url):
 
     def test_start_auth_defaults_to_redirecting_to_discovery_server(self, context, sp_conf):
         resp = self.samlbackend.start_auth(context, InternalData())
-        self.assert_redirect_to_discovery_server(resp, sp_conf)
+        self.assert_redirect_to_discovery_server(resp, sp_conf, DISCOSRV_URL)
+
+    def test_discovery_server_set_in_context(self, context, sp_conf):
+        discosrv_url = 'https://my.org/saml_discovery_service'
+        context.decorate(
+            SAMLBackend.KEY_SAML_DISCOVERY_SERVICE_URL, discosrv_url
+        )
+        resp = self.samlbackend.start_auth(context, InternalData())
+        self.assert_redirect_to_discovery_server(resp, sp_conf, discosrv_url)
 
     def test_full_flow(self, context, idp_conf, sp_conf):
         test_state_key = "test_state_key_456afgrh"
@@ -110,7 +120,7 @@ def test_full_flow(self, context, idp_conf, sp_conf):
 
         # start auth flow (redirecting to discovery server)
         resp = self.samlbackend.start_auth(context, InternalData())
-        self.assert_redirect_to_discovery_server(resp, sp_conf)
+        self.assert_redirect_to_discovery_server(resp, sp_conf, DISCOSRV_URL)
 
         # fake response from discovery server
         disco_resp = parse_qs(urlparse(resp.message).query)
@@ -166,7 +176,7 @@ def test_always_redirect_to_discovery_service_if_using_mdq(self, context, sp_con
         samlbackend = SAMLBackend(None, INTERNAL_ATTRIBUTES, {"sp_config": sp_conf, "disco_srv": DISCOSRV_URL,},
                                   "base_url", "saml_backend")
         resp = samlbackend.start_auth(context, InternalData())
-        self.assert_redirect_to_discovery_server(resp, sp_conf)
+        self.assert_redirect_to_discovery_server(resp, sp_conf, DISCOSRV_URL)
 
     def test_authn_request(self, context, idp_conf):
         resp = self.samlbackend.authn_request(context, idp_conf["entityid"])