From a71502d0b2ab42620c722e5415a759244a5cc98e Mon Sep 17 00:00:00 2001
From: Rainnny7 <braydonrainnny@gmail.com>
Date: Thu, 19 Sep 2024 23:34:11 -0400
Subject: [PATCH] Require your TFA pin to disable TFA

---
 .../user/settings/tfa/tfa-setting.tsx         |  51 +++--
 .../user/settings/tfa/tfa-setup-form.tsx      |   3 +
 src/components/ui/dialog.tsx                  |   5 +-
 src/components/user/user-tfa-prompt.tsx       | 191 ++++++++++++++++++
 4 files changed, 228 insertions(+), 22 deletions(-)
 create mode 100644 src/components/user/user-tfa-prompt.tsx

diff --git a/src/components/dashboard/user/settings/tfa/tfa-setting.tsx b/src/components/dashboard/user/settings/tfa/tfa-setting.tsx
index 719ed7b..9e8ff3c 100644
--- a/src/components/dashboard/user/settings/tfa/tfa-setting.tsx
+++ b/src/components/dashboard/user/settings/tfa/tfa-setting.tsx
@@ -21,9 +21,8 @@ import {
 } from "@/components/ui/dialog";
 import { Button } from "@/components/ui/button";
 import TfaSetupForm from "@/components/dashboard/user/settings/tfa/tfa-setup-form";
-import { AppRouterInstance } from "next/dist/shared/lib/app-router-context.shared-runtime";
-import { useRouter } from "next/navigation";
 import { toast } from "sonner";
+import UserTfaPrompt from "@/components/user/user-tfa-prompt";
 
 /**
  * The setting that allows a
@@ -39,13 +38,11 @@ const TFASetting = (): ReactElement => {
     const user: User | undefined = useUserContext(
         (state: UserState) => state.user
     );
-    const router: AppRouterInstance = useRouter();
 
     const [tfaResponse, setTfaResponse] = useState<
         UserSetupTfaResponse | undefined
     >(undefined);
     const [enabledTfa, setEnabledTfa] = useState<boolean>(false);
-    const [disabling, setDisabling] = useState<boolean>(false);
 
     const onDialogStateChange = async (open: boolean) => {
         if (open) {
@@ -77,20 +74,30 @@ const TFASetting = (): ReactElement => {
     /**
      * Disable two-factor auth.
      */
