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

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

diff --git a/mellon/adapters.py b/mellon/adapters.py
index 2aacdf1..21542b7 100644
--- a/mellon/adapters.py
+++ b/mellon/adapters.py
@@ -182,25 +182,12 @@ class DefaultAdapter(object):
             user.save()
 
     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 6462f81..ee8b8a5 100644
--- a/mellon/utils.py
+++ b/mellon/utils.py
@@ -271,3 +271,21 @@ def get_local_path(request, url):
     if request.META.get('SCRIPT_NAME'):
         path = path[len(request.META['SCRIPT_NAME']):]
     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
diff --git a/mellon/views.py b/mellon/views.py
index 4a3da73..d2d3f87 100644
--- a/mellon/views.py
+++ b/mellon/views.py
@@ -218,6 +218,8 @@ class LoginView(ProfileMixin, LogMixin, View):
         if user is not None:
             if user.is_active:
                 utils.login(request, user)
+                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