Added support for entity categories.

10 files changed, 186 insertions, 28 deletions

diff --git a/src/saml2/assertion.py b/src/saml2/assertion.py
index 9a6ea93c..dcfae4c2 100644
--- a/src/saml2/assertion.py
+++ b/src/saml2/assertion.py
@@ -14,6 +14,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import importlib
 import logging
 
 import re
@@ -21,14 +22,15 @@ from saml2.saml import NAME_FORMAT_URI
 import xmlenc
 
 from saml2 import saml
+from saml2 import entity_category
 
 from saml2.time_util import instant, in_a_while
 from saml2.attribute_converter import from_local
-
 from saml2.s_utils import sid, MissingValue
 from saml2.s_utils import factory
 from saml2.s_utils import assertion_factory
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -286,7 +288,19 @@ class Policy(object):
         for _, spec in self._restrictions.items():
             if spec is None:
                 continue
-            
+
+            try:
+                _entcat = spec["entity_categories"]
+            except KeyError:
+                pass
+            else:
+                ecs = []
+                for cat in _entcat:
+                    _mod = importlib.import_module(
+                        "saml2.entity_category.%s" % cat)
+                    ecs.append(_mod.RELEASE)
+                spec["entity_categories"] = ecs
+
             try:
                 restr = spec["attribute_restrictions"]
             except KeyError:
@@ -383,7 +397,53 @@ class Policy(object):
             restrictions = None
         
         return restrictions
-    
+
+    def get_entity_categories_restriction(self, sp_entity_id, mds):
+        if not self._restrictions:
+            return None
+
+        restrictions = {}
+        ec_maps = []
+        try:
+            try:
+                ec_maps = self._restrictions[sp_entity_id]["entity_categories"]
+            except KeyError:
+                try:
+                    ec_maps = self._restrictions["default"]["entity_categories"]
+                except KeyError:
+                    pass
+        except KeyError:
+            pass
+
+        if ec_maps:
+            # always released
+            for ec_map in ec_maps:
+                try:
+                    attrs = ec_map[""]
+                except KeyError:
+                    pass
+                else:
+                    for attr in attrs:
+                        restrictions[attr] = None
+
+            try:
+                ecs = mds.entity_categories(sp_entity_id)
+            except KeyError:
+                pass
+            else:
+                for ec in ecs:
+                    for ec_map in ec_maps:
+                        try:
+                            attrs = ec_map[ec]
+                        except KeyError:
+                            pass
+                        else:
+                            for attr in attrs:
+                                restrictions[attr] = None
+
+        return restrictions
+
+
     def not_on_or_after(self, sp_entity_id):
         """ When the assertion stops being valid, should not be
         used after this time.
@@ -394,7 +454,7 @@ class Policy(object):
         
         return in_a_while(**self.get_lifetime(sp_entity_id))
     
-    def filter(self, ava, sp_entity_id, required=None, optional=None):
+    def filter(self, ava, sp_entity_id, mdstore, required=None, optional=None):
         """ What attribute and attribute values returns depends on what
         the SP has said it wants in the request or in the metadata file and
         what the IdP/AA wants to release. An assumption is that what the SP
