From 71152b664be2713167dfdac8ae814578469db8ca Mon Sep 17 00:00:00 2001
From: Mohammed Imran <mohammedimran86992@gmail.com>
Date: Fri, 17 Oct 2025 23:08:54 +0530
Subject: [PATCH] feature: enhance 2FA management UI and logic in profile
 settings

- Replaced AlertDialog with Dialog for managing 2FA settings, improving user experience.
- Introduced multi-step dialog flow for verifying identity, managing actions, and displaying backup codes.
---
 .../settings/profile/disable-2fa.tsx          | 355 ++++++++++++++----
 .../settings/profile/profile-form.tsx         |  48 +--
 2 files changed, 309 insertions(+), 94 deletions(-)

diff --git a/apps/dokploy/components/dashboard/settings/profile/disable-2fa.tsx b/apps/dokploy/components/dashboard/settings/profile/disable-2fa.tsx
index 4055d4079..aea342b18 100644
--- a/apps/dokploy/components/dashboard/settings/profile/disable-2fa.tsx
+++ b/apps/dokploy/components/dashboard/settings/profile/disable-2fa.tsx
@@ -1,17 +1,28 @@
 import { zodResolver } from "@hookform/resolvers/zod";
-import { useState } from "react";
+import { KeyRound, RefreshCw, ShieldOff } from "lucide-react";
+import { useEffect, useState } from "react";
 import { useForm } from "react-hook-form";
 import { toast } from "sonner";
 import { z } from "zod";
 import {
 	AlertDialog,
+	AlertDialogAction,
+	AlertDialogCancel,
 	AlertDialogContent,
 	AlertDialogDescription,
+	AlertDialogFooter,
 	AlertDialogHeader,
 	AlertDialogTitle,
-	AlertDialogTrigger,
 } from "@/components/ui/alert-dialog";
 import { Button } from "@/components/ui/button";
