Login via Managed Identity token

assemblyline_ui/api/v4/user.py

@@ -420,6 +420,8 @@ def get_user_account(username, **kwargs):
         user['avatar'] = STORAGE.user_avatar.get(username)
 
     user['roles'] = load_roles(user['type'], user.get('roles', None))
+    if ROLES.administration not in kwargs['user']['roles']:
+        user.pop('identity_id')
 
     return make_api_response(user)
 

assemblyline_ui/app.py

@@ -45,6 +45,7 @@
 from assemblyline_ui import config
 
 AL_UNSECURED_UI = os.environ.get('AL_UNSECURED_UI', 'false').lower() == 'true'
+THREADED = os.environ.get('THREADED', 'true').lower() == 'true'
 AL_SESSION_COOKIE_SAMESITE = os.environ.get("AL_SESSION_COOKIE_SAMESITE", None)
 AL_HSTS_MAX_AGE = os.environ.get('AL_HSTS_MAX_AGE', None)
 CERT_BUNDLE = (
@@ -185,7 +186,7 @@ def main():
         wlog.addHandler(h)
 
     app.jinja_env.cache = {}
-    app.run(host="0.0.0.0", debug=False, ssl_context=ssl_context)
+    app.run(host="0.0.0.0", debug=False, ssl_context=ssl_context, threaded=THREADED)
 
 
 if __name__ == '__main__':

assemblyline_ui/helper/oauth.py

@@ -3,6 +3,8 @@
 import re
 import requests
 
+from authlib.integrations.flask_client import FlaskRemoteApp
+from assemblyline.odm.models.config import OAuthProvider
 from assemblyline.odm.models.user import load_roles, USER_TYPE_DEP
 from assemblyline.common.random_user import random_user
 from assemblyline_ui.config import config, CLASSIFICATION as cl_engine
@@ -16,8 +18,7 @@ def reorder_name(name):
 
     return " ".join(name.split(", ", 1)[::-1])
 
-
-def parse_profile(profile, provider):
+def get_profile_identifiers(profile: dict, provider: OAuthProvider):
     # Find email address and normalize it for further processing
     email_adr = None
     for email_key in provider.email_fields:
@@ -33,6 +34,18 @@ def parse_profile(profile, provider):
         if "@" not in email_adr:
             email_adr = None
 
+    # Find identity ID
+    identity_id = profile.get(provider.identity_id_field, None)
+
+    return dict(
+        email=email_adr,
+        identity_id=identity_id
+    )
+
+def parse_profile(profile: dict, provider: OAuthProvider):
+    profile_identifiers = get_profile_identifiers(profile, provider)
+    email_adr = profile_identifiers['email']
+
     # Find the name of the user
     name = reorder_name(profile.get('name', profile.get('displayName', None)))
 
@@ -191,13 +204,14 @@ def parse_profile(profile, provider):
         uname=uname,
         name=name,
         email=email_adr,
+        identity_id=profile_identifiers['identity_id'],
         password="__NO_PASSWORD__",
         avatar=profile.get('picture', alternate),
         **quotas
     )
 
 
-def fetch_avatar(url, provider, provider_config):
+def fetch_avatar(url: str, provider: FlaskRemoteApp, provider_config:OAuthProvider):
     if url.startswith(provider_config.api_base_url):
         resp = provider.get(url[len(provider_config.api_base_url):])
         if resp.ok and resp.headers.get("content-type") is not None:

assemblyline_ui/security/oauth_auth.py

@@ -5,8 +5,8 @@
 from copy import copy
 
 from assemblyline.odm.models.user import load_roles_form_acls
-from assemblyline_ui.config import config, get_token_store, STORAGE, CACHE
-from assemblyline_ui.helper.oauth import parse_profile
+from assemblyline_ui.config import config, get_token_store, STORAGE, CACHE, LOGGER
+from assemblyline_ui.helper.oauth import get_profile_identifiers
 from assemblyline_ui.http_exceptions import AuthenticationException
 
 
@@ -82,19 +82,24 @@ def validate_oauth_token(oauth_token, oauth_provider, return_user=False):
             except jwt.PyJWTError as e:
                 raise AuthenticationException(f"Invalid token - {str(e)}")
 
-            # Get user's email from profile
-            email = parse_profile(jwt_data, oauth_provider_config).get('email', None)
-            if email is not None:
-                # Get user from it's email
-                users = STORAGE.user.search(f"email:{email}", fl="*", as_obj=False, rows=1)['items']
-                if users:
-                    # Limit user logging in from external token to only user READ/WRITE APIs
-                    roles = load_roles_form_acls(["R", "W"], [])
+            profile_identifiers = get_profile_identifiers(jwt_data, oauth_provider_config)
 
-                    if return_user:
-                        return users[0], roles
-                    return users[0]['uname'], roles
+            # Lookup user via its profile identifiers (email or identity_id)
+            for k, v in profile_identifiers.items():
+                if v is not None:
+                    # Get user from it's email
+                    users = STORAGE.user.search(f"{k}:{v}", fl="*", as_obj=False, rows=1)['items']
+                    if users:
+                        # Limit user logging in from external token to only user READ/WRITE APIs
+                        roles = load_roles_form_acls(["R", "W"], [])
 
-        raise AuthenticationException("Invalid token")
+                        if return_user:
+                            return users[0], roles
+                        return users[0]['uname'], roles
+            msg = ", ".join([f"{k}={v}" for k, v in profile_identifiers.items() if v is not None])
+            raise AuthenticationException(f"User not found - No matching user for the following identifiers ({msg})")
+
+
+        raise AuthenticationException("Invalid token - No matching signing key found")
 
     return None, None