feat: make 2fa great again (#26557)

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

frontend/src/layout/GlobalModals.tsx

@@ -1,10 +1,9 @@
-import { LemonModal } from '@posthog/lemon-ui'
 import { actions, kea, path, reducers, useActions, useValues } from 'kea'
 import { ConfirmUpgradeModal } from 'lib/components/ConfirmUpgradeModal/ConfirmUpgradeModal'
 import { HedgehogBuddyWithLogic } from 'lib/components/HedgehogBuddy/HedgehogBuddyWithLogic'
 import { TimeSensitiveAuthenticationModal } from 'lib/components/TimeSensitiveAuthentication/TimeSensitiveAuthentication'
 import { UpgradeModal } from 'lib/components/UpgradeModal/UpgradeModal'
-import { Setup2FA } from 'scenes/authentication/Setup2FA'
+import { TwoFactorSetupModal } from 'scenes/authentication/TwoFactorSetupModal'
 import { CreateOrganizationModal } from 'scenes/organization/CreateOrganizationModal'
 import { membersLogic } from 'scenes/organization/membersLogic'
 import { CreateEnvironmentModal } from 'scenes/project/CreateEnvironmentModal'
@@ -73,22 +72,15 @@ export function GlobalModals(): JSX.Element {
             <SessionPlayerModal />
             <PreviewingCustomCssModal />
             {user && user.organization?.enforce_2fa && !user.is_2fa_enabled && (
-                <LemonModal title="Set up 2FA" closable={false}>
-                    <p>
-                        <b>Your organization requires you to set up 2FA.</b>
-                    </p>
-                    <p>
-                        <b>
-                            Use an authenticator app like Google Authenticator or 1Password to scan the QR code below.
-                        </b>
-                    </p>
-                    <Setup2FA
-                        onSuccess={() => {
-                            userLogic.actions.loadUser()
-                            membersLogic.actions.loadAllMembers()
-                        }}
-                    />
-                </LemonModal>
+                <TwoFactorSetupModal
+                    onSuccess={() => {
+                        userLogic.actions.loadUser()
+                        membersLogic.actions.loadAllMembers()
+                    }}
+                    forceOpen
+                    closable={false}
+                    required={true}
+                />
             )}
             <HedgehogBuddyWithLogic />
         </>

frontend/src/lib/constants.tsx

@@ -233,7 +233,6 @@ export const FEATURE_FLAGS = {
     EXPERIMENTS_MULTIPLE_METRICS: 'experiments-multiple-metrics', // owner: @jurajmajerik #team-experiments
     WEB_ANALYTICS_WARN_CUSTOM_EVENT_NO_SESSION: 'web-analytics-warn-custom-event-no-session', // owner: @robbie-c #team-web-analytics
     REMOTE_CONFIG: 'remote-config', // owner: @benjackwhite
-    TWO_FACTOR_UI: 'two-factor-ui', // owner: @zach
     SITE_DESTINATIONS: 'site-destinations', // owner: @mariusandra #team-cdp
     SITE_APP_FUNCTIONS: 'site-app-functions', // owner: @mariusandra #team-cdp
     REPLAY_HOGQL_FILTERS: 'replay-hogql-filters', // owner: @pauldambra #team-replay

frontend/src/scenes/authentication/Login2FA.tsx

@@ -25,7 +25,7 @@ export function Login2FA(): JSX.Element {
         >
             <div className="space-y-2">
                 <h2>Two-Factor Authentication</h2>
-                <p>Enter a token from your authenticator app.</p>
+                <p>Enter a token from your authenticator app or a backup code.</p>
 
                 <Form logic={login2FALogic} formKey="twofactortoken" enableFormOnSubmit className="space-y-4">
                     {generalError && <LemonBanner type="error">{generalError.detail}</LemonBanner>}

frontend/src/scenes/authentication/Setup2FA.tsx → frontend/src/scenes/authentication/TwoFactorSetup.tsx

@@ -6,17 +6,17 @@ import { Form } from 'kea-forms'
 import { LemonBanner } from 'lib/lemon-ui/LemonBanner'
 import { LemonField } from 'lib/lemon-ui/LemonField'
 
-import { setup2FALogic } from './setup2FALogic'
+import { twoFactorLogic } from './twoFactorLogic'
 
-export function Setup2FA({ onSuccess }: { onSuccess: () => void }): JSX.Element | null {
-    const { startSetupLoading, generalError } = useValues(setup2FALogic({ onSuccess }))
+export function TwoFactorSetup({ onSuccess }: { onSuccess: () => void }): JSX.Element | null {
+    const { startSetupLoading, generalError } = useValues(twoFactorLogic({ onSuccess }))
     if (startSetupLoading) {
         return null
     }
 
     return (
         <>
-            <Form logic={setup2FALogic} formKey="token" enableFormOnSubmit className="flex flex-col space-y-4">
+            <Form logic={twoFactorLogic} formKey="token" enableFormOnSubmit className="flex flex-col space-y-4">
                 <div className="bg-white ml-auto mr-auto mt-2">
                     <img src="/account/two_factor/qrcode/" className="Setup2FA__image" />
                 </div>

frontend/src/scenes/authentication/TwoFactorSetupModal.tsx

@@ -0,0 +1,49 @@
+import { useActions, useValues } from 'kea'
+import { LemonBanner } from 'lib/lemon-ui/LemonBanner'
+import { LemonModal } from 'lib/lemon-ui/LemonModal'
+
+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)
+
+    return (
+        <LemonModal
+            title="Set up two-factor authentication"
+            isOpen={isTwoFactorSetupModalOpen || forceOpen}
+            onClose={closable ? () => toggleTwoFactorSetupModal(false) : undefined}
+            closable={closable}
+        >
+            <div className="max-w-md">
+                {required && (
+                    <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()
+                        }
+                    }}
+                />
+            </div>
+        </LemonModal>
+    )
+}

frontend/src/scenes/authentication/login2FALogic.ts

@@ -45,11 +45,7 @@ export const login2FALogic = kea<login2FALogicType>([
         twofactortoken: {
             defaults: { token: '' } as TwoFactorForm,
             errors: ({ token }) => ({
-                token: !token
-                    ? 'Please enter a token to continue'
-                    : token.length !== 6 || isNaN(parseInt(token))
-                    ? 'A token must consist of 6 digits'
-                    : null,
+                token: !token ? 'Please enter a token to continue' : null,
             }),
             submit: async ({ token }, breakpoint) => {
                 breakpoint()

frontend/src/scenes/authentication/setup2FALogic.ts → frontend/src/scenes/authentication/twoFactorLogic.ts

@@ -6,7 +6,7 @@ import api from 'lib/api'
 import { featureFlagLogic } from 'lib/logic/featureFlagLogic'
 import { preflightLogic } from 'scenes/PreflightCheck/preflightLogic'
 
-import type { setup2FALogicType } from './setup2FALogicType'
+import type { twoFactorLogicType } from './twoFactorLogicType'
 
 export interface TwoFactorForm {
     token: number | null
@@ -18,24 +18,45 @@ export interface TwoFactorStatus {
     method: string | null
 }
 
-export interface Setup2FALogicProps {
+export interface TwoFactorLogicProps {
     onSuccess?: () => void
 }
 
-export const setup2FALogic = kea<setup2FALogicType>([
+export const twoFactorLogic = kea<twoFactorLogicType>([
     path(['scenes', 'authentication', 'loginLogic']),
-    props({} as Setup2FALogicProps),
+    props({} as TwoFactorLogicProps),
     connect({
         values: [preflightLogic, ['preflight'], featureFlagLogic, ['featureFlags']],
     }),
     actions({
         setGeneralError: (code: string, detail: string) => ({ code, detail }),
         clearGeneralError: true,
-        setup: true,
         loadStatus: true,
         generateBackupCodes: true,
+        disable2FA: true,
+        toggleTwoFactorSetupModal: (open: boolean) => ({ open }),
+        toggleDisable2FAModal: (open: boolean) => ({ open }),
+        toggleBackupCodesModal: (open: boolean) => ({ open }),
     }),
     reducers({
+        isTwoFactorSetupModalOpen: [
+            false,
+            {
+                toggleTwoFactorSetupModal: (_, { open }) => open,
+            },
+        ],
+        isDisable2FAModalOpen: [
+            false,
+            {
+                toggleDisable2FAModal: (_, { open }) => open,
+            },
+        ],
+        isBackupCodesModalOpen: [
+            false,
+            {
+                toggleBackupCodesModal: (_, { open }) => open,
+            },
+        ],
         generalError: [
             null as { code: string; detail: string } | null,
             {
@@ -67,9 +88,11 @@ export const setup2FALogic = kea<setup2FALogicType>([
         startSetup: [
             {},
             {
-                setup: async (_, breakpoint) => {
-                    breakpoint()
-                    await api.get('api/users/@me/start_2fa_setup/')
+                toggleTwoFactorSetupModal: async ({ open }, breakpoint) => {
+                    if (open) {
+                        breakpoint()
+                        await api.get('api/users/@me/two_factor_start_setup/')
+                    }
                     return { status: 'completed' }
                 },
             },
@@ -90,20 +113,6 @@ export const setup2FALogic = kea<setup2FALogicType>([
                 },
             },
         ],
-        disable2FA: [
-            false,
-            {
-                disable2FA: async () => {
-                    try {
-                        await api.create<any>('api/users/@me/two_factor_disable/')
-                        return true
-                    } catch (e) {
-                        const { code, detail } = e as Record<string, any>
-                        throw { code, detail }
-                    }
-                },
-            },
-        ],
     })),
     forms(({ actions }) => ({
         token: {
@@ -114,7 +123,7 @@ export const setup2FALogic = kea<setup2FALogicType>([
             submit: async ({ token }, breakpoint) => {
                 breakpoint()
                 try {
-                    return await api.create<any>('api/users/@me/validate_2fa/', { token })
+                    return await api.create<any>('api/users/@me/two_factor_validate/', { token })
                 } catch (e) {
                     const { code, detail } = e as Record<string, any>
                     actions.setGeneralError(code, detail)
@@ -129,16 +138,29 @@ export const setup2FALogic = kea<setup2FALogicType>([
             actions.loadStatus()
             props.onSuccess?.()
         },
-        disable2FASuccess: () => {
-            lemonToast.success('2FA disabled successfully')
+        disable2FA: async () => {
+            try {
+                await api.create<any>('api/users/@me/two_factor_disable/')
+                lemonToast.success('2FA disabled successfully')
+                actions.loadStatus()
+            } catch (e) {
+                const { code, detail } = e as Record<string, any>
+                actions.setGeneralError(code, detail)
+                throw e
+            }
         },
         generateBackupCodesSuccess: () => {
             lemonToast.success('Backup codes generated successfully')
         },
+        toggleTwoFactorSetupModal: ({ open }) => {
+            if (!open) {
+                // Clear the form when closing the modal
+                actions.resetToken()
+            }
+        },
     })),
 
     afterMount(({ actions }) => {
-        actions.setup()
         actions.loadStatus()
     }),
 ])

frontend/src/scenes/settings/SettingsMap.tsx

@@ -59,7 +59,7 @@ import { HedgehogModeSettings } from './user/HedgehogModeSettings'
 import { OptOutCapture } from './user/OptOutCapture'
 import { PersonalAPIKeys } from './user/PersonalAPIKeys'
 import { ThemeSwitcher } from './user/ThemeSwitcher'
-import { TwoFactorAuthentication } from './user/TwoFactorAuthentication'
+import { TwoFactorSettings } from './user/TwoFactorSettings'
 import { UpdateEmailPreferences } from './user/UpdateEmailPreferences'
 import { UserDetails } from './user/UserDetails'
 
@@ -473,7 +473,7 @@ export const SETTINGS_MAP: SettingSection[] = [
             {
                 id: '2fa',
                 title: 'Two-factor authentication',
-                component: <TwoFactorAuthentication />,
+                component: <TwoFactorSettings />,
             },
         ],
     },

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

@@ -1,4 +1,4 @@
-import { LemonInput, LemonModal, LemonSwitch } from '@posthog/lemon-ui'
+import { LemonInput, LemonSwitch } from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { PayGateMini } from 'lib/components/PayGateMini/PayGateMini'
 import { TZLabel } from 'lib/components/TZLabel'
@@ -17,8 +17,9 @@ import {
     membershipLevelToName,
     organizationMembershipLevelIntegers,
 } from 'lib/utils/permissioning'
-import { useEffect, useState } from 'react'
-import { Setup2FA } from 'scenes/authentication/Setup2FA'
+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'
@@ -138,13 +139,14 @@ function ActionsComponent(_: any, member: OrganizationMemberType): JSX.Element |
 
 export function Members(): JSX.Element | null {
     const { filteredMembers, membersLoading, search } = useValues(membersLogic)
-    const { setSearch, ensureAllMembersLoaded, loadAllMembers } = useActions(membersLogic)
     const { currentOrganization } = useValues(organizationLogic)
-    const { updateOrganization } = useActions(organizationLogic)
-    const [is2FAModalVisible, set2FAModalVisible] = useState(false)
     const { preflight } = useValues(preflightLogic)
     const { user } = useValues(userLogic)
 
+    const { setSearch, ensureAllMembersLoaded, loadAllMembers } = useActions(membersLogic)
+    const { updateOrganization } = useActions(organizationLogic)
+    const { toggleTwoFactorSetupModal } = useActions(twoFactorLogic)
+
     useEffect(() => {
         ensureAllMembersLoaded()
     }, [])
@@ -210,16 +212,13 @@ export function Members(): JSX.Element | null {
             render: function LevelRender(_, member) {
                 return (
                     <>
-                        {member.user.uuid == user.uuid && is2FAModalVisible && (
-                            <LemonModal title="Set up or manage 2FA" onClose={() => set2FAModalVisible(false)}>
-                                <Setup2FA
-                                    onSuccess={() => {
-                                        set2FAModalVisible(false)
-                                        userLogic.actions.updateUser({})
-                                        loadAllMembers()
-                                    }}
-                                />
-                            </LemonModal>
+                        {member.user.uuid == user.uuid && (
+                            <TwoFactorSetupModal
+                                onSuccess={() => {
+                                    userLogic.actions.updateUser({})
+                                    loadAllMembers()
+                                }}
+                            />
                         )}
                         <Tooltip
                             title={
@@ -231,7 +230,7 @@ export function Members(): JSX.Element | null {
                             <LemonTag
                                 onClick={
                                     member.user.uuid == user.uuid && !member.is_2fa_enabled
-                                        ? () => set2FAModalVisible(true)
+                                        ? () => toggleTwoFactorSetupModal(true)
                                         : undefined
                                 }
                                 data-attr="2fa-enabled"