From af912665e6451683d2b3638a063ef63e46a4420a Mon Sep 17 00:00:00 2001
From: Valentin Deniaud <vdeniaud@entrouvert.com>
Date: Wed, 24 Apr 2019 11:41:09 +0200
Subject: [PATCH 2/3] views: save is_staff in session

---
 mellon/adapters.py | 24 ++++++------------------
 mellon/utils.py    | 18 ++++++++++++++++++
 mellon/views.py    |  2 ++
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/mellon/adapters.py b/mellon/adapters.py
index 2aacdf1..96dad0c 100644
--- a/mellon/adapters.py
+++ b/mellon/adapters.py
@@ -183,24 +183,12 @@ class DefaultAdapter(object):
 
     def provision_superuser(self, user, idp, saml_attributes):
         superuser_mapping = utils.get_setting(idp, 'SUPERUSER_MAPPING')
-        if not superuser_mapping:
-            return
-        for key, values in superuser_mapping.items():
-            if key in saml_attributes:
-                if not isinstance(values, (tuple, list)):
-                    values = [values]
-                values = set(values)
-                attribute_values = saml_attributes[key]
-                if not isinstance(attribute_values, (tuple, list)):
-                    attribute_values = [attribute_values]
-                attribute_values = set(attribute_values)
-                if attribute_values & values:
-                    if not (user.is_staff and user.is_superuser):
-                        user.is_staff = True
-                        user.is_superuser = True
-                        user.save()
-                        self.logger.info('flag is_staff and is_superuser added to user %s', user)
-                    break
+        if utils.has_superuser_flag(idp, saml_attributes):
+            if not (user.is_staff and user.is_superuser):
+                user.is_staff = True
+                user.is_superuser = True
+                user.save()
+                self.logger.info('flag is_staff and is_superuser added to user %s', user)
         else:
             self.remove_superuser(user)
 
diff --git a/mellon/utils.py b/mellon/utils.py
index c7ffe3d..754f071 100644
--- a/mellon/utils.py
+++ b/mellon/utils.py
@@ -276,6 +276,24 @@ def get_local_path(request, url):
     return path
 
 
+def has_superuser_flag(idp, saml_attributes):
+    superuser_mapping = get_setting(idp, 'SUPERUSER_MAPPING')
+    if not superuser_mapping:
+        return False
+    for key, values in superuser_mapping.items():
+        if key in saml_attributes:
+            if not isinstance(values, (tuple, list)):
+                values = [values]
+            values = set(values)
+            attribute_values = saml_attributes[key]
+            if not isinstance(attribute_values, (tuple, list)):
+                attribute_values = [attribute_values]
+            attribute_values = set(attribute_values)
+            if attribute_values & values:
+                return True
+    return False
+
+
 def user_has_role(request, role_id):
     try:
         group = request.user.groups.get(id=role_id)
diff --git a/mellon/views.py b/mellon/views.py
index 7b0b1fe..f5f5edd 100644
--- a/mellon/views.py
+++ b/mellon/views.py
@@ -218,6 +218,8 @@ class LoginView(ProfileMixin, LogMixin, View):
             if user.is_active:
                 utils.login(request, user)
                 request.session['role_uuids'] = dict.fromkeys(attributes['role-slug'])
+                idp = self.get_idp(request)
+                request.session['is_staff'] = utils.has_superuser_flag(idp, attributes)
                 self.log.info('user %s (NameID is %r) logged in using SAML', user,
                               attributes['name_id_content'])
                 request.session['mellon_session'] = utils.flatten_datetime(attributes)
-- 
2.20.1