-    const disableTfa = async () => {
-        setDisabling(true);
-        await apiRequest<void>({
+    const disableTfa = async ({
+        pin,
+        setError,
+    }: {
+        pin: string;
+        setError: (error: string | undefined) => void;
+    }) => {
+        const { error } = await apiRequest<void>({
             endpoint: "/user/disable-tfa",
             method: "POST",
             session,
+            body: { pin },
         });
-        toast("Two-Factor Auth", {
-            icon: "🔓",
-            description: "Two-factor auth has been disabled for your account.",
-        });
-        setTimeout(() => {
-            window.location.reload();
-        }, 1500);
+        setError(error?.message);
+        if (!error) {
+            toast("Two-Factor Auth", {
+                icon: "🔓",
+                description:
+                    "Two-factor auth has been disabled for your account.",
+            });
+            setTimeout(() => {
+                window.location.reload();
+            }, 1500);
+        }
     };
 
     return (
@@ -107,14 +114,15 @@ const TFASetting = (): ReactElement => {
 
             {/* Setting */}
             {hasFlag(user as User, UserFlag.TFA_ENABLED) ? (
-                <Button
-                    size="sm"
-                    variant="destructive"
-                    onClick={disableTfa}
-                    disabled={disabling}
+                <UserTfaPrompt
+                    message="Please verify it's you before disabling two-factor auth."
+                    submitButtonText="Disable TFA"
+                    onSubmit={disableTfa}
                 >
-                    Disable
-                </Button>
+                    <Button size="sm" variant="destructive">
+                        Disable
+                    </Button>
+                </UserTfaPrompt>
             ) : (
                 <Dialog onOpenChange={onDialogStateChange}>
                     <DialogTrigger>
@@ -122,6 +130,7 @@ const TFASetting = (): ReactElement => {
                             className="bg-background/30"
                             size="sm"
                             variant="outline"
+                            disabled={enabledTfa}
                         >
                             Setup
                         </Button>
diff --git a/src/components/dashboard/user/settings/tfa/tfa-setup-form.tsx b/src/components/dashboard/user/settings/tfa/tfa-setup-form.tsx
index b446e56..c72eee6 100644
--- a/src/components/dashboard/user/settings/tfa/tfa-setup-form.tsx
+++ b/src/components/dashboard/user/settings/tfa/tfa-setup-form.tsx
@@ -1,3 +1,5 @@
+"use client";
+
 import { ReactElement, useState } from "react";
 import { UserSetupTfaResponse } from "@/app/types/user/response/user-setup-tfa-response";
 import { useForm } from "react-hook-form";
@@ -25,6 +27,7 @@ const FormSchema = z.object({
  * the setup of TFA for a user.
  *
  * @param tfaResponse the tfa setup response
+ * @param setEnabledTfa the function to invoke to indicate tfa is enabled
  * @return the form jsx
  */
 const TfaSetupForm = ({
diff --git a/src/components/ui/dialog.tsx b/src/components/ui/dialog.tsx
index a33c04d..1415cc1 100644
--- a/src/components/ui/dialog.tsx
+++ b/src/components/ui/dialog.tsx
@@ -44,7 +44,10 @@ const DialogContent = React.forwardRef<
             {...props}
         >
             {children}
-            <DialogPrimitive.Close className="absolute right-4 top-4 rounded-sm opacity-70 ring-offset-background transition-opacity hover:opacity-100 focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2 disabled:pointer-events-none data-[state=open]:bg-accent data-[state=open]:text-muted-foreground">
+            <DialogPrimitive.Close
+                id="closeDialog"
+                className="absolute right-4 top-4 rounded-sm opacity-70 ring-offset-background transition-opacity hover:opacity-100 focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2 disabled:pointer-events-none data-[state=open]:bg-accent data-[state=open]:text-muted-foreground"
+            >
                 <Cross2Icon className="h-4 w-4" />
                 <span className="sr-only">Close</span>
             </DialogPrimitive.Close>
diff --git a/src/components/user/user-tfa-prompt.tsx b/src/components/user/user-tfa-prompt.tsx
new file mode 100644
index 0000000..04a4668
--- /dev/null
+++ b/src/components/user/user-tfa-prompt.tsx
@@ -0,0 +1,191 @@
+"use client";
+
+import { ReactElement, ReactNode, useState } from "react";
+import {
+    Dialog,
+    DialogContent,
+    DialogDescription,
+    DialogHeader,
+    DialogTitle,
+    DialogTrigger,
+} from "@/components/ui/dialog";
+import { useForm } from "react-hook-form";
+import { zodResolver } from "@hookform/resolvers/zod";
+import { z } from "zod";
+import {
+    InputOTP,
+    InputOTPGroup,
+    InputOTPSeparator,
+    InputOTPSlot,
+} from "@/components/ui/input-otp";
+import { Button } from "@/components/ui/button";
+import { ArrowPathIcon } from "@heroicons/react/24/outline";
+
+/**
+ * The props for the TFA prompt.
+ */
+type UserTfaPromptProps = {
+    /**
+     * The message to display under the title.
+     */
+    message: string;
+
+    /**
+     * The text to display for the submit button.
+     */
+    submitButtonText: string;
+
+    /**
+     * Invoked when the form is submitted.
+     *
+     * @param pin the pin to submit
+     * @param setError the error setter
+     * @param closePrompt the prompt closer
+     */
+    onSubmit: ({
+        pin,
+        setError,
+        closePrompt,
+    }: {
+        pin: string;
+        setError: (error: string | undefined) => void;
+        closePrompt: () => void;
+    }) => Promise<any>;
+
+    /**
+     * The elements to trigger the prompt.
+     */
+    children: ReactNode;
+};
+
+const FormSchema = z.object({
+    pin: z.string(),
+});
+
+/**
+ * A universal prompt for a user's TFA pin.
+ *
+ * @param message the message to display
+ * @param submitButtonText the submit text
+ * @param onSubmit the function to invoke when the form is submitted
+ * @param children the elements to trigger the prompt.
+ * @return the prompt jsx
+ */
+const UserTfaPrompt = ({
+    message,
+    submitButtonText,
+    onSubmit,
+    children,
+}: UserTfaPromptProps): ReactElement => {
+    const [pin, setPin] = useState<string>("");
+    const [loading, setLoading] = useState<boolean>(false);
+    const [error, setError] = useState<string | undefined>();
+
+    // Build the form
+    const {
+        handleSubmit,
+        formState: { errors },
+    } = useForm({
+        resolver: zodResolver(FormSchema),
+    });
+
+    /**
+     * Handle submitting the form.
+     *
+     * @param pin the submitted pin
+     */
+    const handleFormSubmit = async ({ pin }: any) => {
+        if (pin.length === 6) {
+            setLoading(true);
+            await onSubmit({
+                pin,
+                setError: (error: string | undefined) => {
+                    setError(error);
+                    if (error) {
+                        setLoading(false);
+                    }
+                },
+                closePrompt: () =>
+                    document.getElementById("closeDialog")?.click(),
+            });
+        }
+    };
+
+    /**
+     * Cleanup the prompt when closed.
+     */
+    const cleanup = () => {
+        setPin("");
+        setLoading(false);
+        setError(undefined);
+    };
+
+    // Render the prompt
+    return (
+        <Dialog
+            onOpenChange={(open: boolean) => {
+                if (!open) {
+                    cleanup();
+                }
+            }}
+        >
+            <DialogTrigger>{children}</DialogTrigger>
+            <DialogContent>
+                {/* Header */}
+                <DialogHeader>
+                    <DialogTitle>Two-Factor Border</DialogTitle>
+                    <DialogDescription>{message}</DialogDescription>
+                </DialogHeader>
+
+                {/* Content */}
+                <form
+                    className="flex flex-col gap-4 items-center"
+                    onSubmit={async (event) => {
+                        event.preventDefault();
+                        await handleFormSubmit({ pin });
+                    }}
+                >
+                    {/* Input */}
+                    <InputOTP
+                        maxLength={6}
+                        value={pin}
+                        onChange={async (pin: string) => {
+                            await handleFormSubmit({ pin });
+                            setPin(pin);
+                        }}
+                    >
+                        <InputOTPGroup>
+                            {[0, 1, 2].map((index: number) => (
+                                <InputOTPSlot key={index} index={index} />
+                            ))}
+                        </InputOTPGroup>
+                        <InputOTPSeparator />
+                        <InputOTPGroup>
+                            {[3, 4, 5].map((index: number) => (
+                                <InputOTPSlot key={index} index={index} />
+                            ))}
+                        </InputOTPGroup>
+                    </InputOTP>
+
+                    {/* Display the error if it exists */}
+                    {error && (
+                        <p className="-mt-2 pb-0.5 text-red-500">{error}</p>
+                    )}
+
+                    {/* Verify Pin */}
+                    <Button
+                        className="w-32 flex gap-2.5 items-center text-white"
+                        type="submit"
+                        disabled={loading}
+                    >
+                        {loading && (
+                            <ArrowPathIcon className="w-4 h-4 animate-spin" />
+                        )}
+                        {submitButtonText}
+                    </Button>
+                </form>
+            </DialogContent>
+        </Dialog>
+    );
+};
+export default UserTfaPrompt;