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

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@5554
Revision 5554, 9.9 KB checked in by pjkersha, 11 years ago (diff)
  • Started adding ElementTree based parsers for SAML classes in ndg.security.common.saml.xml.etree.
  • ndg.security.common.utils.prettyPrint needs a bug fix for namespace declarations
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
16from cStringIO import StringIO
17
18import unittest
19
20from xml.etree.ElementTree import iselement
21from xml.etree import ElementTree
22
23from ndg.security.common.saml import Assertion, Attribute, AttributeValue, \
24    AttributeStatement, SAMLVersion, XSStringAttributeValue, \
25    XSGroupRoleAttributeValue, AttributeQuery, Issuer, Subject, NameID
26from ndg.security.common.saml.xml import XMLConstants
27from ndg.security.common.saml.xml.etree import AssertionETreeObject, \
28    XSGroupRoleAttributeValueETreeObject, AttributeQueryETreeObject
29
30
31class SAMLUtil(object):
32    """SAML utility class based on ANL examples for Earth System Grid:
33    http://www.ci.uchicago.edu/wiki/bin/view/ESGProject/ESGSAMLAttributes#ESG_Attribute_Service
34    """
35   
36    def __init__(self):
37        """Set-up ESG core attributes, Group/Role and miscellaneous
38        attributes lists
39        """
40        self.firstName = None
41        self.lastName = None
42        self.emailAddress = None
43       
44        self.__groupRoleList = []
45        self.__miscAttrList = []
46
47    def addGroupRole(self, group, role):
48        """Add an ESG Group/Role attribute
49        @type group: basestring
50        @param group: group name
51        @type role: basestring
52        @param role: role name
53        """
54        self.__groupRoleList.append((group, role))
55   
56    def addAttribute(self, name, value):
57        """Add a generic attribute
58        @type name: basestring
59        @param name: attribute name
60        @type value: basestring
61        @param value: attribute value
62        """
63        self.__miscAttrList.append((name, value))
64
65    def buildAssertion(self):
66        """Create a SAML Assertion containing ESG core attributes: First
67        Name, Last Name, e-mail Address; ESG Group/Role type attributes
68        and generic attributes
69        @rtype: ndg.security.common.saml.Assertion
70        @return: new SAML Assertion object
71        """
72       
73        assertion = Assertion()
74        assertion.version = SAMLVersion(SAMLVersion.VERSION_20)
75        assertion.id = str(uuid4())
76        assertion.issueInstant = datetime.utcnow()
77
78        attributeStatement = AttributeStatement()
79       
80        for attribute in self.createAttributes():
81            attributeStatement.attributes.append(attribute)
82           
83        return assertion
84
85    def buildAttributeQuery(self, issuer, subjectNameID):
86       
87        attributeQuery = AttributeQuery()
88        attributeQuery.version = SAMLVersion(SAMLVersion.VERSION_20)
89        attributeQuery.id = str(uuid4())
90        attributeQuery.issueInstant = datetime.utcnow()
91       
92        attributeQuery.issuer = Issuer()
93        attributeQuery.issuer.format = Issuer.X509_SUBJECT
94        attributeQuery.issuer.value = issuer
95                       
96        attributeQuery.subject = Subject() 
97        attributeQuery.subject.nameID = NameID()
98        attributeQuery.subject.nameID.format = "urn:esg:openid"
99        attributeQuery.subject.nameID.value = subjectNameID
100                                   
101        attributeQuery.attributes = self.createAttributes()
102       
103        return attributeQuery
104   
105    def createAttributes(self):
106       
107        attributes = []
108        if self.firstName is not None:   
109            # special case handling for 'FirstName' attribute
110            fnAttribute = Attribute()
111            fnAttribute.name = "urn:esg:first:name"
112            fnAttribute.nameFormat = "http://www.w3.org/2001/XMLSchema#string"
113            fnAttribute.friendlyName = "FirstName"
114
115            firstName = XSStringAttributeValue()
116            firstName.value = self.firstName
117            fnAttribute.attributeValues.append(firstName)
118
119            attributes.append(fnAttribute)
120       
121
122        if self.lastName is not None:
123            # special case handling for 'LastName' attribute
124            lnAttribute = Attribute()
125            lnAttribute.name = "urn:esg:last:name"
126            lnAttribute.nameFormat = "http://www.w3.org/2001/XMLSchema#string"
127            lnAttribute.friendlyName = "LastName"
128
129            lastName = XSStringAttributeValue()
130            lastName.value = self.lastName
131            lnAttribute.attributeValues.append(lastName)
132
133            attributes.append(lnAttribute)
134       
135
136        if self.emailAddress is not None:
137            # special case handling for 'LastName' attribute
138            emailAddressAttribute = Attribute()
139            emailAddressAttribute.name = "urn:esg:email:address"
140            emailAddressAttribute.nameFormat = XMLConstants.XSD_NS+"#"+\
141                                        XSStringAttributeValue.TYPE_LOCAL_NAME
142            emailAddressAttribute.friendlyName = "emailAddress"
143
144            emailAddress = XSStringAttributeValue()
145            emailAddress.value = self.emailAddress
146            emailAddressAttribute.attributeValues.append(emailAddress)
147
148            attributes.append(emailAddressAttribute)
149       
150        if len(self.__groupRoleList) > 0:
151            # custom group/role attribute to be added to attr statement
152            groupRoleAttribute = Attribute()
153            groupRoleAttribute.name = "GroupRole"
154            groupRoleAttribute.nameFormat = \
155                                    XSGroupRoleAttributeValue.TYPE_LOCAL_NAME
156
157            for group, role in self.__groupRoleList:
158                groupRole = XSGroupRoleAttributeValue()
159                groupRole.group = group
160                groupRole.role = role
161
162                groupRoleAttribute.attributeValues.append(groupRole)
163           
164            attributes.append(groupRoleAttribute)
165       
166        for name, value in self.__miscAttrList:
167            attribute = Attribute()
168            attribute.name = name
169            attribute.nameFormat = "http://www.w3.org/2001/XMLSchema#string"
170
171            stringAttributeValue = XSStringAttributeValue()
172            stringAttributeValue.value = value
173            attribute.attributeValues.append(stringAttributeValue)
174
175            attributes.append(attribute)
176           
177        return attributes
178
179
180class SAMLTestCase(unittest.TestCase):
181    """Test SAML implementation for use with CMIP5 federation"""
182   
183    def test01CreateAssertion(self):
184        samlUtil = SAMLUtil()
185       
186        # ESG core attributes
187        samlUtil.firstName = "Philip"
188        samlUtil.lastName = "Kershaw"
189        samlUtil.emailAddress = "p.j.k@somewhere"
190       
191        # BADC specific attributes
192        badcRoleList = (
193            'urn:badc:security:authz:1.0:attr:admin', 
194            'urn:badc:security:authz:1.0:attr:rapid', 
195            'urn:badc:security:authz:1.0:attr:coapec', 
196            'urn:badc:security:authz:1.0:attr:midas', 
197            'urn:badc:security:authz:1.0:attr:quest', 
198            'urn:badc:security:authz:1.0:attr:staff'
199        )
200        for role in badcRoleList:
201            samlUtil.addAttribute("urn:badc:security:authz:1.0:attr", role)
202       
203        # ESG Group/Role type list
204        esgGroupRoleList = (
205            ("ESG-NCAR", "admin"),
206            ("ESG-PCMDI", "testUser"),
207        )
208        for group, role in esgGroupRoleList:
209            samlUtil.addGroupRole(group, role)
210       
211        # Make an assertion object
212        assertion = samlUtil.buildAssertion()
213       
214        # Create XML rendering class using the ElementTree implementation
215        assertionETreeObject = AssertionETreeObject()
216       
217        # Add mapping for ESG Group/Role Attribute Value to enable ElementTree
218        # Attribute Value factory to render the XML output
219        attributeValueETreeObjectClassMap = {
220            XSGroupRoleAttributeValue: XSGroupRoleAttributeValueETreeObject           
221        }
222       
223        # Create ElementTree Assertion Element
224        assertionElem = assertionETreeObject.create(assertion,
225                            customClassMap=attributeValueETreeObjectClassMap)
226       
227        self.assert_(iselement(assertionElem))
228       
229        # Serialise to output
230        xmlOutput = assertionETreeObject.prettyPrint()
231        self.assert_(len(xmlOutput))
232        print(xmlOutput)
233
234    def test02CreateAttributeQuery(self):
235        samlUtil = SAMLUtil()
236        samlUtil.firstName = ''
237        samlUtil.lastName = ''
238        samlUtil.emailAddress = ''
239        attributeQuery = samlUtil.buildAttributeQuery(
240                        "/O=NDG/OU=BADC/CN=attributeauthority.badc.rl.ac.uk",
241                        "https://openid.localhost/philip.kershaw")
242       
243        elem = AttributeQueryETreeObject.create(attributeQuery)
244        xmlOutput = AttributeQueryETreeObject.prettyPrint(elem)
245        print(xmlOutput)
246
247    def test03ParseAttributeQuery(self):
248        samlUtil = SAMLUtil()
249        samlUtil.firstName = ''
250        samlUtil.lastName = ''
251        samlUtil.emailAddress = ''
252        attributeQuery = samlUtil.buildAttributeQuery(
253                        "/O=NDG/OU=BADC/CN=attributeauthority.badc.rl.ac.uk",
254                        "https://openid.localhost/philip.kershaw")
255       
256        elem = AttributeQueryETreeObject.create(attributeQuery)
257#        xmlOutput = AttributeQueryETreeObject.prettyPrint(elem)
258        xmlOutput = AttributeQueryETreeObject.serialize(elem)
259       
260        attributeQueryStream = StringIO()
261        attributeQueryStream.write(xmlOutput)
262        attributeQueryStream.seek(0)
263        print ElementTree.tostring(elem)
264        tree = ElementTree.parse(attributeQueryStream)
265        elem2 = tree.getroot()
266       
267        attributeQuery2 = AttributeQueryETreeObject.parse(elem2)
268        xmlOutput2 = AttributeQueryETreeObject.serialize(elem2)
269        self.assert_(xmlOutput == xmlOutput2)
270       
271       
272if __name__ == "__main__":
273    unittest.main()       
Note: See TracBrowser for help on using the repository browser.