Merge pull request #613 from skoranda/more_flexible_entity_category_import

Make entity category imports more flexible

4 files changed, 144 insertions, 2 deletions

diff --git a/src/saml2/assertion.py b/src/saml2/assertion.py
index 8984db59..ed043d9a 100644
--- a/src/saml2/assertion.py
+++ b/src/saml2/assertion.py
@@ -353,8 +353,11 @@ class Policy(object):
             else:
                 ecs = []
                 for cat in items:
-                    _mod = importlib.import_module(
-                        "saml2.entity_category.%s" % cat)
+                    try:
+                        _mod = importlib.import_module(cat)
+                    except ImportError:
+                        _mod = importlib.import_module(
+                            "saml2.entity_category.%s" % cat)
                     _ec = {}
                     for key, items in _mod.RELEASE.items():
                         alist = [k.lower() for k in items]
diff --git a/tests/entity_cat_rs.xml b/tests/entity_cat_rs.xml
new file mode 100644
index 00000000..5f3e00f8
--- /dev/null
+++ b/tests/entity_cat_rs.xml
@@ -0,0 +1,84 @@
+<?xml version='1.0' encoding='UTF-8'?>
+<ns0:EntityDescriptor xmlns:ns0="urn:oasis:names:tc:SAML:2.0:metadata"
+                      xmlns:xs="http://www.w3.org/2001/XMLSchema"
+                      xmlns:ns1="urn:oasis:names:tc:SAML:metadata:attribute"
+                      xmlns:ns2="urn:oasis:names:tc:SAML:2.0:assertion"
+                      xmlns:ns5="http://www.w3.org/2000/09/xmldsig#"
+                      xmlns:ns4="urn:oasis:names:tc:SAML:profiles:SSO:idp-discovery-protocol"
+                      xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+                      entityID="urn:mace:example.com:saml:roland:sp">
+    <ns0:Extensions>
+        <ns1:EntityAttributes>
+            <ns2:Attribute Name="http://macedir.org/entity-category">
+                <ns2:AttributeValue xsi:type="xs:string">
+                    http://refeds.org/category/research-and-scholarship
+                </ns2:AttributeValue>
+            </ns2:Attribute>
+        </ns1:EntityAttributes>
+    </ns0:Extensions>
+    <ns0:SPSSODescriptor AuthnRequestsSigned="false" WantAssertionsSigned="true"
+                         protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">
+        <ns0:Extensions>
+            <ns4:DiscoveryResponse
+                    Binding="urn:oasis:names:tc:SAML:profiles:SSO:idp-discovery-protocol"
+                    Location="https://xenosmilus2.umdc.umu.se:8086/disco"
+                    index="1"/>
+        </ns0:Extensions>
+        <ns0:KeyDescriptor use="encryption">
+            <ns5:KeyInfo>
+                <ns5:X509Data>
+                    <ns5:X509Certificate>
+                        MIIC8jCCAlugAwIBAgIJAJHg2V5J31I8MA0GCSqGSIb3DQEBBQUAMFoxCzAJBgNV
+                        BAYTAlNFMQ0wCwYDVQQHEwRVbWVhMRgwFgYDVQQKEw9VbWVhIFVuaXZlcnNpdHkx
+                        EDAOBgNVBAsTB0lUIFVuaXQxEDAOBgNVBAMTB1Rlc3QgU1AwHhcNMDkxMDI2MTMz
+                        MTE1WhcNMTAxMDI2MTMzMTE1WjBaMQswCQYDVQQGEwJTRTENMAsGA1UEBxMEVW1l
+                        YTEYMBYGA1UEChMPVW1lYSBVbml2ZXJzaXR5MRAwDgYDVQQLEwdJVCBVbml0MRAw
+                        DgYDVQQDEwdUZXN0IFNQMIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDkJWP7
+                        bwOxtH+E15VTaulNzVQ/0cSbM5G7abqeqSNSs0l0veHr6/ROgW96ZeQ57fzVy2MC
+                        FiQRw2fzBs0n7leEmDJyVVtBTavYlhAVXDNa3stgvh43qCfLx+clUlOvtnsoMiiR
+                        mo7qf0BoPKTj7c0uLKpDpEbAHQT4OF1HRYVxMwIDAQABo4G/MIG8MB0GA1UdDgQW
+                        BBQ7RgbMJFDGRBu9o3tDQDuSoBy7JjCBjAYDVR0jBIGEMIGBgBQ7RgbMJFDGRBu9
+                        o3tDQDuSoBy7JqFepFwwWjELMAkGA1UEBhMCU0UxDTALBgNVBAcTBFVtZWExGDAW
+                        BgNVBAoTD1VtZWEgVW5pdmVyc2l0eTEQMA4GA1UECxMHSVQgVW5pdDEQMA4GA1UE
+                        AxMHVGVzdCBTUIIJAJHg2V5J31I8MAwGA1UdEwQFMAMBAf8wDQYJKoZIhvcNAQEF
+                        BQADgYEAMuRwwXRnsiyWzmRikpwinnhTmbooKm5TINPE7A7gSQ710RxioQePPhZO
+                        zkM27NnHTrCe2rBVg0EGz7QTd1JIwLPvgoj4VTi/fSha/tXrYUaqc9AqU1kWI4WN
+                        +vffBGQ09mo+6CffuFTZYeOhzP/2stAPwCTU4kxEoiy0KpZMANI=
+                    </ns5:X509Certificate>
+                </ns5:X509Data>
+            </ns5:KeyInfo>
+        </ns0:KeyDescriptor>
+        <ns0:KeyDescriptor use="signing">
+            <ns5:KeyInfo>
+                <ns5:X509Data>
+                    <ns5:X509Certificate>
+                        MIIC8jCCAlugAwIBAgIJAJHg2V5J31I8MA0GCSqGSIb3DQEBBQUAMFoxCzAJBgNV
+                        BAYTAlNFMQ0wCwYDVQQHEwRVbWVhMRgwFgYDVQQKEw9VbWVhIFVuaXZlcnNpdHkx
+                        EDAOBgNVBAsTB0lUIFVuaXQxEDAOBgNVBAMTB1Rlc3QgU1AwHhcNMDkxMDI2MTMz
+                        MTE1WhcNMTAxMDI2MTMzMTE1WjBaMQswCQYDVQQGEwJTRTENMAsGA1UEBxMEVW1l
+                        YTEYMBYGA1UEChMPVW1lYSBVbml2ZXJzaXR5MRAwDgYDVQQLEwdJVCBVbml0MRAw
+                        DgYDVQQDEwdUZXN0IFNQMIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDkJWP7
+                        bwOxtH+E15VTaulNzVQ/0cSbM5G7abqeqSNSs0l0veHr6/ROgW96ZeQ57fzVy2MC
+                        FiQRw2fzBs0n7leEmDJyVVtBTavYlhAVXDNa3stgvh43qCfLx+clUlOvtnsoMiiR
+                        mo7qf0BoPKTj7c0uLKpDpEbAHQT4OF1HRYVxMwIDAQABo4G/MIG8MB0GA1UdDgQW
+                        BBQ7RgbMJFDGRBu9o3tDQDuSoBy7JjCBjAYDVR0jBIGEMIGBgBQ7RgbMJFDGRBu9
+                        o3tDQDuSoBy7JqFepFwwWjELMAkGA1UEBhMCU0UxDTALBgNVBAcTBFVtZWExGDAW
+                        BgNVBAoTD1VtZWEgVW5pdmVyc2l0eTEQMA4GA1UECxMHSVQgVW5pdDEQMA4GA1UE
+                        AxMHVGVzdCBTUIIJAJHg2V5J31I8MAwGA1UdEwQFMAMBAf8wDQYJKoZIhvcNAQEF
+                        BQADgYEAMuRwwXRnsiyWzmRikpwinnhTmbooKm5TINPE7A7gSQ710RxioQePPhZO
+                        zkM27NnHTrCe2rBVg0EGz7QTd1JIwLPvgoj4VTi/fSha/tXrYUaqc9AqU1kWI4WN
+                        +vffBGQ09mo+6CffuFTZYeOhzP/2stAPwCTU4kxEoiy0KpZMANI=
+                    </ns5:X509Certificate>
+                </ns5:X509Data>
+            </ns5:KeyInfo>
+        </ns0:KeyDescriptor>
+        <ns0:AssertionConsumerService
+                Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect"
+                Location="https://xenosmilus2.umdc.umu.se:8086/acs/sfs/re_nren/redirect"
+                index="1"/>
+        <ns0:AssertionConsumerService
+                Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST"
+                Location="https://xenosmilus2.umdc.umu.se:8086/acs/sfs/re_nren/post"
+                index="2"/>
+    </ns0:SPSSODescriptor>
+</ns0:EntityDescriptor>
diff --git a/tests/myentitycategory.py b/tests/myentitycategory.py
new file mode 100644
index 00000000..9ec55bf9
--- /dev/null
+++ b/tests/myentitycategory.py
@@ -0,0 +1,16 @@
+CUSTOM_R_AND_S = ['eduPersonTargetedID',
+           'eduPersonPrincipalName',
+           'mail',
+           'displayName',
+           'givenName',
+           'sn',
+           'eduPersonScopedAffiliation',
+           'eduPersonUniqueId'
+           ]
+
+RESEARCH_AND_SCHOLARSHIP = "http://refeds.org/category/research-and-scholarship"
+
+RELEASE = {
+    "": ["eduPersonTargetedID"],
+    RESEARCH_AND_SCHOLARSHIP: CUSTOM_R_AND_S,
+}
diff --git a/tests/test_37_entity_categories.py b/tests/test_37_entity_categories.py
index 625caaa1..839030fd 100644
--- a/tests/test_37_entity_categories.py
+++ b/tests/test_37_entity_categories.py
@@ -152,5 +152,44 @@ def test_idp_policy_filter():
             "eduPersonTargetedID"]  # because no entity category
 
 
