Added support for entity categories.

17 files changed, 364 insertions, 72 deletions

diff --git a/example/idp2/idp_conf.py b/example/idp2/idp_conf.py
index 4a291164..d6a95145 100644
--- a/example/idp2/idp_conf.py
+++ b/example/idp2/idp_conf.py
@@ -22,12 +22,12 @@ else:
 #BASE = "http://lingon.catalogix.se:8088"
 BASE = "http://localhost:8088"
 
-CONFIG={
-    "entityid" : "%s/idp.xml" % BASE,
+CONFIG = {
+    "entityid": "%s/idp.xml" % BASE,
     "description": "My IDP",
     "service": {
         "aa": {
-            "endpoints" : {
+            "endpoints": {
                 "attribute_service": [
                     ("%s/attr" % BASE, BINDING_SOAP)
                 ]
@@ -36,16 +36,16 @@ CONFIG={
                                NAMEID_FORMAT_PERSISTENT]
         },
         "aq": {
-            "endpoints" : {
+            "endpoints": {
                 "authn_query_service": [
                     ("%s/aqs" % BASE, BINDING_SOAP)
                 ]
             },
         },
         "idp": {
-            "name" : "Rolands IdP",
-            "endpoints" : {
-                "single_sign_on_service" : [
+            "name": "Rolands IdP",
+            "endpoints": {
+                "single_sign_on_service": [
                     ("%s/sso/redirect" % BASE, BINDING_HTTP_REDIRECT),
                     ("%s/sso/post" % BASE, BINDING_HTTP_POST),
                     ("%s/sso/art" % BASE, BINDING_HTTP_ARTIFACT),
@@ -56,27 +56,28 @@ CONFIG={
                     ("%s/slo/post" % BASE, BINDING_HTTP_POST),
                     ("%s/slo/redirect" % BASE, BINDING_HTTP_REDIRECT)
                 ],
-                "artifact_resolve_service":[
+                "artifact_resolve_service": [
                     ("%s/ars" % BASE, BINDING_SOAP)
                 ],
                 "assertion_id_request_service": [
                     ("%s/airs" % BASE, BINDING_URI)
                 ],
-                "manage_name_id_service":[
+                "manage_name_id_service": [
                     ("%s/mni/soap" % BASE, BINDING_SOAP),
                     ("%s/mni/post" % BASE, BINDING_HTTP_POST),
                     ("%s/mni/redirect" % BASE, BINDING_HTTP_REDIRECT),
                     ("%s/mni/art" % BASE, BINDING_HTTP_ARTIFACT)
                 ],
-                "name_id_mapping_service":[
+                "name_id_mapping_service": [
                     ("%s/nim" % BASE, BINDING_SOAP),
                 ],
             },
             "policy": {
                 "default": {
-                    "lifetime": {"minutes":15},
+                    "lifetime": {"minutes": 15},
                     "attribute_restrictions": None, # means all I have
-                    "name_form": NAME_FORMAT_URI
+                    "name_form": NAME_FORMAT_URI,
+                    "entity_categories": ["swami", "edugain"]
                 },
             },
             "subject_data": "./idp.subject",
@@ -84,10 +85,10 @@ CONFIG={
                                NAMEID_FORMAT_PERSISTENT]
         },
     },
-    "debug" : 1,
-    "key_file" : "pki/mykey.pem",
-    "cert_file" : "pki/mycert.pem",
-    "metadata" : {
+    "debug": 1,
+    "key_file": "pki/mykey.pem",
+    "cert_file": "pki/mycert.pem",
+    "metadata": {
         "local": ["../sp/sp.xml"],
     },
     "organization": {
@@ -95,27 +96,28 @@ CONFIG={
         "name": "Rolands Identiteter",
         "url": "http://www.example.com",
     },
-    "contact_person": [{
-        "contact_type": "technical",
-        "given_name": "Roland",
-        "sur_name": "Hedberg",
-        "email_address": "technical@example.com"
-    },{
-        "contact_type": "support",
-        "given_name": "Support",
-        "email_address": "support@example.com"
-    },
+    "contact_person": [
+        {
+            "contact_type": "technical",
+            "given_name": "Roland",
+            "sur_name": "Hedberg",
+            "email_address": "technical@example.com"
+        }, {
+            "contact_type": "support",
+            "given_name": "Support",
+            "email_address": "support@example.com"
+        },
     ],
     # This database holds the map between a subjects local identifier and
     # the identifier returned to a SP
     "xmlsec_binary": xmlsec_path,
-    "attribute_map_dir" : "../attributemaps",
+    "attribute_map_dir": "../attributemaps",
     "logger": {
         "rotating": {
             "filename": "idp.log",
             "maxBytes": 500000,
             "backupCount": 5,
-            },
+        },
         "loglevel": "debug",
     }
 }
diff --git a/setup.py b/setup.py
index a0f6767c..9a7bebee 100755
--- a/setup.py
+++ b/setup.py
@@ -73,7 +73,8 @@ setup(
 
     packages=['saml2', 'xmldsig', 'xmlenc', 's2repoze', 's2repoze.plugins',
               "saml2/profile", "saml2/schema", "saml2/extension",
-              "saml2/attributemaps", "saml2/authn_context"],
+              "saml2/attributemaps", "saml2/authn_context",
+              "saml2/entity_category"],
 
     package_dir={'': 'src'},
     package_data={'': ['xml/*.xml']},
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()
 
diff --git a/tests/servera_conf.py b/tests/servera_conf.py
index a79cbc4d..7aec90da 100644
--- a/tests/servera_conf.py
+++ b/tests/servera_conf.py
@@ -11,9 +11,9 @@ from pathutils import full_path, xmlsec_path
 
 BASE = "http://lingon.catalogix.se:8087"
 
-CONFIG={
-    "entityid" : "urn:mace:example.com:saml:roland:sp",
-    "name" : "urn:mace:example.com:saml:roland:sp",
+CONFIG = {
+    "entityid": "urn:mace:example.com:saml:roland:sp",
+    "name": "urn:mace:example.com:saml:roland:sp",
     "description": "My own SP",
     "service": {
         "sp": {
@@ -22,10 +22,10 @@ CONFIG={
                     ("%s/" % BASE, BINDING_HTTP_POST),
                     ("%s/paos" % BASE, BINDING_PAOS),
                     ("%s/redirect" % BASE, BINDING_HTTP_REDIRECT)],
-                "artifact_resolution_service":[
+                "artifact_resolution_service": [
                     ("%s/ars" % BASE, BINDING_SOAP)
                 ],
-                "manage_name_id_service":[
+                "manage_name_id_service": [
                     ("%s/mni/soap" % BASE, BINDING_SOAP),
                     ("%s/mni/post" % BASE, BINDING_HTTP_POST),
                     ("%s/mni/redirect" % BASE, BINDING_HTTP_REDIRECT),
@@ -34,26 +34,27 @@ CONFIG={
                 "single_logout_service": [
                     ("%s/sls" % BASE, BINDING_SOAP)
                 ],
-                "discovery_response":[
+                "discovery_response": [
                     ("%s/disco" % BASE, BINDING_DISCO)
                 ]
             },
             "required_attributes": ["surName", "givenName", "mail"],
             "optional_attributes": ["title", "eduPersonAffiliation"],
             "idp": ["urn:mace:example.com:saml:roland:idp"],
-            "name_id_format":[NAMEID_FORMAT_TRANSIENT, NAMEID_FORMAT_PERSISTENT]
+            "name_id_format": [NAMEID_FORMAT_TRANSIENT,
+                               NAMEID_FORMAT_PERSISTENT]
         }
     },
     "debug": 1,
     "key_file": full_path("test.key"),
     "cert_file": full_path("test.pem"),
     "ca_certs": full_path("cacerts.txt"),
-    "xmlsec_binary" : xmlsec_path,
+    "xmlsec_binary": xmlsec_path,
     "metadata": {
         "local": [full_path("idp_all.xml"), full_path("vo_metadata.xml")],
     },
     "virtual_organization": {
-        "urn:mace:example.com:it:tek":{
+        "urn:mace:example.com:it:tek": {
             "nameid_format": "urn:oid:1.3.6.1.4.1.1466.115.121.1.15-NameID",
             "common_identifier": "umuselin",
         }
@@ -61,12 +62,15 @@ CONFIG={
     "subject_data": "subject_data.db",
     "accepted_time_diff": 60,
     "attribute_map_dir": full_path("attributemaps"),
+    "entity_category": ["http://www.swamid.se/category/sfs-1993-1153",
+                        #"http://www.swamid.se/category/research-and-education",
+                        "http://www.swamid.se/category/hei-service"],
     #"valid_for": 6,
     "organization": {
         "name": ("AB Exempel", "se"),
         "display_name": ("AB Exempel", "se"),
         "url": "http://www.example.org",
-        },
+    },
     "contact_person": [
         {
             "given_name": "Roland",
diff --git a/tests/test_20_assertion.py b/tests/test_20_assertion.py
index 043b0d3c..be04b45a 100644
--- a/tests/test_20_assertion.py
+++ b/tests/test_20_assertion.py
@@ -731,5 +731,24 @@ def test_assertion_with_noop_attribute_conv():
             assert attr.attribute_value[0].text == "Roland"
 
 
+def test_filter_ava_5():
+    policy = Policy({
+        "default": {
+            "lifetime": {"minutes": 15},
+            #"attribute_restrictions": None  # means all I have
+            "entity_categories": ["swami", "edugain"]
+        }
+    })
+
+    ava = {"givenName": ["Derek"], "surName": ["Jeter"],
+           "mail": ["derek@nyy.mlb.com", "dj@example.com"]}
+
+    # No restrictions apply
+    ava = policy.filter(ava, "urn:mace:example.com:saml:curt:sp", [], [])
+
+    assert _eq(ava.keys(), ['mail', 'givenName', 'surName'])
+    assert _eq(ava["mail"], ["derek@nyy.mlb.com", "dj@example.com"])
+
+
 if __name__ == "__main__":
     test_assertion_with_noop_attribute_conv()
\ No newline at end of file
diff --git a/tests/test_30_mdstore.py b/tests/test_30_mdstore.py
index 78ea7e9c..f300a94f 100644
--- a/tests/test_30_mdstore.py
+++ b/tests/test_30_mdstore.py
@@ -12,7 +12,6 @@ from saml2 import BINDING_HTTP_REDIRECT
 from saml2 import BINDING_HTTP_POST
 from saml2 import BINDING_HTTP_ARTIFACT
 from saml2 import saml
-from saml2 import sigver
 from saml2 import config
 from saml2.attribute_converter import ac_factory
 from saml2.attribute_converter import d_to_local_name
diff --git a/tests/test_37_entity_categories.py b/tests/test_37_entity_categories.py
new file mode 100644
index 00000000..739725eb
--- /dev/null
+++ b/tests/test_37_entity_categories.py
@@ -0,0 +1,105 @@
+from saml2 import saml, sigver
+from saml2 import md
+from saml2 import config
+from saml2.assertion import Policy
+from saml2.attribute_converter import ac_factory
+from saml2.extension import mdui
+from saml2.extension import idpdisc
+from saml2.extension import dri
+from saml2.extension import mdattr
+from saml2.extension import ui
+from pathutils import full_path
+from saml2.mdstore import MetadataStore
+import xmldsig
+import xmlenc
+
+ONTS = {
+    saml.NAMESPACE: saml,
+    mdui.NAMESPACE: mdui,
+    mdattr.NAMESPACE: mdattr,
+    dri.NAMESPACE: dri,
+    ui.NAMESPACE: ui,
+    idpdisc.NAMESPACE: idpdisc,
+    md.NAMESPACE: md,
+    xmldsig.NAMESPACE: xmldsig,
+    xmlenc.NAMESPACE: xmlenc
+}
+
+ATTRCONV = ac_factory(full_path("attributemaps"))
+sec_config = config.Config()
+sec_config.xmlsec_binary = sigver.get_xmlsec_binary(["/opt/local/bin"])
+
+__author__ = 'rolandh'
+
+MDS = MetadataStore(ONTS.values(), ATTRCONV, sec_config,
+                    disable_ssl_certificate_validation=True)
+MDS.imp({"mdfile": [full_path("swamid.md")]})
+
+
+def _eq(l1, l2):
+    return set(l1) == set(l2)
+
+
+def test_filter_ava():
+    policy = Policy({
+        "default": {
+            "lifetime": {"minutes": 15},
+            #"attribute_restrictions": None  # means all I have
+            "entity_categories": ["swamid"]
+        }
+    })
+
+    ava = {"givenName": ["Derek"], "surname": ["Jeter"],
+           "email": ["derek@nyy.mlb.com", "dj@example.com"], "c": ["USA"]}
+
+    ava = policy.filter(ava, "https://connect.sunet.se/shibboleth", MDS)
+
+    assert _eq(ava.keys(), ['email', 'givenName', 'surname', 'c'])
+    assert _eq(ava["email"], ["derek@nyy.mlb.com", "dj@example.com"])
+
+
+def test_filter_ava2():
+    policy = Policy({
+        "default": {
+            "lifetime": {"minutes": 15},
+            #"attribute_restrictions": None  # means all I have
+            "entity_categories": ["edugain"]
+        }
+    })
+
+    ava = {"givenName": ["Derek"], "surname": ["Jeter"],
+           "email": ["derek@nyy.mlb.com"], "c": ["USA"],
+           "eduPersonTargetedID": "foo!bar!xyz"}
+
+    ava = policy.filter(ava, "https://connect.sunet.se/shibboleth", MDS)
+
+    # Mismatch, policy deals with eduGAIN, metadata says SWAMID
+    # So only minimum should come out
+    assert _eq(ava.keys(), ['eduPersonTargetedID'])
+
+
+def test_filter_ava3():
+    policy = Policy({
+        "default": {
+            "lifetime": {"minutes": 15},
+            #"attribute_restrictions": None  # means all I have
+            "entity_categories": ["swamid"]
+        }
+    })
+
+    mds = MetadataStore(ONTS.values(), ATTRCONV, sec_config,
+                        disable_ssl_certificate_validation=True)
+    mds.imp({"local": [full_path("entity_cat_sfs_hei.xml")]})
+
+    ava = {"givenName": ["Derek"], "surname": ["Jeter"],
+           "email": ["derek@nyy.mlb.com"], "c": ["USA"],
+           "eduPersonTargetedID": "foo!bar!xyz",
+           "norEduPersonNIN": "19800101134"}
+
+    ava = policy.filter(ava, "urn:mace:example.com:saml:roland:sp", mds)
+
+    assert _eq(ava.keys(), ['eduPersonTargetedID', "norEduPersonNIN"])
+
+
+if __name__ == "__main__":
+    test_filter_ava2()
\ No newline at end of file
diff --git a/tools/make_metadata.py b/tools/make_metadata.py
index 68110ec1..1eaaca05 100755
--- a/tools/make_metadata.py
+++ b/tools/make_metadata.py
@@ -6,7 +6,7 @@ from saml2.metadata import entity_descriptor
 from saml2.metadata import entities_descriptor
 from saml2.metadata import sign_entity_descriptor
 
-from saml2.sigver import SecurityContext
+from saml2.sigver import SecurityContext, CryptoBackendXmlSec1
 from saml2.sigver import get_xmlsec_cryptobackend
 from saml2.sigver import get_xmlsec_binary
 from saml2.validate import valid_instance
@@ -61,9 +61,13 @@ for filespec in args.config:
     cnf = Config().load_file(fil, metadata_construction=True)
     eds.append(entity_descriptor(cnf))
 
-crypto = get_xmlsec_cryptobackend()
-secc = SecurityContext(crypto, key_file=args.keyfile,
-                              cert_file=args.cert, debug=1)
+if not xmlsec:
+    crypto = get_xmlsec_cryptobackend()
+else:
+    crypto = CryptoBackendXmlSec1(xmlsec)
+
+secc = SecurityContext(crypto, key_file=args.keyfile, cert_file=args.cert,
+                       debug=1)
 
 if args.id:
     desc = entities_descriptor(eds, valid_for, args.name, args.id,