From e575389a661db528c8234723882d166b8605372d Mon Sep 17 00:00:00 2001
From: Lea <lea@janderedev.xyz>
Date: Thu, 10 Aug 2023 16:57:00 +0200
Subject: [PATCH] feat: MFA management

---
 components/pages/inspector/AccountActions.tsx | 73 +++++++++++++++++++
 lib/accessPermissions.ts                      |  3 +-
 lib/actions.ts                                | 34 +++++++++
 lib/db.ts                                     | 40 +++++++---
 4 files changed, 140 insertions(+), 10 deletions(-)

diff --git a/components/pages/inspector/AccountActions.tsx b/components/pages/inspector/AccountActions.tsx
index 7d599d9..e893fbd 100644
--- a/components/pages/inspector/AccountActions.tsx
+++ b/components/pages/inspector/AccountActions.tsx
@@ -8,7 +8,9 @@ import { User } from "revolt-api";
 import {
   cancelAccountDeletion,
   changeAccountEmail,
+  deleteMFARecoveryCodes,
   disableAccount,
+  disableMFA,
   queueAccountDeletion,
   restoreAccount,
   verifyAccountEmail,
@@ -137,6 +139,77 @@ export function AccountActions({
           </AlertDialogFooter>
         </AlertDialogContent>
       </AlertDialog>
+      
+      <AlertDialog>
+        <AlertDialogTrigger asChild>
+          <Button
+            className="flex-1"
+            disabled={!accountDraft.mfa?.totp_token?.status}
+          >
+            MFA {accountDraft.mfa?.totp_token?.status.toLowerCase() || "disabled"}
+          </Button>
+        </AlertDialogTrigger>
+        <AlertDialogContent>
+          <AlertDialogHeader>
+            <AlertDialogTitle>
+              Manage MFA
+            </AlertDialogTitle>
+            <AlertDialogDescription>
+            MFA is currently {
+              accountDraft.mfa?.totp_token?.status == "Pending"
+                ? "pending setup"
+                : (accountDraft.mfa?.totp_token?.status.toLowerCase() || "disabled")
+            }.
+            <br />
+            The account has {accountDraft.mfa?.recovery_codes ?? "no"} recovery codes.
+            </AlertDialogDescription>
+          </AlertDialogHeader>
+          <AlertDialogFooter>
+            <AlertDialogAction
+              className="hover:bg-red-800"
+              disabled={accountDraft.mfa?.recovery_codes == null}
+              onClick={async () => {
+                try {
+                  await deleteMFARecoveryCodes(account._id);
+                  toast({
+                    title: "MFA recovery codes cleared",
+                  });
+                  accountDraft.mfa!.recovery_codes = undefined;
+                } catch(e) {
+                  toast({
+                    title: "Failed to clear recovery codes",
+                    description: String(e),
+                    variant: "destructive",
+                  })
+                }
+              }}
+            >
+              Clear recovery codes
+            </AlertDialogAction>
+            <AlertDialogAction
+              className="hover:bg-red-800"
+              onClick={async () => {
+                try {
+                  await disableMFA(account._id);
+                  toast({
+                    title: "MFA disabled",
+                  });
+                  accountDraft.mfa!.totp_token = undefined;
+                } catch(e) {
+                  toast({
+                    title: " Failed to disable MFA",
+                    description: String(e),
+                    variant: "destructive",
+                  })
+                }
+              }}
+            >
+              Disable MFA
+            </AlertDialogAction>
+            <AlertDialogCancel>Close</AlertDialogCancel>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialog>
 
       <AlertDialog>
         <AlertDialogTrigger asChild>
diff --git a/lib/accessPermissions.ts b/lib/accessPermissions.ts
index a150873..b0909d1 100644
--- a/lib/accessPermissions.ts
+++ b/lib/accessPermissions.ts
@@ -8,7 +8,7 @@ type Permission =
   | `accounts${
       | ""
       | `/fetch${"" | "/by-id" | "/by-email"}`
-      | "/update/email"
+      | `/update${"" | "/email" | "/mfa"}`
       | "/disable"
       | "/restore"
       | `/deletion${"" | "/queue" | "/cancel"}`}`
@@ -107,6 +107,7 @@ const PermissionSets = {
     "accounts/deletion/queue",
     "accounts/deletion/cancel",
     "accounts/update/email",
+    "accounts/update/mfa",
   ] as Permission[],
 
   // Moderate users
diff --git a/lib/actions.ts b/lib/actions.ts
index 00ebbc5..f55318e 100644
--- a/lib/actions.ts
+++ b/lib/actions.ts
@@ -195,6 +195,40 @@ export async function disableAccount(userId: string) {
   });
 }
 
+export async function deleteMFARecoveryCodes(userId: string) {
+  if (RESTRICT_ACCESS_LIST.includes(userId)) throw "restricted access";
+  await checkPermission("accounts/update/mfa", userId);
+
+  await mongo()
+    .db("revolt")
+    .collection<Account>("accounts")
+    .updateOne(
+      { _id: userId },
+      {
+        $unset: {
+          "mfa.recovery_codes": 1,
+        },
+      },
+    )
+}
+
+export async function disableMFA(userId: string) {
+  if (RESTRICT_ACCESS_LIST.includes(userId)) throw "restricted access";
+  await checkPermission("accounts/update/mfa", userId);
+
+  await mongo()
+    .db("revolt")
+    .collection<Account>("accounts")
+    .updateOne(
+      { _id: userId },
+      {
+        $unset: {
+          "mfa.totp_token": 1,
+        },
+      },
+    )
+}
+
 export async function changeAccountEmail(userId: string, email: string) {
   if (RESTRICT_ACCESS_LIST.includes(userId)) throw "restricted access";
   await checkPermission("accounts/update/email", userId);
diff --git a/lib/db.ts b/lib/db.ts
index d687c06..14d810c 100644
--- a/lib/db.ts
+++ b/lib/db.ts
@@ -1,6 +1,6 @@
 "use server";
 
-import { Filter, MongoClient } from "mongodb";
+import { Filter, MongoClient, WithId } from "mongodb";
 import type {
   AccountStrike,
   Bot,
@@ -104,6 +104,12 @@ export type Account = {
     attempts: number;
     expiry: string;
   };
+  mfa?: {
+    totp_token?: {
+      status: "Pending" | "Enabled";
+    };
+    recovery_codes?: number;
+  };
 };
 
 export async function fetchAccountById(id: string) {
@@ -112,15 +118,31 @@ export async function fetchAccountById(id: string) {
   return await mongo()
     .db("revolt")
     .collection<Account>("accounts")
-    .findOne(
-      { _id: id },
-      {
-        projection: {
-          password: 0,
-          mfa: 0,
+    .aggregate(
+      [
+        {
+          $match: { _id: id },
         },
-      }
-    );
+        {
+          $project: {
+            password: 0,
+            "mfa.totp_token.secret": 0,
+          }
+        },
+        {
+          $set: {
+            // Replace recovery code array with amount of codes
+            "mfa.recovery_codes": {
+              $cond: {
+                if: { $isArray: "$mfa.recovery_codes" },
+                then: { $size: "$mfa.recovery_codes", },
+                else: undefined,
+              }
+            }
+          }
+        }
+      ]
+    ).next() as WithId<Account>;
 }
 
 export async function fetchSessionsByAccount(accountId: string) {