From 437ba4375f68e16aa2fb7274a2a8087d99869656 Mon Sep 17 00:00:00 2001
From: peintnermax <max@caos.ch>
Date: Fri, 12 Apr 2024 16:38:38 +0200
Subject: [PATCH] mfa selection, setup page

---
 apps/login/app/(login)/mfa/create/page.tsx | 35 --------------------
 apps/login/app/(login)/mfa/set/page.tsx    |  4 +--
 apps/login/app/(login)/totp/set/page.tsx   | 37 ++++++++++++++++++++++
 apps/login/app/api/session/route.ts        | 22 +++++++++++--
 apps/login/lib/zitadel.ts                  | 13 ++++++--
 apps/login/ui/PasswordForm.tsx             | 19 ++++++++++-
 packages/zitadel-server/src/index.ts       |  2 ++
 7 files changed, 89 insertions(+), 43 deletions(-)
 delete mode 100644 apps/login/app/(login)/mfa/create/page.tsx
 create mode 100644 apps/login/app/(login)/totp/set/page.tsx

diff --git a/apps/login/app/(login)/mfa/create/page.tsx b/apps/login/app/(login)/mfa/create/page.tsx
deleted file mode 100644
index c0ce2fa4d8d..00000000000
--- a/apps/login/app/(login)/mfa/create/page.tsx
+++ /dev/null
@@ -1,35 +0,0 @@
-"use client";
-import { Button, ButtonVariants } from "#/ui/Button";
-import { TextInput } from "#/ui/Input";
-import UserAvatar from "#/ui/UserAvatar";
-import { useRouter } from "next/navigation";
-
-export default function Page() {
-  const router = useRouter();
-
-  return (
-    <div className="flex flex-col items-center space-y-4">
-      <h1>Password</h1>
-      <p className="ztdl-p mb-6 block">Enter your password.</p>
-
-      <UserAvatar
-        showDropdown
-        displayName="Max Peintner"
-        loginName="max@zitadel.com"
-      ></UserAvatar>
-
-      <div className="w-full">
-        <TextInput type="password" label="Password" />
-      </div>
-      <div className="flex w-full flex-row items-center justify-between">
-        <Button
-          onClick={() => router.back()}
-          variant={ButtonVariants.Secondary}
-        >
-          back
-        </Button>
-        <Button variant={ButtonVariants.Primary}>continue</Button>
-      </div>
-    </div>
-  );
-}
diff --git a/apps/login/app/(login)/mfa/set/page.tsx b/apps/login/app/(login)/mfa/set/page.tsx
index 57b7e625d18..cda98ee4a37 100644
--- a/apps/login/app/(login)/mfa/set/page.tsx
+++ b/apps/login/app/(login)/mfa/set/page.tsx
@@ -9,8 +9,8 @@ export default function Page() {
 
   return (
     <div className="flex flex-col items-center space-y-4">
-      <h1>Password</h1>
-      <p className="ztdl-p mb-6 block">Enter your password.</p>
+      <h1>Second Factor</h1>
+      <p className="ztdl-p mb-6 block">Please select a second factor.</p>
 
       <UserAvatar
         showDropdown
diff --git a/apps/login/app/(login)/totp/set/page.tsx b/apps/login/app/(login)/totp/set/page.tsx
new file mode 100644
index 00000000000..f457bcebed3
--- /dev/null
+++ b/apps/login/app/(login)/totp/set/page.tsx
@@ -0,0 +1,37 @@
+import { getBrandingSettings, getLoginSettings, server } from "#/lib/zitadel";
+import DynamicTheme from "#/ui/DynamicTheme";
+import TOTPForm from "#/ui/TOTPForm";
+
+export default async function Page({
+  searchParams,
+}: {
+  searchParams: Record<string | number | symbol, string | undefined>;
+}) {
+  const { loginName, authRequestId, sessionId, organization, code, submit } =
+    searchParams;
+
+  const branding = await getBrandingSettings(server, organization);
+  const loginSettings = await getLoginSettings(server, organization);
+
+  return (
+    <DynamicTheme branding={branding}>
+      <div className="flex flex-col items-center space-y-4">
+        <h1>Verify 2-Factor</h1>
+        <p className="ztdl-p">Enter the code from your authenticator app. </p>
+
+        <div>
+          {loginSettings?.secondFactors.map((factor) => {
+            return (
+              <div>
+                {factor === 1 && <div>TOTP</div>}
+                {factor === 2 && <div>U2F</div>}
+                {factor === 3 && <div>OTP Email</div>}
+                {factor === 4 && <div>OTP Sms</div>}
+              </div>
+            );
+          })}
+        </div>
+      </div>
+    </DynamicTheme>
+  );
+}
diff --git a/apps/login/app/api/session/route.ts b/apps/login/app/api/session/route.ts
index 36537d2a7b5..d63498e8389 100644
--- a/apps/login/app/api/session/route.ts
+++ b/apps/login/app/api/session/route.ts
@@ -1,4 +1,9 @@
-import { server, deleteSession } from "#/lib/zitadel";
+import {
+  server,
+  deleteSession,
+  getUserById,
+  listHumanAuthFactors,
+} from "#/lib/zitadel";
 import {
   SessionCookie,
   getMostRecentSessionCookie,
@@ -100,12 +105,23 @@ export async function PUT(request: NextRequest) {
           challenges,
           undefined,
           authRequestId
-        ).then((session) => {
-          console.log(session);
+        ).then(async (session) => {
+          // if password, check if user has MFA methods
+          let authFactors;
+          if (password && session.factors?.user?.id) {
+            const response = await listHumanAuthFactors(
+              server,
+              session.factors?.user?.id
+            );
+            if (response.result && response.result.length) {
+              authFactors = response.result;
+            }
+          }
           return NextResponse.json({
             sessionId: session.id,
             factors: session.factors,
             challenges: session.challenges,
+            authFactors,
           });
         });
       })
diff --git a/apps/login/lib/zitadel.ts b/apps/login/lib/zitadel.ts
index c9267bb6095..90cd991e615 100644
--- a/apps/login/lib/zitadel.ts
+++ b/apps/login/lib/zitadel.ts
@@ -1,4 +1,6 @@
 import {
+  LegalAndSupportSettings,
+  PasswordComplexitySettings,
   ZitadelServer,
   ZitadelServerOptions,
   user,
@@ -15,8 +17,6 @@ import {
   AddHumanUserResponse,
   BrandingSettings,
   ListSessionsResponse,
-  LegalAndSupportSettings,
-  PasswordComplexitySettings,
   GetSessionResponse,
   VerifyEmailResponse,
   Checks,
@@ -40,6 +40,7 @@ import {
   CreateCallbackResponse,
   RequestChallenges,
   TextQueryMethod,
+  ListHumanAuthFactorsResponse,
   AddHumanUserRequest,
 } from "@zitadel/server";
 
@@ -257,6 +258,14 @@ export async function addHumanUser(
   );
 }
 
+export async function listHumanAuthFactors(
+  server: ZitadelServer,
+  userId: string
+): Promise<ListHumanAuthFactorsResponse> {
+  const managementService = management.getManagement(server);
+  return managementService.listHumanAuthFactors({ userId }, {});
+}
+
 export async function listUsers(
   userName: string,
   organizationId: string
diff --git a/apps/login/ui/PasswordForm.tsx b/apps/login/ui/PasswordForm.tsx
index 7c53240f0b8..8f3a077a49f 100644
--- a/apps/login/ui/PasswordForm.tsx
+++ b/apps/login/ui/PasswordForm.tsx
@@ -89,7 +89,7 @@ export default function PasswordForm({
         let continueWithMfa = undefined;
         if (
           loginSettings?.forceMfa &&
-          loginSettings.secondFactors?.length >= 1 // TODO replace with user methods - if forceMFA is set and no user methods prompt to add method (/mfa/add)
+          resp.authFactors?.length >= 1 // TODO if forceMFA is set and no user methods prompt to add method (/mfa/add)
         ) {
           if (loginSettings.secondFactors?.length === 1) {
             continueWithMfa = loginSettings.secondFactors[0];
@@ -97,6 +97,23 @@ export default function PasswordForm({
             // continueWithMfa = loginSettings.secondFactors[0];
             // render selection page for mfa (/mfa/select)
           }
+        } else if (loginSettings?.forceMfa && resp.authFactors?.length === 0) {
+          const params = new URLSearchParams(
+            authRequestId
+              ? {
+                  loginName: resp.factors.user.loginName,
+                  authRequestId,
+                }
+              : {
+                  loginName: resp.factors.user.loginName,
+                }
+          );
+
+          if (organization) {
+            params.append("organization", organization);
+          }
+
+          return router.push(`/mfa/set?` + params);
         }
         // OIDC flows
         if (authRequestId && resp && resp.sessionId) {
diff --git a/packages/zitadel-server/src/index.ts b/packages/zitadel-server/src/index.ts
index 33cb28ab097..174a06187bf 100644
--- a/packages/zitadel-server/src/index.ts
+++ b/packages/zitadel-server/src/index.ts
@@ -68,6 +68,7 @@ export { TextQueryMethod } from "./proto/server/zitadel/object/v2beta/object";
 export {
   AddHumanUserResponse,
   AddHumanUserRequest,
+  GetUserByIDResponse,
   VerifyEmailResponse,
   VerifyPasskeyRegistrationRequest,
   VerifyPasskeyRegistrationResponse,
@@ -89,6 +90,7 @@ export {
   SetHumanPasswordResponse,
   SetHumanPasswordRequest,
   GetOrgByDomainGlobalResponse,
+  ListHumanAuthFactorsResponse,
 } from "./proto/server/zitadel/management";
 export * from "./proto/server/zitadel/idp";
 export { type LegalAndSupportSettings } from "./proto/server/zitadel/settings/v2beta/legal_settings";