Add mdstore methods to extract mdui uiinfo subelements

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

1 files changed, 108 insertions, 0 deletions

diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index 83a89581..fe57bc8e 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -53,6 +53,15 @@ classnames = {
         ns=NS_MDATTR, tag=EntityAttributes.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),
+    "mdui_uiinfo_information_url": "{ns}&{tag}".format(
+        ns=NS_MDUI, tag=InformationURL.c_tag
+    ),
+    "mdui_uiinfo_privacy_statement_url": "{ns}&{tag}".format(
+        ns=NS_MDUI, tag=PrivacyStatementURL.c_tag
+    ),
+    "mdui_uiinfo_logo": "{ns}&{tag}".format(ns=NS_MDUI, tag=Logo.c_tag),
 }
 
 ENTITY_CATEGORY = "http://macedir.org/entity-category"
@@ -1290,6 +1299,105 @@ class MetadataStore(MetaData):
         uiinfos = list(self._mdui_uiinfo(entity_id))
         return uiinfos
 
+    def _mdui_uiinfo_i18n_elements(self, entity_id, langpref, element_hint, lookup):
+        uiinfos = self._mdui_uiinfo(entity_id)
+        elements = lookup(uiinfos, element_hint)
+        lang_elements = (
+            element
+            for element in elements
+            if element.get("lang") == langpref
+        )
+        values = (
+            value
+            for element in lang_elements
+            for value in [element.get("text")]
+        )
+        return values
+
+    def _mdui_uiinfo_lookup_elements_by_cls(self, uiinfos, element_cls):
+        elements = (
+            element
+            for uiinfo in uiinfos
+            for element_key, elements in uiinfo.items()
+            if element_key != "__class__"
+            for element in elements
+            if element.get("__class__") == element_cls
+        )
+        return elements
+
+    def _mdui_uiinfo_lookup_elements_by_key(self, uiinfos, element_key):
+        elements = (
+            element
+            for uiinfo in uiinfos
+            for elements in [uiinfo.get(element_key, [])]
+            for element in elements
+        )
+        return elements
+
+    def mdui_uiinfo_i18n_element_cls(self, entity_id, langpref, element_cls):
+        values = self._mdui_uiinfo_i18n_elements(
+            entity_id, langpref, element_cls, self._mdui_uiinfo_lookup_elements_by_cls
+        )
+        return values
+
+    def mdui_uiinfo_i18n_element_key(self, entity_id, langpref, element_key):
+        values = self._mdui_uiinfo_i18n_elements(
+            entity_id, langpref, element_key, self._mdui_uiinfo_lookup_elements_by_key
+        )
+        return values
+
+    def _mdui_uiinfo_display_name(self, entity_id, langpref="en"):
+        cls = classnames["mdui_uiinfo_display_name"]
+        values = self.mdui_uiinfo_i18n_element_cls(entity_id, langpref, cls)
+        return values
+
+    def mdui_uiinfo_display_name(self, entity_id, langpref="en"):
+        values = list(self._mdui_uiinfo_display_name(entity_id, langpref))
+        return values
+
+    def _mdui_uiinfo_description(self, entity_id, langpref="en"):
+        cls = classnames["mdui_uiinfo_description"]
+        values = self.mdui_uiinfo_i18n_element_cls(entity_id, langpref, cls)
+        return values
+
+    def mdui_uiinfo_description(self, entity_id, langpref="en"):
+        values = list(self._mdui_uiinfo_description(entity_id, langpref))
+        return values
+
+    def _mdui_uiinfo_information_url(self, entity_id, langpref="en"):
+        cls = classnames["mdui_uiinfo_information_url"]
+        values = self.mdui_uiinfo_i18n_element_cls(entity_id, langpref, cls)
+        return values
+
+    def mdui_uiinfo_information_url(self, entity_id, langpref="en"):
+        values = list(self._mdui_uiinfo_information_url(entity_id, langpref))
+        return values
+
+    def _mdui_uiinfo_privacy_statement_url(self, entity_id, langpref="en"):
+        cls = classnames["mdui_uiinfo_privacy_statement_url"]
+        values = self.mdui_uiinfo_i18n_element_cls(entity_id, langpref, cls)
+        return values
+
+    def mdui_uiinfo_privacy_statement_url(self, entity_id, langpref="en"):
+        values = list(self._mdui_uiinfo_privacy_statement_url(entity_id, langpref))
+        return values
+
+    def _mdui_uiinfo_logo(self, entity_id, width, height):
+        uiinfos = self._mdui_uiinfo(entity_id)
+        cls = classnames["mdui_uiinfo_logo"]
+        elements = self._mdui_uiinfo_lookup_elements_by_cls(uiinfos, cls)
+        values = (
+            element
+            for element in elements
+            if width is None or element.get("width") == width
+            if height is None or element.get("height") == height
+        )
+        return values
+
+    def mdui_uiinfo_logo(self, entity_id, width=None, height=None):
+        values = list(self._mdui_uiinfo_logo(entity_id, width, height))
+        return values
+
     def bindings(self, entity_id, typ, service):
         for _md in self.metadata.values():
             if entity_id in _md.items():