Added oauth RBAC via regex on any fields in the openid profile output

CybercentreCanada · Apr 6, 2020 · 35afe22 · 35afe22
1 parent 54a31b6
commit 35afe22
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 39 deletions.
diff --git a/assemblyline_ui/helper/oauth.py b/assemblyline_ui/helper/oauth.py
@@ -1,11 +1,18 @@
 import base64
 import hashlib
+import re
 import requests
 
+from assemblyline.common import forge
 from assemblyline_ui.config import config
 
+cl_engine = forge.get_classification()
+
+
+def parse_profile(profile, auto_prop_list=None):
+    if not auto_prop_list:
+        auto_prop_list = []
 
-def parse_profile(profile):
     email_adr = profile.get('email', profile.get('upn', None))
 
     if config.auth.oauth.gravatar_enabled:
@@ -14,7 +21,26 @@ def parse_profile(profile):
     else:
         alternate = None
 
+    access = True
+    roles = ['user']
+    classification = cl_engine.UNRESTRICTED
+    for auto_prop in auto_prop_list:
+        if auto_prop.type == "access":
+            if auto_prop.value == "True":
+                access = re.match(auto_prop.pattern, profile.get(auto_prop.field, "")) is not None
+            else:
+                access = re.match(auto_prop.pattern, profile.get(auto_prop.field, "")) is None
+        elif auto_prop.type == "role":
+            if re.match(auto_prop.pattern, profile.get(auto_prop.field, "")):
+                roles.append(auto_prop.value)
+        elif auto_prop.type == "classification":
+            if re.match(auto_prop.pattern, profile.get(auto_prop.field, "")):
+                classification = cl_engine.max_classification(classification, auto_prop.value)
+
     return dict(
+        access=access,
+        type=roles,
+        classification=classification,
         uname=profile.get('uname', email_adr),
         name=profile.get('name', None),
         email=email_adr,

diff --git a/assemblyline_ui/views.py b/assemblyline_ui/views.py
@@ -219,46 +219,51 @@ def login():
                     # Get user data
                     resp = provider.get(config.auth.oauth.providers[oauth_provider].user_get)
                     if resp.ok:
-                        data = parse_profile(resp.json())
-                        oauth_avatar = data.pop('avatar', None)
-
-                        # Find if user already exists
-                        users = STORAGE.user.search(f"email:{data['email']}", fl="uname", as_obj=False)['items']
-                        if users:
-                            cur_user = STORAGE.user.get(users[0]['uname'], as_obj=False) or {}
-                            # Do not update username and password from the current user
-                            data['uname'] = cur_user.get('uname', data['uname'])
-                            data['password'] = cur_user.get('password', data['password'])
-                        else:
-                            cur_user = {}
-
-                        username = data['uname']
-
-                        # Make sure the user exists in AL and is in sync
-                        if (not cur_user and config.auth.oauth.providers[oauth_provider].auto_create) or \
-                                (cur_user and config.auth.oauth.providers[oauth_provider].auto_sync):
-
-                            # Update the current user
-                            cur_user.update(data)
-
-                            # Save avatar
-                            if oauth_avatar:
-                                avatar = fetch_avatar(oauth_avatar)
-                                STORAGE.user_avatar.save(username, avatar)
-
-                            # Save updated user
-                            STORAGE.user.save(username, cur_user)
-
-                        if cur_user:
-                            if avatar is None:
-                                avatar = STORAGE.user_avatar.get(username) or "/static/images/user_default.png"
-                            oauth_token = hashlib.sha256(str(token).encode("utf-8", errors='replace')).hexdigest()
-                            get_token_store(username).add(oauth_token)
+                        data = parse_profile(resp.json(), config.auth.oauth.providers[oauth_provider].auto_properties)
+                        has_access = data.pop('access', False)
+                        if has_access:
+                            oauth_avatar = data.pop('avatar', None)
+
+                            # Find if user already exists
+                            users = STORAGE.user.search(f"email:{data['email']}", fl="uname", as_obj=False)['items']
+                            if users:
+                                cur_user = STORAGE.user.get(users[0]['uname'], as_obj=False) or {}
+                                # Do not update username and password from the current user
+                                data['uname'] = cur_user.get('uname', data['uname'])
+                                data['password'] = cur_user.get('password', data['password'])
+                            else:
+                                cur_user = {}
+
+                            username = data['uname']
+
+                            # Make sure the user exists in AL and is in sync
+                            if (not cur_user and config.auth.oauth.providers[oauth_provider].auto_create) or \
+                                    (cur_user and config.auth.oauth.providers[oauth_provider].auto_sync):
+
+                                # Update the current user
+                                cur_user.update(data)
+
+                                # Save avatar
+                                if oauth_avatar:
+                                    avatar = fetch_avatar(oauth_avatar)
+                                    STORAGE.user_avatar.save(username, avatar)
+
+                                # Save updated user
+                                STORAGE.user.save(username, cur_user)
+
+                            if cur_user:
+                                if avatar is None:
+                                    avatar = STORAGE.user_avatar.get(username) or "/static/images/user_default.png"
+                                oauth_token = hashlib.sha256(str(token).encode("utf-8", errors='replace')).hexdigest()
+                                get_token_store(username).add(oauth_token)
+                            else:
+                                oauth_validation = False
+                                avatar = None
+                                username = ''
+                                oauth_error = "User auto-creation is disabled"
                         else:
                             oauth_validation = False
-                            avatar = None
-                            username = ''
-                            oauth_error = "User auto-creation is disabled"
+                            oauth_error = "This user is not allowed access to the system"
 
                 except Exception as _:
                     oauth_validation = False