Support registration info in mdstore

2 files changed, 59 insertions, 0 deletions

diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index c0eb1686..76a981a8 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -56,6 +56,9 @@ from saml2.extension.mdui import Description
 from saml2.extension.mdui import InformationURL
 from saml2.extension.mdui import PrivacyStatementURL
 from saml2.extension.mdui import Logo
+from saml2.extension.mdrpi import NAMESPACE as NS_MDRPI
+from saml2.extension.mdrpi import RegistrationInfo
+from saml2.extension.mdrpi import RegistrationPolicy
 
 
 logger = logging.getLogger(__name__)
@@ -79,6 +82,8 @@ classnames = {
     "service_artifact_resolution": "{ns}&{tag}".format(ns=NS_MD, tag=ArtifactResolutionService.c_tag),
     "service_single_sign_on": "{ns}&{tag}".format(ns=NS_MD, tag=SingleSignOnService.c_tag),
     "service_nameid_mapping": "{ns}&{tag}".format(ns=NS_MD, tag=NameIDMappingService.c_tag),
+    "mdrpi_registration_info": "{ns}&{tag}".format(ns=NS_MDRPI, tag=RegistrationInfo.c_tag),
+    "mdrpi_registration_policy": "{ns}&{tag}".format(ns=NS_MDRPI, tag=RegistrationPolicy.c_tag),
 }
 
 ENTITY_CATEGORY = "http://macedir.org/entity-category"
@@ -1406,6 +1411,45 @@ class MetadataStore(MetaData):
                 res['signing_methods'].append(elem['algorithm'])
         return res
 
+    def registration_info(self, entity_id):
+        """
+        Get all registration info for an entry in the metadata.
+
+        Example return data:
+
+        res = {
+            'registration_authority': 'http://www.example.com',
+            'registration_instant': '2013-06-15T18:15:03Z',
+            'registration_policy': {
+                'en': 'http://www.example.com/policy.html',
+                'sv': 'http://www.example.com/sv/policy.html',
+            }
+        }
+
+        :param entity_id: Entity id
+        :return: dict with keys and value-lists from metadata
+
+        :type entity_id: string
+        :rtype: dict
+        """
+        res = {
+            'registration_authority': None,
+            'registration_instant': None,
+            'registration_policy': {}
+        }
+        try:
+            ext = self.__getitem__(entity_id)["extensions"]
+        except KeyError:
+            return res
+        for elem in ext["extension_elements"]:
+            if elem["__class__"] == classnames["mdrpi_registration_info"]:
+                res["registration_authority"] = elem["registration_authority"]
+                res["registration_instant"] = elem.get("registration_instant")
+                for policy in elem.get('registration_policy'):
+                    if policy["__class__"] == classnames["mdrpi_registration_policy"]:
+                        res['registration_policy'][policy["lang"]] = policy["text"]
+        return res
+
     def _lookup_elements_by_cls(self, root, cls):
         elements = (
             element
diff --git a/tests/test_30_mdstore.py b/tests/test_30_mdstore.py
index d712383f..fede1010 100644
--- a/tests/test_30_mdstore.py
+++ b/tests/test_30_mdstore.py
@@ -152,6 +152,10 @@ METADATACONF = {
         "class": "saml2.mdstore.MetaDataFile",
         "metadata": [(full_path("uu.xml"),)],
     }],
+    "13": [{
+        "class": "saml2.mdstore.MetaDataFile",
+        "metadata": [(full_path("swamid-2.0.xml"),)],
+    }],
 }
 
 
@@ -560,6 +564,17 @@ def test_supported_algorithms():
     assert 'http://www.w3.org/2001/04/xmldsig-more#rsa-sha256' in algs['signing_methods']
 
 
+def test_registration_info():
+    mds = MetadataStore(ATTRCONV, sec_config,
+                        disable_ssl_certificate_validation=True)
+    mds.imp(METADATACONF["13"])
+    registration_info = mds.registration_info(entity_id='https://aai-idp.unibe.ch/idp/shibboleth')
+    assert 'http://rr.aai.switch.ch/' == registration_info['registration_authority']
+    assert '2013-06-15T18:15:03Z' == registration_info['registration_instant']
+    assert 'https://www.switch.ch/aai/federation/switchaai/metadata-registration-practice-statement-20110711.txt' == \
+           registration_info['registration_policy']['en']
+
+
 def test_extension():
     mds = MetadataStore(ATTRCONV, None)
     # use ordered dict to force expected entity to be last