source: TI12-security/trunk/python/ndg.security.test/ndg/security/test/unit/saml/test_saml.py @ 5513

Subversion URL: http://proj.badc.rl.ac.uk/svn/ndg/TI12-security/trunk/python/ndg.security.test/ndg/security/test/unit/saml/test_saml.py@5513
Revision 5513, 7.4 KB checked in by pjkersha, 11 years ago (diff)

Split SAML unit tests into separate ones for SAML and SAML Attribute Authority interface.

Line 
1"""SAML unit test package
2
3NERC DataGrid Project
4"""
5__author__ = "P J Kershaw"
6__date__ = "21/07/09"
7__copyright__ = "(C) 2009 Science and Technology Facilities Council"
8__license__ = "BSD - see LICENSE file in top-level directory"
9__contact__ = "Philip.Kershaw@stfc.ac.uk"
10__revision__ = '$Id$'
11import logging
12logging.basicConfig(level=logging.DEBUG)
13   
14from datetime import datetime
15from uuid import uuid4
16
17import unittest
18
19from xml.etree.ElementTree import iselement
20
21from ndg.security.common.saml import Assertion, Attribute, AttributeValue, \
22    AttributeStatement, SAMLVersion, XSStringAttributeValue, \
23    XSGroupRoleAttributeValue
24from ndg.security.common.saml.xml import XMLConstants
25from ndg.security.common.saml.xml.etree import AssertionETreeObject, \
26    XSGroupRoleAttributeValueETreeObject
27
28
29class SAMLUtil(object):
30    """SAML utility class based on ANL examples for Earth System Grid:
31    http://www.ci.uchicago.edu/wiki/bin/view/ESGProject/ESGSAMLAttributes#ESG_Attribute_Service
32    """
33   
34    def __init__(self):
35        """Set-up ESG core attributes, Group/Role and miscellaneous
36        attributes lists
37        """
38        self.firstName = None
39        self.lastName = None
40        self.emailAddress = None
41       
42        self.__groupRoleList = []
43        self.__miscAttrList = []
44
45    def addGroupRole(self, group, role):
46        """Add an ESG Group/Role attribute
47        @type group: basestring
48        @param group: group name
49        @type role: basestring
50        @param role: role name
51        """
52        self.__groupRoleList.append((group, role))
53   
54    def addAttribute(self, name, value):
55        """Add a generic attribute
56        @type name: basestring
57        @param name: attribute name
58        @type value: basestring
59        @param value: attribute value
60        """
61        self.__miscAttrList.append((name, value))
62
63    def buildAssertion(self):
64        """Create a SAML Assertion containing ESG core attributes: First
65        Name, Last Name, e-mail Address; ESG Group/Role type attributes
66        and generic attributes
67        @rtype: ndg.security.common.saml.Assertion
68        @return: new SAML Assertion object
69        """
70       
71        assertion = Assertion()
72        assertion.version = SAMLVersion(SAMLVersion.VERSION_20)
73        assertion.id = str(uuid4())
74        assertion.issueInstant = datetime.utcnow()
75
76        attributeStatement = AttributeStatement()
77
78        if self.firstName is not None:   
79            # special case handling for 'FirstName' attribute
80            fnAttribute = Attribute()
81            fnAttribute.name = "urn:esg:first:name"
82            fnAttribute.nameFormat = "http://www.w3.org/2001/XMLSchema#string"
83            fnAttribute.friendlyName = "FirstName"
84
85            firstName = XSStringAttributeValue()
86            firstName.value = self.firstName
87            fnAttribute.attributeValues.append(firstName)
88
89            attributeStatement.attributes.append(fnAttribute)
90       
91
92        if self.lastName is not None:
93            # special case handling for 'LastName' attribute
94            lnAttribute = Attribute()
95            lnAttribute.name = "urn:esg:last:name"
96            lnAttribute.nameFormat = "http://www.w3.org/2001/XMLSchema#string"
97            lnAttribute.friendlyName = "LastName"
98
99            lastName = XSStringAttributeValue()
100            lastName.value = self.lastName
101            lnAttribute.attributeValues.append(lastName)
102
103            attributeStatement.attributes.append(lnAttribute)
104       
105
106        if self.emailAddress is not None:
107            # special case handling for 'LastName' attribute
108            emailAddressAttribute = Attribute()
109            emailAddressAttribute.name = "urn:esg:email:address"
110            emailAddressAttribute.nameFormat = XMLConstants.XSD_NS+"#"+\
111                                        XSStringAttributeValue.TYPE_LOCAL_NAME
112            emailAddressAttribute.friendlyName = "emailAddress"
113
114            emailAddress = XSStringAttributeValue()
115            emailAddress.value = self.emailAddress
116            emailAddressAttribute.attributeValues.append(emailAddress)
117
118            attributeStatement.attributes.append(emailAddressAttribute)
119       
120        if len(self.__groupRoleList) > 0:
121            # custom group/role attribute to be added to attr statement
122            groupRoleAttribute = Attribute()
123            groupRoleAttribute.name = "GroupRole"
124            groupRoleAttribute.nameFormat = \
125                                    XSGroupRoleAttributeValue.TYPE_LOCAL_NAME
126
127            for group, role in self.__groupRoleList:
128                groupRole = XSGroupRoleAttributeValue()
129                groupRole.group = group
130                groupRole.role = role
131
132                groupRoleAttribute.attributeValues.append(groupRole)
133           
134            attributeStatement.attributes.append(groupRoleAttribute)
135       
136        for name, value in self.__miscAttrList:
137            attribute = Attribute()
138            attribute.name = name
139            attribute.nameFormat = "http://www.w3.org/2001/XMLSchema#string"
140
141            stringAttributeValue = XSStringAttributeValue()
142            stringAttributeValue.value = value
143            attribute.attributeValues.append(stringAttributeValue)
144
145            # add all attributes to the attribute statement
146            attributeStatement.attributes.append(attribute)
147           
148        assertion.attributeStatements.append(attributeStatement)
149        return assertion
150
151class SAMLTestCase(unittest.TestCase):
152    """Test SAML implementation for use with CMIP5 federation"""
153   
154    def test01CreateAssertion(self):
155        samlUtil = SAMLUtil()
156       
157        # ESG core attributes
158        samlUtil.firstName = "Philip"
159        samlUtil.lastName = "Kershaw"
160        samlUtil.emailAddress = "p.j.k@somewhere"
161       
162        # BADC specific attributes
163        badcRoleList = (
164            'urn:badc:security:authz:1.0:attr:admin', 
165            'urn:badc:security:authz:1.0:attr:rapid', 
166            'urn:badc:security:authz:1.0:attr:coapec', 
167            'urn:badc:security:authz:1.0:attr:midas', 
168            'urn:badc:security:authz:1.0:attr:quest', 
169            'urn:badc:security:authz:1.0:attr:staff'
170        )
171        for role in badcRoleList:
172            samlUtil.addAttribute("urn:badc:security:authz:1.0:attr", role)
173       
174        # ESG Group/Role type list
175        esgGroupRoleList = (
176            ("ESG-NCAR", "admin"),
177            ("ESG-PCMDI", "testUser"),
178        )
179        for group, role in esgGroupRoleList:
180            samlUtil.addGroupRole(group, role)
181       
182        # Make an assertion object
183        assertion = samlUtil.buildAssertion()
184       
185        # Create XML rendering class using the ElementTree implementation
186        assertionETreeObject = AssertionETreeObject()
187       
188        # Add mapping for ESG Group/Role Attribute Value to enable ElementTree
189        # Attribute Value factory to render the XML output
190        attributeValueETreeObjectClassMap = {
191            XSGroupRoleAttributeValue: XSGroupRoleAttributeValueETreeObject           
192        }
193       
194        # Create ElementTree Assertion Element
195        assertionElem = assertionETreeObject.create(assertion,
196                            customClassMap=attributeValueETreeObjectClassMap)
197       
198        self.assert_(iselement(assertionElem))
199       
200        # Serialise to output
201        xmlOutput = assertionETreeObject.prettyPrint()
202        self.assert_(len(xmlOutput))
203        print(xmlOutput)
204   
205
206if __name__ == "__main__":
207    unittest.main()       
Note: See TracBrowser for help on using the repository browser.