From d9038fe0adba119d45acc88353a543ce42c644fc Mon Sep 17 00:00:00 2001
From: Fredrik Thulin <fredrik@thulin.net>
Date: Thu, 6 Mar 2014 14:22:36 +0100
Subject: [PATCH] Add entity_attributes().

eduID needs to be able to get all entity attributes for an entity id,
not just the entity categories.

Conflicts:
	src/saml2/mdstore.py
---
 src/saml2/mdstore.py | 71 +++++++++++++++++++++++++++-----------------
 1 file changed, 44 insertions(+), 27 deletions(-)

diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index 0a7726e..223f9a5 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -775,37 +775,54 @@ class MetadataStore(object):
         return [m["text"] for m in ad["affiliate_member"]]
 
     def entity_categories(self, entity_id):
-        ent = self.__getitem__(entity_id)
-        res = []
-        try:
-            ext = ent["extensions"]
-        except KeyError:
-            pass
-        else:
-            for elem in ext["extension_elements"]:
-                if elem["__class__"] == ENTITYATTRIBUTES:
-                    for attr in elem["attribute"]:
-                        if attr["name"] == ENTITY_CATEGORY:
-                            res.extend([v["text"] for v in
-                                        attr["attribute_value"]])
+        """
+        Get a list of entity categories for an entity id.
 
-        return res
+        :param entity_id: Entity id
+        :return: Entity categories
+
+        :type entity_id: string
+        :rtype: [string]
+        """
+        attributes = self.entity_attributes(entity_id)
+        return attributes.get(ENTITY_CATEGORY, [])
 
     def supported_entity_categories(self, entity_id):
-        ent = self.__getitem__(entity_id)
-        res = []
-        try:
-            ext = ent["extensions"]
-        except KeyError:
-            pass
-        else:
-            for elem in ext["extension_elements"]:
-                if elem["__class__"] == ENTITYATTRIBUTES:
-                    for attr in elem["attribute"]:
-                        if attr["name"] == ENTITY_CATEGORY_SUPPORT:
-                            res.extend([v["text"] for v in
-                                        attr["attribute_value"]])
+        """
+        Get a list of entity category support for an entity id.
 
+        :param entity_id: Entity id
+        :return: Entity category support
+
+        :type entity_id: string
+        :rtype: [string]
+        """
+        attributes = self.entity_attributes(entity_id)
+        return attributes.get(ENTITY_CATEGORY_SUPPORT, [])
+
+    def entity_attributes(self, entity_id):
+        """
+        Get all entity attributes for an entry in the metadata.
+
+        Example return data:
+
+        {'http://macedir.org/entity-category': ['something', 'something2'],
+         'http://example.org/saml-foo': ['bar']}
+
+        :param entity_id: Entity id
+        :return: dict with keys and value-lists from metadata
+
+        :type entity_id: string
+        :rtype: dict
+        """
+        ext = self.__getitem__(entity_id)["extensions"]
+        res = {}
+        for elem in ext["extension_elements"]:
+            if elem["__class__"] == ENTITYATTRIBUTES:
+                for attr in elem["attribute"]:
+                    if attr["name"] not in res:
+                        res[attr["name"]] = []
+                    res[attr["name"]] += [v["text"] for v in attr["attribute_value"]]
         return res
 
     def bindings(self, entity_id, typ, service):