Add mdstore method to extract mdui uiinfo elements

Signed-off-by: Ivan Kanakarakis <ivan.kanak@gmail.com>

1 files changed, 35 insertions, 0 deletions

diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index 021c4a9e..83a89581 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -37,6 +37,13 @@ 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.mdui import NAMESPACE as NS_MDUI
+from saml2.extension.mdui import UIInfo
+from saml2.extension.mdui import DisplayName
+from saml2.extension.mdui import Description
+from saml2.extension.mdui import InformationURL
+from saml2.extension.mdui import PrivacyStatementURL
+from saml2.extension.mdui import Logo
 
 
 logger = logging.getLogger(__name__)
@@ -45,6 +52,7 @@ classnames = {
     "mdattr_entityattributes": "{ns}&{tag}".format(
         ns=NS_MDATTR, tag=EntityAttributes.c_tag
     ),
+    "mdui_uiinfo": "{ns}&{tag}".format(ns=NS_MDUI, tag=UIInfo.c_tag),
 }
 
 ENTITY_CATEGORY = "http://macedir.org/entity-category"
@@ -1255,6 +1263,33 @@ class MetadataStore(MetaData):
                         "attribute_value"]]
         return res
 
+    def _mdui_uiinfo(self, entity_id):
+        descriptor_names = (
+            item
+            for item in self[entity_id].keys()
+            if item.endswith("_descriptor")
+        )
+        descriptors = (
+            descriptor
+            for descriptor_name in descriptor_names
+            for descriptor in self[entity_id].get(descriptor_name, [])
+        )
+        extensions = (
+            extension
+            for descriptor in descriptors
+            for extension in descriptor.get("extensions", {}).get("extension_elements", [])
+        )
+        uiinfos = (
+            extension
+            for extension in extensions
+            if extension.get("__class__") == classnames["mdui_uiinfo"]
+        )
+        return uiinfos
+
+    def mdui_uiinfo(self, entity_id):
+        uiinfos = list(self._mdui_uiinfo(entity_id))
+        return uiinfos
+
     def bindings(self, entity_id, typ, service):
         for _md in self.metadata.values():
             if entity_id in _md.items():