From 9a5aae50ab5d85b425c803f734f8b9077927636c Mon Sep 17 00:00:00 2001
From: zhengxie <xiez1989@gmail.com>
Date: Tue, 1 Aug 2017 15:41:03 +0800
Subject: [PATCH] [shibboleth] Add joker * in affiliation role map

---
 .../thirdpart/shibboleth/test_middleware.py   | 32 ++++++++++++++++++
 thirdpart/shibboleth/middleware.py            | 33 +++++++++++++++----
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/tests/seahub/thirdpart/shibboleth/test_middleware.py b/tests/seahub/thirdpart/shibboleth/test_middleware.py
index 0907c3e3ed..8f72571679 100644
--- a/tests/seahub/thirdpart/shibboleth/test_middleware.py
+++ b/tests/seahub/thirdpart/shibboleth/test_middleware.py
@@ -146,3 +146,35 @@ class ShibbolethRemoteUserMiddlewareTest(BaseTestCase):
 
         assert len(Profile.objects.all()) == 1
         assert Profile.objects.all()[0].nickname == ''
+
+    @override_settings(SHIBBOLETH_AFFILIATION_ROLE_MAP={
+        'employee@school.edu': 'staff',
+        'member@school.edu': 'staff',
+        'student@school.edu': 'student',
+        'patterns': (
+            ('a@*.edu', 'aaa'),
+            ('*@*.edu', 'student'),
+            ('*', 'guest'),
+        )
+    })
+    @patch('shibboleth.middleware.SHIB_ATTRIBUTE_MAP', {
+        "Shibboleth-eppn": (True, "username"),
+        "givenname": (False, "givenname"),
+        "surname": (False, "surname"),
+        "emailaddress": (False, "contact_email"),
+        "organization": (False, "institution"),
+        "Shibboleth-affiliation": (False, "affiliation"),
+        "Shibboleth-displayName": (False, "display_name"),
+    })
+    def test_get_role_by_affiliation(self):
+        obj = ShibbolethRemoteUserMiddleware()
+
+        assert obj._get_role_by_affiliation('employee@school.edu') == 'staff'
+        assert obj._get_role_by_affiliation('member@school.edu') == 'staff'
+        assert obj._get_role_by_affiliation('student@school.edu') == 'student'
+
+        # test jokers
+        assert obj._get_role_by_affiliation('student1@school.edu') == 'student'
+        assert obj._get_role_by_affiliation('a@x.edu') == 'aaa'
+        assert obj._get_role_by_affiliation('a@x.com') == 'guest'
+
diff --git a/thirdpart/shibboleth/middleware.py b/thirdpart/shibboleth/middleware.py
index 78deae205d..3bf719c098 100755
--- a/thirdpart/shibboleth/middleware.py
+++ b/thirdpart/shibboleth/middleware.py
@@ -1,3 +1,5 @@
+from collections import OrderedDict
+from fnmatch import fnmatch
 import logging
 
 from django.conf import settings
@@ -156,18 +158,37 @@ class ShibbolethRemoteUserMiddleware(RemoteUserMiddleware):
 
         p.save()
 
-    def update_user_role(self, user, shib_meta):
-        affiliation = shib_meta.get('affiliation', '')
-        if not affiliation:
-            return
-
+    def _get_role_by_affiliation(self, affiliation):
         try:
             role_map = settings.SHIBBOLETH_AFFILIATION_ROLE_MAP
         except AttributeError:
             return
 
+        role = role_map.get(affiliation)
+        if role:
+            return role
+
+        if role_map.get('patterns') is not None:
+            joker_map = role_map.get('patterns')
+            try:
+                od = OrderedDict(joker_map)
+            except Exception as e:
+                logger.error(e)
+                return
+
+            for k in od:
+                if fnmatch(affiliation, k):
+                    return od[k]
+
+        return None
+
+    def update_user_role(self, user, shib_meta):
+        affiliation = shib_meta.get('affiliation', '')
+        if not affiliation:
+            return
+
         for e in affiliation.split(';'):
-            role = role_map.get(e)
+            role = self._get_role_by_affiliation(e)
             if role:
                 User.objects.update_role(user.email, role)
                 return role