From 83a0485518664c2a8f3965d9b103298a9f4e4b79 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Wed, 17 Jul 2024 11:22:52 +0300
Subject: [PATCH] Fix user credentials reset on apiserver restart

---
 apiserver/mongo/initialize/__init__.py |  2 +-
 apiserver/mongo/initialize/user.py     | 61 ++++++++++++++++++--------
 2 files changed, 43 insertions(+), 20 deletions(-)

diff --git a/apiserver/mongo/initialize/__init__.py b/apiserver/mongo/initialize/__init__.py
index 94fc24c..02f6598 100644
--- a/apiserver/mongo/initialize/__init__.py
+++ b/apiserver/mongo/initialize/__init__.py
@@ -73,7 +73,7 @@ def init_mongo_data():
             }
             internal_user_emails.add(email.lower())
             revoke = fixed_mode and credentials.get("revoke_in_fixed_mode", False)
-            user_id = _ensure_auth_user(user_data, company_id, log=log, revoke=revoke)
+            user_id = _ensure_auth_user(user_data, company_id, log=log, revoke=revoke, internal_user=True)
             if credentials.role == Role.user:
                 _ensure_backend_user(user_id, company_id, credentials.display_name)
 
diff --git a/apiserver/mongo/initialize/user.py b/apiserver/mongo/initialize/user.py
index 90ca85f..12e6d0b 100644
--- a/apiserver/mongo/initialize/user.py
+++ b/apiserver/mongo/initialize/user.py
@@ -10,7 +10,12 @@ from apiserver.service_repo.auth.fixed_user import FixedUser
 
 
 def _ensure_user_credentials(
-    user: AuthUser, key: str, secret: str, log: Logger, revoke: bool = False
+    user: AuthUser,
+    key: str,
+    secret: str,
+    log: Logger,
+    revoke: bool = False,
+    internal_user: bool = False,
 ) -> None:
     if revoke:
         log.info(f"Revoking credentials for existing user {user.id} ({user.name})")
@@ -19,19 +24,34 @@ def _ensure_user_credentials(
         return
 
     if not (key and secret):
-        log.info(f"Resetting credentials for existing user {user.id} ({user.name})")
-        user.credentials = []
-        user.save()
+        if internal_user:
+            log.info(f"Resetting credentials for existing user {user.id} ({user.name})")
+            user.credentials = []
+            user.save()
         return
 
     new_credentials = Credentials(key=key, secret=secret)
-    log.info(f"Setting credentials for existing user {user.id} ({user.name})")
-    user.credentials = [new_credentials]
-    user.save()
-    return
+    if internal_user:
+        log.info(f"Setting credentials for existing user {user.id} ({user.name})")
+        user.credentials = [new_credentials]
+        user.save()
+        return
+
+    if user.credentials is None:
+        user.credentials = []
+    if not any((cred.key, cred.secret) == (key, secret) for cred in user.credentials):
+        log.info(f"Adding credentials for existing user {user.id} ({user.name})")
+        user.credentials.append(new_credentials)
+        user.save()
 
 
-def _ensure_auth_user(user_data: dict, company_id: str, log: Logger, revoke: bool = False) -> str:
+def _ensure_auth_user(
+    user_data: dict,
+    company_id: str,
+    log: Logger,
+    revoke: bool = False,
+    internal_user: bool = False,
+) -> str:
     user_id = user_data.get("id", f"__{user_data['name']}__")
     role = user_data["role"]
     email = user_data["email"]
@@ -40,12 +60,15 @@ def _ensure_auth_user(user_data: dict, company_id: str, log: Logger, revoke: boo
 
     user: AuthUser = AuthUser.objects(id=user_id).first()
     if user:
-        _ensure_user_credentials(user=user, key=key, secret=secret, log=log, revoke=revoke)
-        if (
-            user.role != role
-            or user.email != email
-            or user.autocreated != autocreated
-        ):
+        _ensure_user_credentials(
+            user=user,
+            key=key,
+            secret=secret,
+            log=log,
+            revoke=revoke,
+            internal_user=internal_user,
+        )
+        if user.role != role or user.email != email or user.autocreated != autocreated:
             user.email = email
             user.role = role
             user.autocreated = autocreated
@@ -54,9 +77,7 @@ def _ensure_auth_user(user_data: dict, company_id: str, log: Logger, revoke: boo
         return user.id
 
     credentials = (
-        [Credentials(key=key, secret=secret)]
-        if not revoke and key and secret
-        else []
+        [Credentials(key=key, secret=secret)] if not revoke and key and secret else []
     )
     log.info(f"Creating user: {user_data['name']}")
 
@@ -108,7 +129,9 @@ def ensure_fixed_user(user: FixedUser, log: Logger, emails: set):
         try:
             log.info(f"Updating user name: {user.name}")
             given_name, _, family_name = user.name.partition(" ")
-            db_user.update(name=user.name, given_name=given_name, family_name=family_name)
+            db_user.update(
+                name=user.name, given_name=given_name, family_name=family_name
+            )
         except Exception:
             pass
     else: