From 39c6cbf66a9cb29f711e1af3b96243c31fe143ae Mon Sep 17 00:00:00 2001
From: Anik Dhabal Babu <adhabal2002@gmail.com>
Date: Mon, 19 Feb 2024 11:25:15 +0000
Subject: [PATCH] feat: ability to download 2FA recovery codes

---
 .../2fa/enable-authenticator-app-dialog.tsx   | 24 +++++++++++++++----
 .../forms/2fa/view-recovery-codes-dialog.tsx  | 10 ++++----
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/apps/web/src/components/forms/2fa/enable-authenticator-app-dialog.tsx b/apps/web/src/components/forms/2fa/enable-authenticator-app-dialog.tsx
index 7a181c4cc..671292bde 100644
--- a/apps/web/src/components/forms/2fa/enable-authenticator-app-dialog.tsx
+++ b/apps/web/src/components/forms/2fa/enable-authenticator-app-dialog.tsx
@@ -1,4 +1,4 @@
-import { useMemo } from 'react';
+import { useMemo, useState } from 'react';
 
 import { useRouter } from 'next/navigation';
 
@@ -56,6 +56,7 @@ export const EnableAuthenticatorAppDialog = ({
 }: EnableAuthenticatorAppDialogProps) => {
   const router = useRouter();
   const { toast } = useToast();
+  const [recoveryCodesUrl, setRecoveryCodesUrl] = useState('');
 
   const { mutateAsync: setupTwoFactorAuthentication, data: setupTwoFactorAuthenticationData } =
     trpc.twoFactorAuthentication.setup.useMutation();
@@ -115,6 +116,16 @@ export const EnableAuthenticatorAppDialog = ({
     }
   };
 
+  const downloadRecoveryCodes = () => {
+    if (enableTwoFactorAuthenticationData && enableTwoFactorAuthenticationData.recoveryCodes) {
+      const textBlob = new Blob([enableTwoFactorAuthenticationData.recoveryCodes.join('\n')], {
+        type: 'text/plain',
+      });
+      if (recoveryCodesUrl) URL.revokeObjectURL(recoveryCodesUrl);
+      setRecoveryCodesUrl(URL.createObjectURL(textBlob));
+    }
+  };
+
   const onEnableTwoFactorAuthenticationFormSubmit = async ({
     token,
   }: TEnableTwoFactorAuthenticationForm) => {
@@ -270,10 +281,13 @@ export const EnableAuthenticatorAppDialog = ({
                 <RecoveryCodeList recoveryCodes={enableTwoFactorAuthenticationData.recoveryCodes} />
               )}
 
-              <div className="mt-4 flex w-full flex-row-reverse items-center justify-between">
-                <Button type="button" onClick={() => onCompleteClick()}>
-                  Complete
-                </Button>
+              <div className="mt-4 flex flex-row-reverse items-center gap-2">
+                <Button onClick={() => onOpenChange(false)}>Complete</Button>
+                <a download="documenso-2FA-recovery-codes.txt" href={recoveryCodesUrl}>
+                  <Button variant="secondary" onClick={downloadRecoveryCodes}>
+                    Download
+                  </Button>
+                </a>
               </div>
             </div>
           ))
diff --git a/apps/web/src/components/forms/2fa/view-recovery-codes-dialog.tsx b/apps/web/src/components/forms/2fa/view-recovery-codes-dialog.tsx
index cfdae7015..797b61b84 100644
--- a/apps/web/src/components/forms/2fa/view-recovery-codes-dialog.tsx
+++ b/apps/web/src/components/forms/2fa/view-recovery-codes-dialog.tsx
@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useState } from 'react';
+import { useMemo, useState } from 'react';
 
 import { zodResolver } from '@hookform/resolvers/zod';
 import { useForm } from 'react-hook-form';
@@ -63,7 +63,7 @@ export const ViewRecoveryCodesDialog = ({ open, onOpenChange }: ViewRecoveryCode
     return 'view';
   }, [viewRecoveryCodesData, isViewRecoveryCodesSubmitting]);
 
-  useEffect(() => {
+  const downloadRecoveryCodes = () => {
     if (viewRecoveryCodesData && viewRecoveryCodesData.recoveryCodes) {
       const textBlob = new Blob([viewRecoveryCodesData.recoveryCodes.join('\n')], {
         type: 'text/plain',
@@ -71,7 +71,7 @@ export const ViewRecoveryCodesDialog = ({ open, onOpenChange }: ViewRecoveryCode
       if (recoveryCodesUrl) URL.revokeObjectURL(recoveryCodesUrl);
       setRecoveryCodesUrl(URL.createObjectURL(textBlob));
     }
-  }, [viewRecoveryCodesData]);
+  };
 
   const onViewRecoveryCodesFormSubmit = async ({ password }: TViewRecoveryCodesForm) => {
     try {
@@ -153,7 +153,9 @@ export const ViewRecoveryCodesDialog = ({ open, onOpenChange }: ViewRecoveryCode
               <div className="mt-4 flex flex-row-reverse items-center gap-2">
                 <Button onClick={() => onOpenChange(false)}>Complete</Button>
                 <a download="documenso-2FA-recovery-codes.txt" href={recoveryCodesUrl}>
-                  <Button variant="secondary">Download</Button>
+                  <Button variant="secondary" onClick={downloadRecoveryCodes}>
+                    Download
+                  </Button>
                 </a>
               </div>
             </div>