From 6cb79b580afe4c58da824cae583a7ceadba6db31 Mon Sep 17 00:00:00 2001
From: Tessa Walsh <tessa@bitarchivist.net>
Date: Tue, 31 Jan 2023 15:55:31 -0500
Subject: [PATCH] Fix issue where users are added to default org as admin
 (#534)

Users should only be added as to the default org with Owner permissions
if they are not specifically being invited to another org. This commit
fixes the logic in the post-registration callback to make this the case.
---
 backend/btrixcloud/invites.py |  2 +-
 backend/btrixcloud/orgs.py    |  2 +-
 backend/btrixcloud/users.py   | 24 +++++++++++++++++-------
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/backend/btrixcloud/invites.py b/backend/btrixcloud/invites.py
index 28742648..a0194d88 100644
--- a/backend/btrixcloud/invites.py
+++ b/backend/btrixcloud/invites.py
@@ -81,7 +81,7 @@ class InviteOps:
                 status_code=403, detail="This user has already been invited"
             )
 
-        # Invitations to a specific org via API must invite role, so if it's
+        # Invitations to a specific org via API must include role, so if it's
         # absent assume this is a general invitation from superadmin.
         if not new_user_invite.role:
             new_user_invite.role = UserRole.OWNER
diff --git a/backend/btrixcloud/orgs.py b/backend/btrixcloud/orgs.py
index 60e9e57f..e3b1dd11 100644
--- a/backend/btrixcloud/orgs.py
+++ b/backend/btrixcloud/orgs.py
@@ -278,7 +278,7 @@ class OrgOps:
         new_user_invite = await self.invites.get_valid_invite(invite_token, user.email)
         await self.add_user_by_invite(new_user_invite, user)
         await self.invites.remove_invite(invite_token)
-        return True
+        return new_user_invite
 
     async def add_user_by_invite(self, invite: InvitePending, user: User):
         """Add user to an org from an InvitePending, if any.
diff --git a/backend/btrixcloud/users.py b/backend/btrixcloud/users.py
index 48397b4e..f22eb332 100644
--- a/backend/btrixcloud/users.py
+++ b/backend/btrixcloud/users.py
@@ -236,6 +236,7 @@ class UserManager(BaseUserManager[UserCreate, UserDB]):
         """custom post registration callback, also receive the UserCreate object"""
 
         print(f"User {user.id} has registered.")
+        add_to_default_org = False
 
         if user_create.newOrg is True:
             print(f"Creating new organization for {user.id}")
@@ -249,25 +250,34 @@ class UserManager(BaseUserManager[UserCreate, UserDB]):
                 storage_name="default",
                 user=user,
             )
-        else:
-            default_org = await self.org_ops.get_default_org()
-            if default_org:
-                await self.org_ops.add_user_to_org(default_org, user.id)
 
         is_verified = hasattr(user_create, "is_verified") and user_create.is_verified
 
         if user_create.inviteToken:
+            new_user_invite = None
             try:
-                await self.org_ops.handle_new_user_invite(user_create.inviteToken, user)
+                new_user_invite = await self.org_ops.handle_new_user_invite(
+                    user_create.inviteToken, user
+                )
             except HTTPException as exc:
                 print(exc)
 
+            if new_user_invite and not new_user_invite.oid:
+                add_to_default_org = True
+
             if not is_verified:
                 # if user has been invited, mark as verified immediately
                 await self._update(user, {"is_verified": True})
 
-        elif not is_verified:
-            asyncio.create_task(self.request_verify(user, request))
+        else:
+            add_to_default_org = True
+            if not is_verified:
+                asyncio.create_task(self.request_verify(user, request))
+
+        if add_to_default_org:
+            default_org = await self.org_ops.get_default_org()
+            if default_org:
+                await self.org_ops.add_user_to_org(default_org, user.id)
 
     async def on_after_forgot_password(
         self, user: UserDB, token: str, request: Optional[Request] = None