Add capability of adding MFA to the new auth device wizard

web/packages/teleport/src/Account/Account.story.tsx

@@ -161,7 +161,6 @@ const props: AccountProps = {
       residentKey: false,
     },
   ],
-  onAddPasskey: () => {},
   onPasskeyAdded: () => {},
   isReauthenticationRequired: false,
   passkeyWizardVisible: false,

web/packages/teleport/src/Account/Account.tsx

@@ -34,8 +34,6 @@ import { MfaChallengeScope } from 'teleport/services/auth/auth';
 
 import cfg from 'teleport/config';
 
-import { storageService } from 'teleport/services/storageService';
-
 import { AuthDeviceList } from './ManageDevices/AuthDeviceList/AuthDeviceList';
 import useManageDevices, {
   State as ManageDevicesState,
@@ -45,9 +43,6 @@ import { ActionButton, Header } from './Header';
 import { PasswordBox } from './PasswordBox';
 import { AddAuthDeviceWizard } from './ManageDevices/AddAuthDeviceWizard';
 
-const useNewAddAuthDeviceDialog =
-  storageService.isNewAddAuthDeviceDialogEnabled();
-
 export interface EnterpriseComponentProps {
   // TODO(bl-nero): Consider moving the notifications to its own store and
   // unifying them between this screen and the unified resources screen.
@@ -110,7 +105,6 @@ export function Account({
   setToken,
   onAddDevice,
   onRemoveDevice,
-  onAddPasskey,
   onPasskeyAdded,
   deviceToRemove,
   fetchDevices,
@@ -206,11 +200,7 @@ export function Account({
                         ? 'Passwordless authentication is disabled'
                         : ''
                     }
-                    onClick={() =>
-                      useNewAddAuthDeviceDialog
-                        ? onAddPasskey()
-                        : onAddDevice('passwordless')
-                    }
+                    onClick={() => onAddDevice('passwordless')}
                   >
                     <Icon.Add size={20} />
                     Add a Passkey
@@ -293,6 +283,8 @@ export function Account({
 
       {passkeyWizardVisible && (
         <AddAuthDeviceWizard
+          usage={restrictNewDeviceUsage}
+          auth2faType={cfg.getAuth2faType()}
           privilegeToken={token}
           onClose={closePasskeyWizard}
           onSuccess={onAddPasskeySuccess}

web/packages/teleport/src/Account/ManageDevices/AddAuthDeviceWizard/AddAuthDeviceWizard.story.tsx

@@ -0,0 +1,154 @@
+/**
+ * Teleport
+ * Copyright (C) 2024 Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+import React from 'react';
+
+import { Auth2faType } from 'shared/services';
+
+import Dialog from 'design/Dialog';
+
+import { initialize, mswLoader } from 'msw-storybook-addon';
+
+import { rest } from 'msw';
+
+import { DeviceUsage } from 'teleport/services/mfa';
+import { createTeleportContext } from 'teleport/mocks/contexts';
+import { ContextProvider } from 'teleport/index';
+
+import cfg from 'teleport/config';
+
+import {
+  CreateDeviceStep,
+  ReauthenticateStep,
+  SaveDeviceStep,
+} from './AddAuthDeviceWizard';
+
+export default {
+  title: 'teleport/Account/Manage Devices/Add Device Wizard',
+  loaders: [mswLoader],
+  decorators: [
+    Story => {
+      const ctx = createTeleportContext();
+      return (
+        <ContextProvider ctx={ctx}>
+          <Dialog open={true} dialogCss={() => ({ width: '650px' })}>
+            <Story />
+          </Dialog>
+        </ContextProvider>
+      );
+    },
+  ],
+};
+
+initialize();
+
+export function Reauthenticate() {
+  return <ReauthenticateStep {...stepProps} />;
+}
+
+export function CreatePasskey() {
+  return <CreateDeviceStep {...stepProps} usage="passwordless" />;
+}
+
+export function CreateMfaHardwareDevice() {
+  return (
+    <CreateDeviceStep {...stepProps} usage="mfa" newMfaDeviceType="webauthn" />
+  );
+}
+
+export function CreateMfaAppQrCodeLoading() {
+  return <CreateDeviceStep {...stepProps} usage="mfa" newMfaDeviceType="otp" />;
+}
+CreateMfaAppQrCodeLoading.parameters = {
+  msw: {
+    handlers: [
+      rest.post(
+        cfg.getMfaCreateRegistrationChallengeUrl('privilege-token'),
+        (req, res, ctx) => res(ctx.delay('infinite'))
+      ),
+    ],
+  },
+};
+
+export function CreateMfaAppQrCodeFailed() {
+  return <CreateDeviceStep {...stepProps} usage="mfa" newMfaDeviceType="otp" />;
+}
+CreateMfaAppQrCodeFailed.parameters = {
+  msw: {
+    handlers: [
+      rest.post(
+        cfg.getMfaCreateRegistrationChallengeUrl('privilege-token'),
+        (req, res, ctx) => res(ctx.status(500))
+      ),
+    ],
+  },
+};
+
+const dummyQrCode =
+  'iVBORw0KGgoAAAANSUhEUgAAAB0AAAAdAQMAAABsXfVMAAAABlBMVEUAAAD///+l2Z/dAAAAAnRSTlP//8i138cAAAAJcEhZcwAACxIAAAsSAdLdfvwAAABrSURBVAiZY/gPBAxoxAcxh3qG71fv1zN8iQ8EEReBRACQ+H4ZKPZBFCj7/3v9f4aPU9vqGX4kFtUzfG5mBLK2aNUz/PM3AsmqAk2RNQTquLYLqDdG/z/QlGAgES4CFLu4GygrXF2Pbi+IAADZqFQFAjXZWgAAAABJRU5ErkJggg==';
+
+export function CreateMfaApp() {
+  return <CreateDeviceStep {...stepProps} usage="mfa" newMfaDeviceType="otp" />;
+}
+CreateMfaApp.parameters = {
+  msw: {
+    handlers: [
+      rest.post(
+        cfg.getMfaCreateRegistrationChallengeUrl('privilege-token'),
+        (req, res, ctx) => res(ctx.json({ totp: { qrCode: dummyQrCode } }))
+      ),
+    ],
+  },
+};
+
+export function SavePasskey() {
+  return <SaveDeviceStep {...stepProps} usage="passwordless" />;
+}
+
+export function SaveMfaHardwareDevice() {
+  return (
+    <SaveDeviceStep {...stepProps} usage="mfa" newMfaDeviceType="webauthn" />
+  );
+}
+
+export function SaveMfaAuthenticatorApp() {
+  return <SaveDeviceStep {...stepProps} usage="mfa" newMfaDeviceType="otp" />;
+}
+
+const stepProps = {
+  // StepComponentProps
+  next: () => {},
+  prev: () => {},
+  hasTransitionEnded: true,
+  stepIndex: 0,
+  flowLength: 1,
+  refCallback: () => {},
+
+  // Other props
+  privilegeToken: 'privilege-token',
+  usage: 'passwordless' as DeviceUsage,
+  auth2faType: 'optional' as Auth2faType,
+  credential: { id: 'cred-id', type: 'public-key' },
+  newMfaDeviceType: 'webauthn' as Auth2faType,
+  onNewMfaDeviceTypeChange: () => {},
+  onDeviceCreated: () => {},
+  onAuthenticated: () => {},
+  onClose: () => {},
+  onPasskeyCreated: () => {},
+  onSuccess: () => {},
+};

web/packages/teleport/src/Account/ManageDevices/AddAuthDeviceWizard/AddAuthDeviceWizard.test.tsx

@@ -16,29 +16,35 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import { render, screen, userEvent } from 'design/utils/testing';
+import { render, screen } from 'design/utils/testing';
 import React from 'react';
 
 import { within } from '@testing-library/react';
+import { userEvent, UserEvent } from '@testing-library/user-event';
 
 import TeleportContext from 'teleport/teleportContext';
 import { ContextProvider } from 'teleport';
-import MfaService from 'teleport/services/mfa';
-import cfg from 'teleport/config';
+import MfaService, { DeviceUsage } from 'teleport/services/mfa';
 import auth from 'teleport/services/auth/auth';
 
 import { AddAuthDeviceWizard } from '.';
 
 const dummyCredential: Credential = { id: 'cred-id', type: 'public-key' };
+let ctx: TeleportContext;
+let user: UserEvent;
+let onSuccess: jest.Mock;
 
 beforeEach(() => {
-  jest.replaceProperty(cfg.auth, 'second_factor', 'optional');
+  ctx = new TeleportContext();
+  user = userEvent.setup();
+  onSuccess = jest.fn();
+
   jest
     .spyOn(MfaService.prototype, 'createNewWebAuthnDevice')
     .mockResolvedValueOnce(dummyCredential);
   jest
     .spyOn(MfaService.prototype, 'saveNewWebAuthnDevice')
-    .mockResolvedValueOnce('some-credential');
+    .mockResolvedValueOnce(undefined);
   jest
     .spyOn(auth, 'createPrivilegeTokenWithWebauthn')
     .mockResolvedValueOnce('webauthn-privilege-token');
@@ -47,22 +53,29 @@ beforeEach(() => {
     .mockImplementationOnce(token =>
       Promise.resolve(`totp-privilege-token-${token}`)
     );
+  jest.spyOn(auth, 'createMfaRegistrationChallenge').mockResolvedValueOnce({
+    qrCode: 'dummy-qr-code',
+    webauthnPublicKey: {} as PublicKeyCredentialCreationOptions,
+  });
+  jest
+    .spyOn(MfaService.prototype, 'addNewTotpDevice')
+    .mockResolvedValueOnce(undefined);
 });
 
 afterEach(jest.resetAllMocks);
 
 function TestWizard({
-  ctx,
   privilegeToken,
-  onSuccess,
+  usage,
 }: {
-  ctx: TeleportContext;
   privilegeToken?: string;
-  onSuccess(): void;
+  usage: DeviceUsage;
 }) {
   return (
     <ContextProvider ctx={ctx}>
       <AddAuthDeviceWizard
+        usage={usage}
+        auth2faType="optional"
         privilegeToken={privilegeToken}
         onClose={() => {}}
         onSuccess={onSuccess}
@@ -73,15 +86,8 @@ function TestWizard({
 
 describe('flow without reauthentication', () => {
   test('adds a passkey', async () => {
-    const ctx = new TeleportContext();
-    const user = userEvent.setup();
-    const onSuccess = jest.fn();
     render(
-      <TestWizard
-        ctx={ctx}
-        onSuccess={onSuccess}
-        privilegeToken="privilege-token"
-      />
+      <TestWizard usage="passwordless" privilegeToken="privilege-token" />
     );
 
     const createStep = within(screen.getByTestId('create-step'));
@@ -108,14 +114,67 @@ describe('flow without reauthentication', () => {
     });
     expect(onSuccess).toHaveBeenCalled();
   });
+
+  test('adds a WebAuthn MFA', async () => {
+    render(<TestWizard usage="mfa" privilegeToken="privilege-token" />);
+
+    const createStep = within(screen.getByTestId('create-step'));
+    await user.click(createStep.getByLabelText('Hardware Device'));
+    await user.click(
+      createStep.getByRole('button', { name: 'Create an MFA method' })
+    );
+    expect(ctx.mfaService.createNewWebAuthnDevice).toHaveBeenCalledWith({
+      tokenId: 'privilege-token',
+      deviceUsage: 'mfa',
+    });
+
+    const saveStep = within(screen.getByTestId('save-step'));
+    await user.type(saveStep.getByLabelText('MFA Method Name'), 'new-mfa');
+    await user.click(
+      saveStep.getByRole('button', { name: 'Save the MFA method' })
+    );
+    expect(ctx.mfaService.saveNewWebAuthnDevice).toHaveBeenCalledWith({
+      credential: dummyCredential,
+      addRequest: {
+        deviceName: 'new-mfa',
+        deviceUsage: 'mfa',
+        tokenId: 'privilege-token',
+      },
+    });
+    expect(onSuccess).toHaveBeenCalled();
+  });
+
+  test('adds an authenticator app', async () => {
+    render(<TestWizard usage="mfa" privilegeToken="privilege-token" />);
+
+    const createStep = within(screen.getByTestId('create-step'));
+    await user.click(createStep.getByLabelText('Authenticator App'));
+    expect(createStep.getByRole('img')).toHaveAttribute(
+      'src',
+      'data:image/png;base64,dummy-qr-code'
+    );
+    await user.click(
+      createStep.getByRole('button', { name: 'Create an MFA method' })
+    );
+
+    const saveStep = within(screen.getByTestId('save-step'));
+    await user.type(saveStep.getByLabelText('MFA Method Name'), 'new-mfa');
+    await user.type(saveStep.getByLabelText(/Authenticator Code/), '345678');
+    await user.click(
+      saveStep.getByRole('button', { name: 'Save the MFA method' })
+    );
+    expect(ctx.mfaService.addNewTotpDevice).toHaveBeenCalledWith({
+      tokenId: 'privilege-token',
+      secondFactorToken: '345678',
+      deviceName: 'new-mfa',
+    });
+    expect(onSuccess).toHaveBeenCalled();
+  });
 });
 
 describe('flow with reauthentication', () => {
   test('adds a passkey with WebAuthn reauthentication', async () => {
-    const ctx = new TeleportContext();
-    const user = userEvent.setup();
-    const onSuccess = jest.fn();
-    render(<TestWizard ctx={ctx} onSuccess={onSuccess} />);
+    render(<TestWizard usage="passwordless" />);
 
     const reauthenticateStep = within(
       screen.getByTestId('reauthenticate-step')
@@ -148,10 +207,7 @@ describe('flow with reauthentication', () => {
   });
 
   test('adds a passkey with OTP reauthentication', async () => {
-    const ctx = new TeleportContext();
-    const user = userEvent.setup();
-    const onSuccess = jest.fn();
-    render(<TestWizard ctx={ctx} onSuccess={onSuccess} />);
+    render(<TestWizard usage="passwordless" />);
 
     const reauthenticateStep = within(
       screen.getByTestId('reauthenticate-step')