+import {
+	Dialog,
+	DialogContent,
+	DialogDescription,
+	DialogHeader,
+	DialogTitle,
+	DialogTrigger,
+} from "@/components/ui/dialog";
 import {
 	Form,
 	FormControl,
@@ -32,11 +43,17 @@ const PasswordSchema = z.object({
 });
 
 type PasswordForm = z.infer<typeof PasswordSchema>;
+type Step = "password" | "actions" | "backup-codes";
 
 export const Disable2FA = () => {
 	const utils = api.useUtils();
-	const [isOpen, setIsOpen] = useState(false);
-	const [isLoading, setIsLoading] = useState(false);
+	const [isDialogOpen, setIsDialogOpen] = useState(false);
+	const [step, setStep] = useState<Step>("password");
+	const [password, setPassword] = useState("");
+	const [backupCodes, setBackupCodes] = useState<string[]>([]);
+	const [showDisableConfirm, setShowDisableConfirm] = useState(false);
+	const [isDisabling, setIsDisabling] = useState(false);
+	const [isRegenerating, setIsRegenerating] = useState(false);
 
 	const form = useForm<PasswordForm>({
 		resolver: zodResolver(PasswordSchema),
@@ -45,91 +62,287 @@ export const Disable2FA = () => {
 		},
 	});
 
-	const handleSubmit = async (formData: PasswordForm) => {
-		setIsLoading(true);
+	useEffect(() => {
+		if (!isDialogOpen) {
+			setStep("password");
+			setPassword("");
+			setBackupCodes([]);
+			form.reset();
+		}
+	}, [isDialogOpen, form]);
+
+	const handlePasswordSubmit = async (formData: PasswordForm) => {
+		setIsRegenerating(true);
 		try {
-			const result = await authClient.twoFactor.disable({
+			// Verify password by attempting to generate backup codes
+			// This validates the password and checks if 2FA is enabled
+			const result = await authClient.twoFactor.generateBackupCodes({
 				password: formData.password,
 			});
 
 			if (result.error) {
-				form.setError("password", {
-					message: result.error.message,
-				});
+				form.setError("password", { message: result.error.message });
+				toast.error(result.error.message);
+				return;
+			}
+
+			// If we get here, password is correct
+			setPassword(formData.password);
+			setStep("actions");
+		} catch (error) {
+			form.setError("password", {
+				message: error instanceof Error ? error.message : "Incorrect password",
+			});
+			toast.error("Incorrect password");
+		} finally {
+			setIsRegenerating(false);
+		}
+	};
+
+	const handleRegenerateBackupCodes = async () => {
+		setIsRegenerating(true);
+		try {
+			const result = await authClient.twoFactor.generateBackupCodes({
+				password,
+			});
+
+			if (result.error) {
+				toast.error(result.error.message);
+				return;
+			}
+
+			if (result.data?.backupCodes) {
+				setBackupCodes(result.data.backupCodes);
+				setStep("backup-codes");
+				toast.success("Backup codes regenerated successfully");
+			}
+		} catch (error) {
+			toast.error(
+				error instanceof Error
+					? error.message
+					: "Failed to regenerate backup codes",
+			);
+		} finally {
+			setIsRegenerating(false);
+		}
+	};
+
+	const handleDisable2FA = async () => {
+		setIsDisabling(true);
+		try {
+			const result = await authClient.twoFactor.disable({
+				password,
+			});
+
+			if (result.error) {
 				toast.error(result.error.message);
 				return;
 			}
 
 			toast.success("2FA disabled successfully");
 			utils.user.get.invalidate();
-			setIsOpen(false);
-		} catch {
-			form.setError("password", {
-				message: "Connection error. Please try again.",
-			});
-			toast.error("Connection error. Please try again.");
+			setIsDialogOpen(false);
+			setShowDisableConfirm(false);
+		} catch (error) {
+			toast.error("Failed to disable 2FA. Please try again.");
 		} finally {
-			setIsLoading(false);
+			setIsDisabling(false);
+		}
+	};
+
+	const handleCloseDialog = () => {
+		if (step === "backup-codes") {
+			setStep("actions");
+		} else {
+			setIsDialogOpen(false);
 		}
 	};
 
 	return (
-		<AlertDialog open={isOpen} onOpenChange={setIsOpen}>
-			<AlertDialogTrigger asChild>
-				<Button variant="destructive">Disable 2FA</Button>
-			</AlertDialogTrigger>
-			<AlertDialogContent>
-				<AlertDialogHeader>
-					<AlertDialogTitle>Are you absolutely sure?</AlertDialogTitle>
-					<AlertDialogDescription>
-						This action cannot be undone. This will permanently disable
-						Two-Factor Authentication for your account.
-					</AlertDialogDescription>
-				</AlertDialogHeader>
+		<>
+			<Dialog open={isDialogOpen} onOpenChange={setIsDialogOpen}>
+				<DialogTrigger asChild>
+					<Button variant="secondary">
+						<KeyRound className="size-4 text-muted-foreground" />
+						Manage 2FA
+					</Button>
+				</DialogTrigger>
+				<DialogContent className="sm:max-w-xl">
+					<DialogHeader>
+						<DialogTitle>
+							{step === "password" && "Verify Your Identity"}
+							{step === "actions" && "2FA Configuration"}
+							{step === "backup-codes" && "New Backup Codes"}
+						</DialogTitle>
+						<DialogDescription>
+							{step === "password" &&
+								"Enter your password to manage your 2FA settings"}
+							{step === "actions" &&
+								"Choose an action to manage your two-factor authentication"}
+							{step === "backup-codes" &&
+								"Save these backup codes in a secure place"}
+						</DialogDescription>
+					</DialogHeader>
 
-				<Form {...form}>
-					<form
-						onSubmit={form.handleSubmit(handleSubmit)}
-						className="space-y-4"
-					>
-						<FormField
-							control={form.control}
-							name="password"
-							render={({ field }) => (
-								<FormItem>
-									<FormLabel>Password</FormLabel>
-									<FormControl>
-										<Input
-											type="password"
-											placeholder="Enter your password"
-											{...field}
-										/>
-									</FormControl>
-									<FormDescription>
-										Enter your password to disable 2FA
-									</FormDescription>
-									<FormMessage />
-								</FormItem>
-							)}
-						/>
-						<div className="flex justify-end gap-4">
-							<Button
-								type="button"
-								variant="outline"
-								onClick={() => {
-									form.reset();
-									setIsOpen(false);
-								}}
+					{step === "password" && (
+						<Form {...form}>
+							<form
+								onSubmit={form.handleSubmit(handlePasswordSubmit)}
+								className="space-y-4"
 							>
-								Cancel
-							</Button>
-							<Button type="submit" variant="destructive" isLoading={isLoading}>
-								Disable 2FA
-							</Button>
+								<FormField
+									control={form.control}
+									name="password"
+									render={({ field }) => (
+										<FormItem>
+											<FormLabel>Password</FormLabel>
+											<FormControl>
+												<Input
+													type="password"
+													placeholder="Enter your password"
+													{...field}
+												/>
+											</FormControl>
+											<FormDescription>
+												Enter your password to continue
+											</FormDescription>
+											<FormMessage />
+										</FormItem>
+									)}
+								/>
+								<div className="flex justify-end gap-4">
+									<Button
+										type="button"
+										variant="outline"
+										onClick={() => setIsDialogOpen(false)}
+									>
+										Cancel
+									</Button>
+									<Button type="submit" isLoading={isRegenerating}>
+										Continue
+									</Button>
+								</div>
+							</form>
+						</Form>
+					)}
+
+					{step === "actions" && (
+						<div className="space-y-4">
+							<div className="grid gap-3">
+								<div className="flex flex-col gap-2 p-4 border rounded-lg hover:bg-muted/50 transition-colors">
+									<div className="flex items-start justify-between">
+										<div className="flex-1">
+											<h4 className="font-medium flex items-center gap-2">
+												<RefreshCw className="size-4" />
+												Regenerate Backup Codes
+											</h4>
+											<p className="text-sm text-muted-foreground mt-1">
+												Generate new backup codes to replace your existing ones.
+												This will invalidate all previous backup codes.
+											</p>
+										</div>
+									</div>
+									<Button
+										onClick={handleRegenerateBackupCodes}
+										variant="outline"
+										className="w-full mt-2"
+										isLoading={isRegenerating}
+									>
+										<RefreshCw className="size-4 mr-2" />
+										Regenerate Backup Codes
+									</Button>
+								</div>
+
+								<div className="flex flex-col gap-2 p-4 border border-destructive/50 rounded-lg hover:bg-destructive/5 transition-colors">
+									<div className="flex items-start justify-between">
+										<div className="flex-1">
+											<h4 className="font-medium flex items-center gap-2 text-destructive">
+												<ShieldOff className="size-4" />
+												Disable 2FA
+											</h4>
+											<p className="text-sm text-muted-foreground mt-1">
+												Completely disable two-factor authentication for your
+												account. This will make your account less secure.
+											</p>
+										</div>
+									</div>
+									<Button
+										onClick={() => setShowDisableConfirm(true)}
+										variant="destructive"
+										className="w-full mt-2"
+									>
+										<ShieldOff className="size-4 mr-2" />
+										Disable 2FA
+									</Button>
+								</div>
+							</div>
+
+							<div className="flex justify-end">
+								<Button
+									variant="outline"
+									onClick={() => setIsDialogOpen(false)}
+								>
+									Close
+								</Button>
+							</div>
 						</div>
-					</form>
-				</Form>
-			</AlertDialogContent>
-		</AlertDialog>
+					)}
+
+					{step === "backup-codes" && (
+						<div className="space-y-4">
+							<div className="w-full space-y-3 border rounded-lg p-4 bg-muted/50">
+								<div className="grid grid-cols-2 gap-2">
+									{backupCodes.map((code, index) => (
+										<code
+											key={index}
+											className="bg-background p-2 rounded text-sm font-mono text-center"
+										>
+											{code}
+										</code>
+									))}
+								</div>
+								<p className="text-sm text-muted-foreground">
+									Save these backup codes in a secure place. You can use them to
+									access your account if you lose access to your authenticator
+									device. Each code can only be used once.
+								</p>
+							</div>
+
+							<div className="flex justify-end gap-4">
+								<Button variant="outline" onClick={handleCloseDialog}>
+									Back to Actions
+								</Button>
+								<Button onClick={() => setIsDialogOpen(false)}>Done</Button>
+							</div>
+						</div>
+					)}
+				</DialogContent>
+			</Dialog>
+
+			<AlertDialog
+				open={showDisableConfirm}
+				onOpenChange={setShowDisableConfirm}
+			>
+				<AlertDialogContent>
+					<AlertDialogHeader>
+						<AlertDialogTitle>Are you absolutely sure?</AlertDialogTitle>
+						<AlertDialogDescription>
+							This will permanently disable Two-Factor Authentication for your
+							account. Your account will be less secure without 2FA enabled.
+						</AlertDialogDescription>
+					</AlertDialogHeader>
+					<AlertDialogFooter>
+						<AlertDialogCancel>Cancel</AlertDialogCancel>
+						<AlertDialogAction
+							onClick={handleDisable2FA}
+							className="bg-destructive text-destructive-foreground hover:bg-destructive/90"
+							disabled={isDisabling}
+						>
+							{isDisabling ? "Disabling..." : "Disable 2FA"}
+						</AlertDialogAction>
+					</AlertDialogFooter>
+				</AlertDialogContent>
+			</AlertDialog>
+		</>
 	);
 };
diff --git a/apps/dokploy/components/dashboard/settings/profile/profile-form.tsx b/apps/dokploy/components/dashboard/settings/profile/profile-form.tsx
index c481d5b8b..79f18225e 100644
--- a/apps/dokploy/components/dashboard/settings/profile/profile-form.tsx
+++ b/apps/dokploy/components/dashboard/settings/profile/profile-form.tsx
@@ -62,7 +62,6 @@ const randomImages = [
 ];
 
 export const ProfileForm = () => {
-	const _utils = api.useUtils();
 	const { data, refetch, isLoading } = api.user.get.useQuery();
 	const { data: isCloud } = api.settings.isCloud.useQuery();
 
@@ -120,28 +119,27 @@ export const ProfileForm = () => {
 	}, [form, data]);
 
 	const onSubmit = async (values: Profile) => {
-		await mutateAsync({
-			email: values.email.toLowerCase(),
-			password: values.password || undefined,
-			image: values.image,
-			currentPassword: values.currentPassword || undefined,
-			allowImpersonation: values.allowImpersonation,
-			name: values.name || undefined,
-		})
-			.then(async () => {
-				await refetch();
-				toast.success("Profile Updated");
-				form.reset({
-					email: values.email,
-					password: "",
-					image: values.image,
-					currentPassword: "",
-					name: values.name || "",
-				});
-			})
-			.catch(() => {
-				toast.error("Error updating the profile");
+		try {
+			await mutateAsync({
+				email: values.email.toLowerCase(),
+				password: values.password || undefined,
+				image: values.image,
+				currentPassword: values.currentPassword || undefined,
+				allowImpersonation: values.allowImpersonation,
+				name: values.name || undefined,
 			});
+			await refetch();
+			toast.success("Profile Updated");
+			form.reset({
+				email: values.email,
+				password: "",
+				image: values.image,
+				currentPassword: "",
+				name: values.name || "",
+			});
+		} catch (error) {
+			toast.error("Error updating the profile");
+		}
 	};
 
 	return (
@@ -158,7 +156,9 @@ export const ProfileForm = () => {
 								{t("settings.profile.description")}
 							</CardDescription>
 						</div>
-						{!data?.user.twoFactorEnabled ? <Enable2FA /> : <Disable2FA />}
+						<div className="flex flex-row gap-2">
+							{!data?.user.twoFactorEnabled ? <Enable2FA /> : <Disable2FA />}
+						</div>
 					</CardHeader>
 
 					<CardContent className="space-y-2 py-8 border-t">
@@ -304,6 +304,7 @@ export const ProfileForm = () => {
 																			}
 																		>
 																			{field.value?.startsWith("data:") ? (
+																				// biome-ignore lint/performance/noImgElement: this is an justified use of img element
 																				<img
 																					src={field.value}
 																					alt="Custom avatar"
@@ -362,6 +363,7 @@ export const ProfileForm = () => {
 																				/>
 																			</FormControl>
 
+																			{/* biome-ignore lint/performance/noImgElement: this is an justified use of img element */}
 																			<img
 																				key={image}
 																				src={image}