@@ -408,8 +468,11 @@ class Policy(object):
         :return: A possibly modified AVA
         """
 
-        ava = filter_attribute_value_assertions(
-            ava, self.get_attribute_restriction(sp_entity_id))
+        _rest = self.get_attribute_restriction(sp_entity_id)
+        if _rest is None:
+            _rest = self.get_entity_categories_restriction(sp_entity_id,
+                                                           mdstore)
+        ava = filter_attribute_value_assertions(ava, _rest)
         
         if required or optional:
             ava = filter_on_attributes(ava, required, optional)
@@ -427,8 +490,8 @@ class Policy(object):
         if metadata:
             spec = metadata.attribute_requirement(sp_entity_id)
             if spec:
-                return self.filter(ava, sp_entity_id, spec["required"],
-                                   spec["optional"])
+                ava = self.filter(ava, sp_entity_id, metadata,
+                                  spec["required"], spec["optional"])
 
         return self.filter(ava, sp_entity_id, [], [])
 
@@ -447,19 +510,6 @@ class Policy(object):
                            audience=factory(saml.Audience,
                                             text=sp_entity_id))])
 
-NAME = ["givenName", "surname", "initials", "displayName", "schacSn1",
-        "schacSn2"]
-STATIC_ORG_INFO = ["organizationName", ""]
-
-RESEARCH_AND_EDUCATION = "http://www.swamid.se/category/research-and-education"
-SFS_1993_1153 = "http://www.swamid.se/category/sfs-1993-1153"
-
-# EC_RELEASE = {
-#     "eduPersonPrincipalName", "eduPersonTargetedID", "mail", "email",
-#         "eduPersonScopedAffiliation"
-#     ]),
-#     "http://www.swamid.se/category/sfs-1993-1153": ["norEduPersonNIN"]
-# }
 
 
 class EntityCategories(object):
diff --git a/src/saml2/config.py b/src/saml2/config.py
index 8edd8b2f..fb0f440d 100644
--- a/src/saml2/config.py
+++ b/src/saml2/config.py
@@ -60,7 +60,8 @@ COMMON_ARGS = [
     "logout_requests_signed",
     "disable_ssl_certificate_validation",
     "referred_binding",
-    "session_storage"
+    "session_storage",
+    "entity_category"
 ]
 
 SP_ARGS = [
@@ -189,6 +190,7 @@ class Config(object):
         self.preferred_binding = PREFERRED_BINDING
         self.domain = ""
         self.name_qualifier = ""
+        self.entity_category = ""
 
     def setattr(self, context, attr, val):
         if context == "":
diff --git a/src/saml2/discovery.py b/src/saml2/discovery.py
index 9b28ddce..383b32e9 100644
--- a/src/saml2/discovery.py
+++ b/src/saml2/discovery.py
@@ -7,6 +7,7 @@ __author__ = 'rolandh'
 
 IDPDISC_POLICY = "urn:oasis:names:tc:SAML:profiles:SSO:idp-discovery-protocol:single"
 
+
 class DiscoveryServer(Entity):
     def __init__(self, config=None, config_file=""):
         Entity.__init__(self, "disco", config, config_file)
@@ -65,7 +66,7 @@ class DiscoveryServer(Entity):
                                           returnIDParam="entityID",
                                           entity_id=None):
         if entity_id:
-            qp = urlencode({returnIDParam:entity_id})
+            qp = urlencode({returnIDParam: entity_id})
 
             part = urlparse(return_url)
             if part.query:
diff --git a/src/saml2/entity_category/__init__.py b/src/saml2/entity_category/__init__.py
new file mode 100644
index 00000000..d78e3b99
--- /dev/null
+++ b/src/saml2/entity_category/__init__.py
@@ -0,0 +1,14 @@
+__author__ = 'rolandh'
+
+ENTITYATTRIBUTES = "urn:oasis:names:tc:SAML:metadata:attribute&EntityAttributes"
+
+
+def entity_categories(md):
+    res = []
+    if "extensions" in md:
+        for elem in md["extensions"]["extension_elements"]:
+            if elem["__class__"] == ENTITYATTRIBUTES:
+                for attr in elem["attribute"]:
+                    res.append(attr["text"])
+
+    return res
\ No newline at end of file
diff --git a/src/saml2/entity_category/edugain.py b/src/saml2/entity_category/edugain.py
new file mode 100644
index 00000000..5a094a73
--- /dev/null
+++ b/src/saml2/entity_category/edugain.py
@@ -0,0 +1,10 @@
+__author__ = 'rolandh'
+
+COC = "http://www.edugain.org/dataprotection/coc-eu-01-draft"
+
+RELEASE = {
+    "": ["eduPersonTargetedID"],
+    COC: ["eduPersonPrincipalName", "eduPersonScopedAffiliation", "email",
+          "givenName", "surname", "displayName", "schacHomeOrganization"]
+}
+
diff --git a/src/saml2/entity_category/swamid.py b/src/saml2/entity_category/swamid.py
new file mode 100644
index 00000000..899086f7
--- /dev/null
+++ b/src/saml2/entity_category/swamid.py
@@ -0,0 +1,20 @@
+__author__ = 'rolandh'
+
+
+NAME = ["givenName", "surname", "initials", "displayName"]
+STATIC_ORG_INFO = ["c", "o", "ou"]
+OTHER = ["eduPersonPrincipalName", "eduPersonScopedAffiliation", "email"]
+
+RESEARCH_AND_EDUCATION = "http://www.swamid.se/category/research-and-education"
+SFS_1993_1153 = "http://www.swamid.se/category/sfs-1993-1153"
+
+EU = "http://www.swamid.se/category/eu-adequate-protection"
+NREN = "http://www.swamid.se/category/nren-service"
+HEI = "http://www.swamid.se/category/hei-service"
+
+RELEASE = {
+    "": ["eduPersonTargetedID"],
+    SFS_1993_1153: ["norEduPersonNIN"],
+    RESEARCH_AND_EDUCATION: NAME + STATIC_ORG_INFO + OTHER,
+}
+
diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index 31e0e913..9c50478d 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -46,6 +46,9 @@ REQ2SRV = {
     "discovery_service_request": "discovery_response"
 }
 
+
+ENTITYATTRIBUTES = "urn:oasis:names:tc:SAML:metadata:attribute&EntityAttributes"
+
 # ---------------------------------------------------
 
 
@@ -321,6 +324,16 @@ class MetaData(object):
 
         return res
 
+    def entity_categories(self, entity_id):
+        res = []
+        if "extensions" in self[entity_id]:
+            for elem in self[entity_id]["extensions"]["extension_elements"]:
+                if elem["__class__"] == ENTITYATTRIBUTES:
+                    for attr in elem["attribute"]:
+                        res.append(attr["text"])
+
+        return res
+
 
 class MetaDataFile(MetaData):
     """
