fix: clean up two factor (#26901)

Co-authored-by: Surbhi <[email protected]>
Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

cypress/e2e/dashboard.cy.ts

@@ -379,6 +379,8 @@ describe('Dashboard', () => {
         cy.get('[data-attr="date-filter"]').click()
         cy.contains('span', 'Last 14 days').click()
 
+        cy.wait(2000)
+
         // insight meta should be updated to show new date range
         cy.get('h5').contains('Last 14 days').should('exist')
 

frontend/src/layout/GlobalModals.tsx

@@ -5,14 +5,12 @@ import { TimeSensitiveAuthenticationModal } from 'lib/components/TimeSensitiveAu
 import { UpgradeModal } from 'lib/components/UpgradeModal/UpgradeModal'
 import { TwoFactorSetupModal } from 'scenes/authentication/TwoFactorSetupModal'
 import { CreateOrganizationModal } from 'scenes/organization/CreateOrganizationModal'
-import { membersLogic } from 'scenes/organization/membersLogic'
 import { CreateEnvironmentModal } from 'scenes/project/CreateEnvironmentModal'
 import { CreateProjectModal } from 'scenes/project/CreateProjectModal'
 import { SessionPlayerModal } from 'scenes/session-recordings/player/modal/SessionPlayerModal'
 import { inviteLogic } from 'scenes/settings/organization/inviteLogic'
 import { InviteModal } from 'scenes/settings/organization/InviteModal'
 import { PreviewingCustomCssModal } from 'scenes/themes/PreviewingCustomCssModal'
-import { userLogic } from 'scenes/userLogic'
 
 import type { globalModalsLogicType } from './GlobalModalsType'
 
@@ -58,7 +56,6 @@ export function GlobalModals(): JSX.Element {
         useActions(globalModalsLogic)
     const { isInviteModalShown } = useValues(inviteLogic)
     const { hideInviteModal } = useActions(inviteLogic)
-    const { user } = useValues(userLogic)
 
     return (
         <>
@@ -71,17 +68,7 @@ export function GlobalModals(): JSX.Element {
             <TimeSensitiveAuthenticationModal />
             <SessionPlayerModal />
             <PreviewingCustomCssModal />
-            {user && user.organization?.enforce_2fa && !user.is_2fa_enabled && (
-                <TwoFactorSetupModal
-                    onSuccess={() => {
-                        userLogic.actions.loadUser()
-                        membersLogic.actions.loadAllMembers()
-                    }}
-                    forceOpen
-                    closable={false}
-                    required={true}
-                />
-            )}
+            <TwoFactorSetupModal />
             <HedgehogBuddyWithLogic />
         </>
     )

frontend/src/mocks/handlers.ts

@@ -119,6 +119,7 @@ export const defaultMocks: Mocks = {
                 },
             },
         ],
+        '/api/users/@me/two_factor_status/': () => [200, { is_enabled: true, backup_codes: [], method: 'TOTP' }],
         '/api/environments/@current/': MOCK_DEFAULT_TEAM,
         '/api/projects/@current/': MOCK_DEFAULT_TEAM,
         '/api/projects/:team_id/comments/count': { count: 0 },

frontend/src/scenes/authentication/TwoFactorSetupModal.tsx

@@ -1,46 +1,35 @@
 import { useActions, useValues } from 'kea'
 import { LemonBanner } from 'lib/lemon-ui/LemonBanner'
 import { LemonModal } from 'lib/lemon-ui/LemonModal'
+import { membersLogic } from 'scenes/organization/membersLogic'
+import { userLogic } from 'scenes/userLogic'
 
 import { twoFactorLogic } from './twoFactorLogic'
 import { TwoFactorSetup } from './TwoFactorSetup'
 
-interface TwoFactorSetupModalProps {
-    onSuccess: () => void
-    closable?: boolean
-    required?: boolean
-    forceOpen?: boolean
-}
-
-export function TwoFactorSetupModal({
-    onSuccess,
-    closable = true,
-    required = false,
-    forceOpen = false,
-}: TwoFactorSetupModalProps): JSX.Element {
-    const { isTwoFactorSetupModalOpen } = useValues(twoFactorLogic)
-    const { toggleTwoFactorSetupModal } = useActions(twoFactorLogic)
+export function TwoFactorSetupModal(): JSX.Element {
+    const { isTwoFactorSetupModalOpen, forceOpenTwoFactorSetupModal } = useValues(twoFactorLogic)
+    const { closeTwoFactorSetupModal } = useActions(twoFactorLogic)
 
     return (
         <LemonModal
             title="Set up two-factor authentication"
-            isOpen={isTwoFactorSetupModalOpen || forceOpen}
-            onClose={closable ? () => toggleTwoFactorSetupModal(false) : undefined}
-            closable={closable}
+            isOpen={isTwoFactorSetupModalOpen || forceOpenTwoFactorSetupModal}
+            onClose={!forceOpenTwoFactorSetupModal ? () => closeTwoFactorSetupModal() : undefined}
+            closable={!forceOpenTwoFactorSetupModal}
         >
             <div className="max-w-md">
-                {required && (
+                {forceOpenTwoFactorSetupModal && (
                     <LemonBanner className="mb-4" type="warning">
                         Your organization requires you to set up 2FA.
                     </LemonBanner>
                 )}
                 <p>Use an authenticator app like Google Authenticator or 1Password to scan the QR code below.</p>
                 <TwoFactorSetup
                     onSuccess={() => {
-                        toggleTwoFactorSetupModal(false)
-                        if (onSuccess) {
-                            onSuccess()
-                        }
+                        closeTwoFactorSetupModal()
+                        userLogic.actions.loadUser()
+                        membersLogic.actions.loadAllMembers()
                     }}
                 />
             </div>

frontend/src/scenes/authentication/twoFactorLogic.ts

@@ -4,7 +4,9 @@ import { forms } from 'kea-forms'
 import { loaders } from 'kea-loaders'
 import api from 'lib/api'
 import { featureFlagLogic } from 'lib/logic/featureFlagLogic'
+import { membersLogic } from 'scenes/organization/membersLogic'
 import { preflightLogic } from 'scenes/PreflightCheck/preflightLogic'
+import { userLogic } from 'scenes/userLogic'
 
 import type { twoFactorLogicType } from './twoFactorLogicType'
 
@@ -26,24 +28,33 @@ export const twoFactorLogic = kea<twoFactorLogicType>([
     path(['scenes', 'authentication', 'loginLogic']),
     props({} as TwoFactorLogicProps),
     connect({
-        values: [preflightLogic, ['preflight'], featureFlagLogic, ['featureFlags']],
+        values: [preflightLogic, ['preflight'], featureFlagLogic, ['featureFlags'], userLogic, ['user']],
+        actions: [userLogic, ['loadUser'], membersLogic, ['loadAllMembers']],
     }),
     actions({
         setGeneralError: (code: string, detail: string) => ({ code, detail }),
         clearGeneralError: true,
         loadStatus: true,
         generateBackupCodes: true,
         disable2FA: true,
-        toggleTwoFactorSetupModal: (open: boolean) => ({ open }),
+        openTwoFactorSetupModal: (forceOpen?: boolean) => ({ forceOpen }),
+        closeTwoFactorSetupModal: true,
         toggleDisable2FAModal: (open: boolean) => ({ open }),
         toggleBackupCodesModal: (open: boolean) => ({ open }),
-        startSetup: true,
     }),
     reducers({
         isTwoFactorSetupModalOpen: [
             false,
             {
-                toggleTwoFactorSetupModal: (_, { open }) => open,
+                openTwoFactorSetupModal: () => true,
+                closeTwoFactorSetupModal: () => false,
+            },
+        ],
+        forceOpenTwoFactorSetupModal: [
+            false,
+            {
+                openTwoFactorSetupModal: (_, { forceOpen }) => !!forceOpen,
+                closeTwoFactorSetupModal: () => false,
             },
         ],
         isDisable2FAModalOpen: [
@@ -89,11 +100,9 @@ export const twoFactorLogic = kea<twoFactorLogicType>([
         startSetup: [
             {},
             {
-                toggleTwoFactorSetupModal: async ({ open }, breakpoint) => {
-                    if (open) {
-                        breakpoint()
-                        await api.get('api/users/@me/two_factor_start_setup/')
-                    }
+                openTwoFactorSetupModal: async (_, breakpoint) => {
+                    breakpoint()
+                    await api.get('api/users/@me/two_factor_start_setup/')
                     return { status: 'completed' }
                 },
             },
@@ -144,6 +153,10 @@ export const twoFactorLogic = kea<twoFactorLogicType>([
                 await api.create<any>('api/users/@me/two_factor_disable/')
                 lemonToast.success('2FA disabled successfully')
                 actions.loadStatus()
+
+                // Refresh user and members
+                actions.loadUser()
+                actions.loadAllMembers()
             } catch (e) {
                 const { code, detail } = e as Record<string, any>
                 actions.setGeneralError(code, detail)
@@ -153,19 +166,17 @@ export const twoFactorLogic = kea<twoFactorLogicType>([
         generateBackupCodesSuccess: () => {
             lemonToast.success('Backup codes generated successfully')
         },
-        toggleTwoFactorSetupModal: ({ open }) => {
-            if (!open) {
-                // Clear the form when closing the modal
-                actions.resetToken()
-            }
-        },
-        startSetup: async () => {
-            await api.get('api/users/@me/two_factor_start_setup/')
+        closeTwoFactorSetupModal: () => {
+            // Clear the form when closing the modal
+            actions.resetToken()
         },
     })),
 
-    afterMount(({ actions }) => {
-        actions.startSetup()
+    afterMount(({ actions, values }) => {
         actions.loadStatus()
+
+        if (values.user && values.user.organization?.enforce_2fa && !values.user.is_2fa_enabled) {
+            actions.openTwoFactorSetupModal(true)
+        }
     }),
 ])

frontend/src/scenes/settings/organization/Members.tsx

@@ -19,7 +19,6 @@ import {
 } from 'lib/utils/permissioning'
 import { useEffect } from 'react'
 import { twoFactorLogic } from 'scenes/authentication/twoFactorLogic'
-import { TwoFactorSetupModal } from 'scenes/authentication/TwoFactorSetupModal'
 import { membersLogic } from 'scenes/organization/membersLogic'
 import { organizationLogic } from 'scenes/organizationLogic'
 import { preflightLogic } from 'scenes/PreflightCheck/preflightLogic'
@@ -143,9 +142,9 @@ export function Members(): JSX.Element | null {
     const { preflight } = useValues(preflightLogic)
     const { user } = useValues(userLogic)
 
-    const { setSearch, ensureAllMembersLoaded, loadAllMembers } = useActions(membersLogic)
+    const { setSearch, ensureAllMembersLoaded } = useActions(membersLogic)
     const { updateOrganization } = useActions(organizationLogic)
-    const { toggleTwoFactorSetupModal } = useActions(twoFactorLogic)
+    const { openTwoFactorSetupModal } = useActions(twoFactorLogic)
 
     useEffect(() => {
         ensureAllMembersLoaded()
@@ -212,14 +211,6 @@ export function Members(): JSX.Element | null {
             render: function LevelRender(_, member) {
                 return (
                     <>
-                        {member.user.uuid == user.uuid && (
-                            <TwoFactorSetupModal
-                                onSuccess={() => {
-                                    userLogic.actions.updateUser({})
-                                    loadAllMembers()
-                                }}
-                            />
-                        )}
                         <Tooltip
                             title={
                                 member.user.uuid == user.uuid && !member.is_2fa_enabled
@@ -230,7 +221,7 @@ export function Members(): JSX.Element | null {
                             <LemonTag
                                 onClick={
                                     member.user.uuid == user.uuid && !member.is_2fa_enabled
-                                        ? () => toggleTwoFactorSetupModal(true)
+                                        ? () => openTwoFactorSetupModal()
                                         : undefined
                                 }
                                 data-attr="2fa-enabled"

frontend/src/scenes/settings/user/TwoFactorSettings.tsx

@@ -3,7 +3,6 @@ import { LemonButton, LemonModal } from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { copyToClipboard } from 'lib/utils/copyToClipboard'
 import { twoFactorLogic } from 'scenes/authentication/twoFactorLogic'
-import { TwoFactorSetupModal } from 'scenes/authentication/TwoFactorSetupModal'
 import { membersLogic } from 'scenes/organization/membersLogic'
 import { userLogic } from 'scenes/userLogic'
 
@@ -13,13 +12,8 @@ export function TwoFactorSettings(): JSX.Element {
 
     const { updateUser } = useActions(userLogic)
     const { loadMemberUpdates } = useActions(membersLogic)
-    const {
-        generateBackupCodes,
-        disable2FA,
-        toggleTwoFactorSetupModal,
-        toggleDisable2FAModal,
-        toggleBackupCodesModal,
-    } = useActions(twoFactorLogic)
+    const { generateBackupCodes, disable2FA, openTwoFactorSetupModal, toggleDisable2FAModal, toggleBackupCodesModal } =
+        useActions(twoFactorLogic)
 
     const handleSuccess = (): void => {
         updateUser({})
@@ -28,8 +22,6 @@ export function TwoFactorSettings(): JSX.Element {
 
     return (
         <div className="flex flex-col items-start">
-            <TwoFactorSetupModal onSuccess={handleSuccess} />
-
             {isDisable2FAModalOpen && (
                 <LemonModal
                     title="Disable 2FA"
@@ -118,7 +110,7 @@ export function TwoFactorSettings(): JSX.Element {
                         <IconWarning color="orange" className="text-xl" />
                         <span className="font-medium">2FA is not enabled</span>
                     </div>
-                    <LemonButton type="primary" onClick={() => toggleTwoFactorSetupModal(true)}>
+                    <LemonButton type="primary" onClick={() => openTwoFactorSetupModal()}>
                         Set up 2FA
                     </LemonButton>
                 </div>