Handled possible exception in the entity_categories method and add a supported_entity_categories method.

Author: Roland Hedberg

src/saml2/mdstore.py

@@ -56,6 +56,8 @@ class ToOld(Exception):
 
 
 ENTITYATTRIBUTES = "urn:oasis:names:tc:SAML:metadata:attribute&EntityAttributes"
+ENTITY_CATEGORY = "http://macedir.org/entity-category"
+ENTITY_CATEGORY_SUPPORT = "http://macedir.org/entity-category-support"
 
 # ---------------------------------------------------
 
@@ -598,30 +600,30 @@ def single_sign_on_service(self, entity_id, binding=None, typ="idpsso"):
         if binding is None:
             binding = BINDING_HTTP_REDIRECT
         return self.service(entity_id, "idpsso_descriptor",
-                             "single_sign_on_service", binding)
+                            "single_sign_on_service", binding)
 
     def name_id_mapping_service(self, entity_id, binding=None, typ="idpsso"):
         # IDP
         if binding is None:
             binding = BINDING_HTTP_REDIRECT
         return self.service(entity_id, "idpsso_descriptor",
-                             "name_id_mapping_service", binding)
+                            "name_id_mapping_service", binding)
 
     def authn_query_service(self, entity_id, binding=None,
                             typ="authn_authority"):
         # AuthnAuthority
         if binding is None:
             binding = BINDING_SOAP
         return self.service(entity_id, "authn_authority_descriptor",
-                             "authn_query_service", binding)
+                            "authn_query_service", binding)
 
     def attribute_service(self, entity_id, binding=None,
                           typ="attribute_authority"):
         # AttributeAuthority
         if binding is None:
             binding = BINDING_HTTP_REDIRECT
         return self.service(entity_id, "attribute_authority_descriptor",
-                             "attribute_service", binding)
+                            "attribute_service", binding)
 
     def authz_service(self, entity_id, binding=None, typ="pdp"):
         # PDP
@@ -774,13 +776,35 @@ def vo_members(self, entity_id):
 
     def entity_categories(self, entity_id):
         ent = self.__getitem__(entity_id)
-        ext = ent["extensions"]
         res = []
-        for elem in ext["extension_elements"]:
-            if elem["__class__"] == ENTITYATTRIBUTES:
-                for attr in elem["attribute"]:
-                    if attr["name"] == "http://macedir.org/entity-category":
-                        res.extend([v["text"] for v in attr["attribute_value"]])
+        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"]])
+
+        return res
+
+    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"]])
 
         return res