Convenience method for getting supported algorithms from metadata

2 files changed, 48 insertions, 0 deletions

diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index b7f014aa..82bfa138 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -37,6 +37,8 @@ from saml2.validate import NotValid
 from saml2.sigver import security_context
 from saml2.extension.mdattr import NAMESPACE as NS_MDATTR
 from saml2.extension.mdattr import EntityAttributes
+from saml2.extension.algsupport import NAMESPACE as NS_ALGSUPPORT
+from saml2.extension.algsupport import SigningMethod, DigestMethod
 from saml2.extension.mdui import NAMESPACE as NS_MDUI
 from saml2.extension.mdui import UIInfo
 from saml2.extension.mdui import DisplayName
@@ -52,6 +54,8 @@ classnames = {
     "mdattr_entityattributes": "{ns}&{tag}".format(
         ns=NS_MDATTR, tag=EntityAttributes.c_tag
     ),
+    "algsupport_signing_method": "{ns}&{tag}".format(ns=NS_ALGSUPPORT, tag=SigningMethod.c_tag),
+    "algsupport_digest_method": "{ns}&{tag}".format(ns=NS_ALGSUPPORT, tag=DigestMethod.c_tag),
     "mdui_uiinfo": "{ns}&{tag}".format(ns=NS_MDUI, tag=UIInfo.c_tag),
     "mdui_uiinfo_display_name": "{ns}&{tag}".format(ns=NS_MDUI, tag=DisplayName.c_tag),
     "mdui_uiinfo_description": "{ns}&{tag}".format(ns=NS_MDUI, tag=Description.c_tag),
@@ -1282,6 +1286,36 @@ class MetadataStore(MetaData):
                         "attribute_value"]]
         return res
 
+    def supported_algorithms(self, entity_id):
+        """
+        Get all supported algorithms for an entry in the metadata.
+
+        Example return data:
+
+        {'digest_methods': ['http://www.w3.org/2001/04/xmldsig-more#sha224', 'http://www.w3.org/2001/04/xmlenc#sha256'],
+         'signing_methods': ['http://www.w3.org/2001/04/xmldsig-more#rsa-sha256']}
+
+        :param entity_id: Entity id
+        :return: dict with keys and value-lists from metadata
+
+        :type entity_id: string
+        :rtype: dict
+        """
+        res = {
+            'digest_methods': [],
+            'signing_methods': []
+        }
+        try:
+            ext = self.__getitem__(entity_id)["extensions"]
+        except KeyError:
+            return res
+        for elem in ext["extension_elements"]:
+            if elem["__class__"] == classnames["algsupport_digest_method"]:
+                res['digest_methods'].append(elem['algorithm'])
+            elif elem["__class__"] == classnames["algsupport_signing_method"]:
+                res['signing_methods'].append(elem['algorithm'])
+        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 630821db..c77293bb 100644
--- a/tests/test_30_mdstore.py
+++ b/tests/test_30_mdstore.py
@@ -54,6 +54,7 @@ TEST_METADATA_STRING = """
 <EntitiesDescriptor
     xmlns="urn:oasis:names:tc:SAML:2.0:metadata"
     xmlns:md="urn:oasis:names:tc:SAML:2.0:metadata"
+    xmlns:alg="urn:oasis:names:tc:SAML:metadata:algsupport"
     xmlns:shibmeta="urn:mace:shibboleth:metadata:1.0"
     xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
     xmlns:ds="http://www.w3.org/2000/09/xmldsig#"
@@ -61,6 +62,10 @@ TEST_METADATA_STRING = """
   <EntityDescriptor
     entityID="http://xenosmilus.umdc.umu.se/simplesaml/saml2/idp/metadata.php"
     xml:base="swamid-1.0/idp.umu.se-saml2.xml">
+    <md:Extensions>
+        <alg:DigestMethod Algorithm="http://www.w3.org/2001/04/xmlenc#sha256"/>
+        <alg:SigningMethod Algorithm="http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"/>
+    </md:Extensions>
   <IDPSSODescriptor protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">
     <KeyDescriptor>
       <ds:KeyInfo>
@@ -488,6 +493,15 @@ def test_metadata_extension_algsupport():
     assert mds
 
 
+def test_supported_algorithms():
+    mds = MetadataStore(ATTRCONV, sec_config,
+                        disable_ssl_certificate_validation=True)
+    mds.imp(METADATACONF["11"])
+    algs = mds.supported_algorithms(entity_id='http://xenosmilus.umdc.umu.se/simplesaml/saml2/idp/metadata.php')
+    assert 'http://www.w3.org/2001/04/xmlenc#sha256' in algs['digest_methods']
+    assert 'http://www.w3.org/2001/04/xmldsig-more#rsa-sha256' in algs['signing_methods']
+
+
 def test_extension():
     mds = MetadataStore(ATTRCONV, None)
     # use ordered dict to force expected entity to be last