From 7d490de37cd9d7250932e1296b7d3ec73e3656da Mon Sep 17 00:00:00 2001
From: Carson Gee <x@carsongee.com>
Date: Thu, 4 Dec 2014 16:03:07 -0500
Subject: [PATCH] Add CAS 3.0 attribute support, and style cleanups

---
 mitx_cas_mapper/__init__.py | 71 ++++++++++++++++++++++++++++---------
 1 file changed, 55 insertions(+), 16 deletions(-)

diff --git a/mitx_cas_mapper/__init__.py b/mitx_cas_mapper/__init__.py
index 6ecbf1b..ba5ddd5 100644
--- a/mitx_cas_mapper/__init__.py
+++ b/mitx_cas_mapper/__init__.py
@@ -1,31 +1,70 @@
 """
 This is used as an attribute mapper
-callable for django-cas per the README at 
+callable for django-cas per the README at
 https://github.com/mitocw/django-cas.
-""" 
+"""
+
+from django.conf import settings
 
-VERSION = "0.0.1"
+VERSION = "0.1.0"
 
 CAS_URI = 'http://www.yale.edu/tp/cas'
 NSMAP = {'cas': CAS_URI}
 CAS = '{%s}' % CAS_URI
 
+
 def populate_user(user, authentication_response):
-    if authentication_response.find(CAS + 'authenticationSuccess/'  + CAS + 'attributes'  , namespaces=NSMAP) is not None:
-        attr = authentication_response.find(CAS + 'authenticationSuccess/'  + CAS + 'attributes'  , namespaces=NSMAP)
+    """
+    This is passed a django user object to be modified
+    and the `authentication_response` from the CAS server.
+
+    It allows you to convert attributes from CAS into local
+    Django user attributes
+    """
+    if settings.CAS_VERSION == '2':
+        if authentication_response.find(
+                CAS + 'authenticationSuccess/' + CAS + 'attributes',
+                namespaces=NSMAP
+        ) is not None:
+            attr = authentication_response.find(
+                CAS + 'authenticationSuccess/' + CAS + 'attributes',
+                namespaces=NSMAP
+            )
+
+            if attr.find(CAS + 'is_staff', NSMAP) is not None:
+                user.is_staff = attr.find(
+                    CAS + 'is_staff',
+                    NSMAP
+                ).text.upper() == 'TRUE'
+
+            if attr.find(CAS + 'is_superuser', NSMAP) is not None:
+                user.is_superuser = attr.find(
+                    CAS + 'is_superuser',
+                    NSMAP
+                    ).text.upper() == 'TRUE'
+
+            if attr.find(CAS + 'givenName', NSMAP) is not None:
+                user.first_name = attr.find(CAS + 'givenName', NSMAP).text
+
+            if attr.find(CAS + 'sn', NSMAP) is not None:
+                user.last_name = attr.find(CAS + 'sn', NSMAP).text
+
+            if attr.find(CAS + 'email', NSMAP) is not None:
+                user.email = attr.find(CAS + 'email', NSMAP).text
 
-        if attr.find(CAS + 'is_staff', NSMAP) is not None:
-            user.is_staff = attr.find(CAS + 'is_staff', NSMAP).text.upper() == 'TRUE'
+    if settings.CAS_VERSION == '3':
+        if authentication_response is not None:
+            if 'is_superuser' in authentication_response:
+                user.is_superuser = authentication_response['is_superuser']
 
-        if attr.find(CAS + 'is_superuser', NSMAP) is not None:
-            user.is_superuser = attr.find(CAS + 'is_superuser', NSMAP).text.upper() == 'TRUE'
+            if 'is_staff' in authentication_response:
+                user.is_staff = authentication_response['is_staff']
 
-        if attr.find(CAS + 'givenName', NSMAP) is not None:
-            user.first_name = attr.find(CAS + 'givenName', NSMAP).text
+            if 'givenName' in authentication_response:
+                user.first_name = authentication_response['givenName']
 
-        if attr.find(CAS + 'sn', NSMAP) is not None:
-            user.last_name = attr.find(CAS + 'sn', NSMAP).text
+            if 'sn' in authentication_response:
+                user.last_name = authentication_response['sn']
 
-        if attr.find(CAS + 'email', NSMAP) is not None:
-            user.email = attr.find(CAS + 'email', NSMAP).text
-    pass
+            if 'email' in authentication_response:
+                user.email = authentication_response['email']