Changes in the MetaDataMDX class to allow for entity id transformation.

3 files changed, 36 insertions, 19 deletions

diff --git a/example/idp2/idp.py b/example/idp2/idp.py
index efa8e92f..940cb49c 100755
--- a/example/idp2/idp.py
+++ b/example/idp2/idp.py
@@ -135,8 +135,8 @@ class Service(object):
                                saml_msg["RelayState"],
                                encrypt_cert=_encrypt_cert)
             except KeyError:
-                # Can live with no relay state
-                return self.do(saml_msg["SAMLRequest"], binding)
+                # Can live with no relay state # TODO or can we, for inacademia?
+                return self.do(saml_msg["SAMLRequest"], binding, saml_msg["RelayState"])
 
     def artifact_operation(self, saml_msg):
         if not saml_msg:
@@ -400,20 +400,29 @@ class SSO(Service):
         """
         logger.info("--- In SSO POST ---")
         saml_msg = self.unpack_either()
-        self.req_info = IDP.parse_authn_request(
-            saml_msg["SAMLRequest"], BINDING_HTTP_POST)
-        _req = self.req_info.message
-        if self.user:
-            if _req.force_authn:
+
+        try:
+            _key = saml_msg["key"]
+            saml_msg = IDP.ticket[_key]
+            self.req_info = saml_msg["req_info"]
+            del IDP.ticket[_key]
+        except KeyError:
+            self.req_info = IDP.parse_authn_request(
+                saml_msg["SAMLRequest"], BINDING_HTTP_POST)
+            _req = self.req_info.message
+            if self.user:
+                if _req.force_authn:
+                    saml_msg["req_info"] = self.req_info
+                    key = self._store_request(saml_msg)
+                    return self.not_authn(key, _req.requested_authn_context)
+                else:
+                    return self.operation(saml_msg, BINDING_HTTP_POST)
+            else:
                 saml_msg["req_info"] = self.req_info
                 key = self._store_request(saml_msg)
                 return self.not_authn(key, _req.requested_authn_context)
-            else:
-                return self.operation(saml_msg, BINDING_HTTP_POST)
         else:
-            saml_msg["req_info"] = self.req_info
-            key = self._store_request(saml_msg)
-            return self.not_authn(key, _req.requested_authn_context)
+            return self.operation(saml_msg, BINDING_HTTP_POST)
 
     # def artifact(self):
     #     # Can be either by HTTP_Redirect or HTTP_POST
diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index c9d3aa09..faab780f 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -121,7 +121,7 @@ class MetaData(object):
         self.entities_descr = None
         self.entity_descr = None
         self.check_validity = check_validity
-
+        
     def items(self):
         return self.entity.items()
 
@@ -569,9 +569,14 @@ SAML_METADATA_CONTENT_TYPE = 'application/samlmetadata+xml'
 
 
 class MetaDataMDX(MetaData):
-
-    def __init__(self, onts, attrc, url, security, cert, http, **kwargs):
+    """ Uses the md protocol to fetch entity information
+    """
+    def __init__(self, entity_transform, onts, attrc, url, security, cert,
+                 http, **kwargs):
         """
+        :params entity_transform: function transforming (e.g. base64 or sha1
+        hash) the entity id. It is applied to the entity id before it is
+        concatenated with the request URL sent to the MDX server.
         :params onts:
         :params attrc:
         :params url:
@@ -584,6 +589,7 @@ class MetaDataMDX(MetaData):
         self.security = security
         self.cert = cert
         self.http = http
+        self.entity_transform = entity_transform
 
     def load(self):
         pass
@@ -592,7 +598,7 @@ class MetaDataMDX(MetaData):
         try:
             return self.entity[item]
         except KeyError:
-            mdx_url = "%s/entities/%s" % (self.url, quote_plus(item))
+            mdx_url = "%s/entities/%s" % (self.url, self.entity_transform(item))
             response = self.http.send(
                 mdx_url, headers={'Accept': SAML_METADATA_CONTENT_TYPE})
             if response.status_code == 200:
@@ -616,7 +622,6 @@ class MetaDataMDX(MetaData):
             raise KeyError
 
 
-
 class MetadataStore(object):
     def __init__(self, onts, attrc, config, ca_certs=None,
                  check_validity=True,
diff --git a/tests/test_30_mdstore.py b/tests/test_30_mdstore.py
index 82ac9907..0d8886b5 100644
--- a/tests/test_30_mdstore.py
+++ b/tests/test_30_mdstore.py
@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 import datetime
 import re
+from urllib import quote_plus
 from saml2.httpbase import HTTPBase
 
 from saml2.mdstore import MetadataStore, MetaDataMDX
@@ -230,7 +231,8 @@ def test_mdx_service():
     sec_config.xmlsec_binary = sigver.get_xmlsec_binary(["/opt/local/bin"])
     http = HTTPBase(verify=False, ca_bundle=None)
 
-    mdx = MetaDataMDX(ONTS.values(), ATTRCONV, "http://pyff-test.nordu.net",
+    mdx = MetaDataMDX(quote_plus, ONTS.values(), ATTRCONV,
+                      "http://pyff-test.nordu.net",
                       sec_config, None, http)
     foo = mdx.service("https://idp.umu.se/saml2/idp/metadata.php",
                       "idpsso_descriptor", "single_sign_on_service")
@@ -243,7 +245,8 @@ def test_mdx_certs():
     sec_config.xmlsec_binary = sigver.get_xmlsec_binary(["/opt/local/bin"])
     http = HTTPBase(verify=False, ca_bundle=None)
 
-    mdx = MetaDataMDX(ONTS.values(), ATTRCONV, "http://pyff-test.nordu.net",
+    mdx = MetaDataMDX(quote_plus, ONTS.values(), ATTRCONV,
+                      "http://pyff-test.nordu.net",
                       sec_config, None, http)
     foo = mdx.certs("https://idp.umu.se/saml2/idp/metadata.php", "idpsso")