@@ -649,6 +662,17 @@ class MetadataStore(object):
         ad = self.__getitem__(entity_id)["affiliation_descriptor"]
         return [m["text"] for m in ad["affiliate_member"]]
 
+    def entity_categories(self, entity_id):
+        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"] == "http://macedir.org/entity-category":
+                        res.extend([v["text"] for v in attr["attribute_value"]])
+
+        return res
+
     def bindings(self, entity_id, typ, service):
         for md in self.metadata.values():
             if entity_id in md.items():
diff --git a/src/saml2/metadata.py b/src/saml2/metadata.py
index 665f5689..c01713bb 100644
--- a/src/saml2/metadata.py
+++ b/src/saml2/metadata.py
@@ -1,9 +1,9 @@
 #!/usr/bin/env python
 from saml2.time_util import in_a_while
-from saml2.extension import mdui, idpdisc, shibmd
-from saml2.saml import NAME_FORMAT_URI
+from saml2.extension import mdui, idpdisc, shibmd, mdattr
+from saml2.saml import NAME_FORMAT_URI, AttributeValue, Attribute
 from saml2.attribute_converter import from_local_name
-from saml2 import md
+from saml2 import md, saml
 from saml2 import BINDING_HTTP_POST
 from saml2 import BINDING_HTTP_REDIRECT
 from saml2 import BINDING_SOAP
@@ -549,6 +549,14 @@ def entity_descriptor(confd):
     if confd.contact_person is not None:
         entd.contact_person = do_contact_person_info(confd.contact_person)
 
+    if confd.entity_category:
+        entd.extensions = md.Extensions()
+        ava = [AttributeValue(text=c) for c in confd.entity_category]
+        attr = Attribute(attribute_value=ava,
+                         name="http://macedir.org/entity-category")
+        item = mdattr.EntityAttributes(attribute=attr)
+        entd.extensions.add_extension_element(item)
+
     serves = confd.serves
     if not serves:
         raise Exception(
diff --git a/src/saml2/s_utils.py b/src/saml2/s_utils.py
index 457bf853..34552cb0 100644
--- a/src/saml2/s_utils.py
+++ b/src/saml2/s_utils.py
@@ -9,6 +9,8 @@ import sys
 import hmac
 
 # from python 2.5
+import imp
+
 if sys.version_info >= (2, 5):
     import hashlib
 else:  # before python 2.5
@@ -406,4 +408,31 @@ def fticks_log(sp, logf, idp_entity_id, user_id, secret, assertion):
         "PN": csum.hexdigest(),
         "AM": assertion.AuthnStatement.AuthnContext.AuthnContextClassRef.text
     }
-    logf.info(FTICKS_FORMAT % "#".join(["%s=%s" % (a,v) for a,v in info]))
\ No newline at end of file
+    logf.info(FTICKS_FORMAT % "#".join(["%s=%s" % (a,v) for a,v in info]))
+
+
+def dynamic_importer(name, class_name=None):
+    """
+    Dynamically imports modules / classes
+    """
+    try:
+        fp, pathname, description = imp.find_module(name)
+    except ImportError:
+        print "unable to locate module: " + name
+        return None, None
+
+    try:
+        package = imp.load_module(name, fp, pathname, description)
+    except Exception, e:
+        raise
+
+    if class_name:
+        try:
+            _class = imp.load_module("%s.%s" % (name, class_name), fp,
+                                      pathname, description)
+        except Exception, e:
+            raise
+
+        return package, _class
+    else:
+        return package, None
diff --git a/src/saml2/server.py b/src/saml2/server.py
index 06e88101..77f17a9a 100644
--- a/src/saml2/server.py
+++ b/src/saml2/server.py
@@ -344,7 +344,7 @@ class Server(Entity):
             ast = Assertion(identity)
             policy = self.config.getattr("policy", "aa")
             if policy:
-                ast.apply_policy(sp_entity_id, policy)
+                ast.apply_policy(sp_entity_id, policy, self.metadata)
             else:
                 policy = Policy()