+def test_entity_category_import_from_path():
+    # The entity category module myentitycategory.py is in the tests
+    # directory which is on the standard module search path.
+    # The module uses a custom interpretation of the REFEDs R&S entity category
+    # by adding eduPersonUniqueId.
+    policy = Policy({
+        "default": {
+            "lifetime": {"minutes": 15},
+            "entity_categories": ["myentitycategory"]
+        }
+    })
+
+    mds = MetadataStore(ATTRCONV, sec_config,
+                        disable_ssl_certificate_validation=True)
+
+    # The file entity_cat_rs.xml contains the SAML metadata for an SP
+    # tagged with the REFEDs R&S entity category.
+    mds.imp([{"class": "saml2.mdstore.MetaDataFile",
+              "metadata": [(full_path("entity_cat_rs.xml"),)]}])
+
+    ava = {"givenName": ["Derek"], "sn": ["Jeter"],
+           "displayName": "Derek Jeter",
+           "mail": ["derek@nyy.mlb.com"], "c": ["USA"],
+           "eduPersonTargetedID": "foo!bar!xyz",
+           "eduPersonUniqueId": "R13ET7UD68K0HGR153KE@my.org",
+           "eduPersonScopedAffiliation": "member@my.org",
+           "eduPersonPrincipalName": "user01@my.org",
+           "norEduPersonNIN": "19800101134"}
+
+    ava = policy.filter(ava, "urn:mace:example.com:saml:roland:sp", mds)
+
+    # We expect c and norEduPersonNIN to be filtered out since they are not
+    # part of the custom entity category.
+    assert _eq(list(ava.keys()),
+               ["eduPersonTargetedID", "eduPersonPrincipalName",
+                "eduPersonUniqueId", "displayName", "givenName",
+                "eduPersonScopedAffiliation", "mail", "sn"])
+
+
 if __name__ == "__main__":
     test_filter_ava3()