feat: TOTP (two-factor) Authentication (#55)

* Working on some initial prototype stuff for TOTP

* Fixed a bug that prevented the change password menu from working

* Enable/disable totp working

* Added the new login procedure including TOTP! :)

* misc: Changed bad description for the TOTP_SECRET env var

* I forgot to include the migration for the new prisma stuff

* fix: refresh user context instead refreshing the page

* refactor: simplify totp error handling

* Removed U2F tab + format schema

* fix: tokens not saved in cookies

* refactor: deleted commented out code

* refactor: move password text to input description

* refactor: remove tabler icon package

Co-authored-by: Elias Schneider <login@eliasschneider.com>
Co-authored-by: Elias Schneider <58886915+stonith404@users.noreply.github.com>

frontend/src/components/account/showEnableTotpModal.tsx

@@ -0,0 +1,131 @@
+import {
+  Button,
+  Center,
+  Col,
+  Grid,
+  Image,
+  Stack,
+  Text,
+  TextInput,
+  Title,
+  Tooltip,
+} from "@mantine/core";
+import { useForm, yupResolver } from "@mantine/form";
+import { useModals } from "@mantine/modals";
+import { ModalsContextProps } from "@mantine/modals/lib/context";
+import * as yup from "yup";
+import useUser from "../../hooks/user.hook";
+import authService from "../../services/auth.service";
+import toast from "../../utils/toast.util";
+
+const showEnableTotpModal = (
+  modals: ModalsContextProps,
+  refreshUser: () => {},
+  options: {
+    qrCode: string;
+    secret: string;
+    password: string;
+  }
+) => {
+  return modals.openModal({
+    title: <Title order={4}>Enable TOTP</Title>,
+    children: (
+      <CreateEnableTotpModal options={options} refreshUser={refreshUser} />
+    ),
+  });
+};
+
+const CreateEnableTotpModal = ({
+  options,
+  refreshUser,
+}: {
+  options: {
+    qrCode: string;
+    secret: string;
+    password: string;
+  };
+  refreshUser: () => {};
+}) => {
+  const modals = useModals();
+  const user = useUser();
+
+  console.log(user.user);
+
+  const validationSchema = yup.object().shape({
+    code: yup
+      .string()
+      .min(6)
+      .max(6)
+      .required()
+      .matches(/^[0-9]+$/, { message: "Code must be a number" }),
+  });
+
+  const form = useForm({
+    initialValues: {
+      code: "",
+    },
+    validate: yupResolver(validationSchema),
+  });
+
+  return (
+    <div>
+      <Center>
+        <Stack>
+          <Text>Step 1: Add your authenticator</Text>
+          <Image src={options.qrCode} alt="QR Code" />
+
+          <Center>
+            <span>OR</span>
+          </Center>
+
+          <Tooltip label="Click to copy">
+            <Button
+              onClick={() => {
+                navigator.clipboard.writeText(options.secret);
+                toast.success("Copied to clipboard");
+              }}
+            >
+              {options.secret}
+            </Button>
+          </Tooltip>
+          <Center>
+            <Text fz="xs">Enter manually</Text>
+          </Center>
+
+          <Text>Step 2: Validate your code</Text>
+
+          <form
+            onSubmit={form.onSubmit((values) => {
+              authService
+                .verifyTOTP(values.code, options.password)
+                .then(() => {
+                  toast.success("Successfully enabled TOTP");
+                  modals.closeAll();
+                  refreshUser();
+                })
+                .catch(toast.axiosError);
+            })}
+          >
+            <Grid align="flex-end">
+              <Col xs={9}>
+                <TextInput
+                  variant="filled"
+                  label="Code"
+                  placeholder="******"
+                  {...form.getInputProps("code")}
+                />
+              </Col>
+              <Col xs={3}>
+                <Button variant="outline" type="submit">
+                  Verify
+                </Button>
+              </Col>
+            </Grid>
+          </form>
+        </Stack>
+      </Center>
+    </div>
+  );
+};
+
+export default showEnableTotpModal;

frontend/src/components/auth/SignInForm.tsx

@@ -9,7 +9,11 @@ import {
   Title,
 } from "@mantine/core";
 import { useForm, yupResolver } from "@mantine/form";
+import { showNotification } from "@mantine/notifications";
+import { setCookie } from "cookies-next";
 import Link from "next/link";
+import React from "react";
+import { TbInfoCircle } from "react-icons/tb";
 import * as yup from "yup";
 import useConfig from "../../hooks/config.hook";
 import authService from "../../services/auth.service";
@@ -17,16 +21,24 @@ import toast from "../../utils/toast.util";
 
 const SignInForm = () => {
   const config = useConfig();
+  const [showTotp, setShowTotp] = React.useState(false);
+  const [loginToken, setLoginToken] = React.useState("");
 
   const validationSchema = yup.object().shape({
     emailOrUsername: yup.string().required(),
     password: yup.string().min(8).required(),
+    totp: yup.string().when("totpRequired", {
+      is: true,
+      then: yup.string().min(6).max(6).required(),
+      otherwise: yup.string(),
+    }),
   });
 
   const form = useForm({
     initialValues: {
       emailOrUsername: "",
       password: "",
+      totp: "",
     },
     validate: yupResolver(validationSchema),
   });
@@ -34,10 +46,47 @@ const SignInForm = () => {
   const signIn = (email: string, password: string) => {
     authService
       .signIn(email, password)
-      .then(() => window.location.replace("/"))
+      .then((response) => {
+        if (response.data["loginToken"]) {
+          // Prompt the user to enter their totp code
+          setShowTotp(true);
+          showNotification({
+            icon: <TbInfoCircle />,
+            color: "blue",
+            radius: "md",
+            title: "Two-factor authentication required",
+            message: "Please enter your two-factor authentication code",
+          });
+          setLoginToken(response.data["loginToken"]);
+        } else {
+          setCookie("access_token", response.data.accessToken);
+          setCookie("refresh_token", response.data.refreshToken);
+          window.location.replace("/");
+        }
+      })
       .catch(toast.axiosError);
   };
 
+  const signInTotp = (email: string, password: string, totp: string) => {
+    authService
+      .signInTotp(email, password, totp, loginToken)
+      .then((response) => {
+        setCookie("access_token", response.data.accessToken);
+        setCookie("refresh_token", response.data.refreshToken);
+        window.location.replace("/");
+      })
+      .catch((error) => {
+        if (error?.response?.data?.message == "Login token expired") {
+          toast.error("Login token expired");
+          // Refresh the page to start over
+          window.location.reload();
+        }
+
+        toast.axiosError(error);
+        form.setValues({ totp: "" });
+      });
+  };
+
   return (
     <Container size={420} my={40}>
       <Title
@@ -59,9 +108,11 @@ const SignInForm = () => {
       )}
       <Paper withBorder shadow="md" p={30} mt={30} radius="md">
         <form
-          onSubmit={form.onSubmit((values) =>
-            signIn(values.emailOrUsername, values.password)
-          )}
+          onSubmit={form.onSubmit((values) => {
+            if (showTotp)
+              signInTotp(values.emailOrUsername, values.password, values.totp);
+            else signIn(values.emailOrUsername, values.password);
+          })}
         >
           <TextInput
             label="Email or username"
@@ -74,6 +125,15 @@ const SignInForm = () => {
             mt="md"
             {...form.getInputProps("password")}
           />
+          {showTotp && (
+            <TextInput
+              variant="filled"
+              label="Code"
+              placeholder="******"
+              mt="md"
+              {...form.getInputProps("totp")}
+            />
+          )}
           <Button fullWidth mt="xl" type="submit">
             Sign in
           </Button>

frontend/src/components/auth/SignUpForm.tsx

@@ -9,6 +9,7 @@ import {
   Title,
 } from "@mantine/core";
 import { useForm, yupResolver } from "@mantine/form";
+import { setCookie } from "cookies-next";
 import Link from "next/link";
 import * as yup from "yup";
 import useConfig from "../../hooks/config.hook";
@@ -33,16 +34,14 @@ const SignUpForm = () => {
     validate: yupResolver(validationSchema),
   });
 
-  const signIn = (email: string, password: string) => {
-    authService
-      .signIn(email, password)
-      .then(() => window.location.replace("/"))
-      .catch(toast.axiosError);
-  };
   const signUp = (email: string, username: string, password: string) => {
     authService
       .signUp(email, username, password)
-      .then(() => signIn(email, password))
+      .then((response) => {
+        setCookie("access_token", response.data.accessToken);
+        setCookie("refresh_token", response.data.refreshToken);
+        window.location.replace("/");
+      })
       .catch(toast.axiosError);
   };
 

frontend/src/components/navBar/ActionAvatar.tsx

@@ -5,7 +5,7 @@ import useUser from "../../hooks/user.hook";
 import authService from "../../services/auth.service";
 
 const ActionAvatar = () => {
-  const user = useUser();
+  const { user } = useUser();
 
   return (
     <Menu position="bottom-start" withinPortal>

frontend/src/components/navBar/NavBar.tsx

@@ -107,7 +107,7 @@ const useStyles = createStyles((theme) => ({
 }));
 
 const NavBar = () => {
-  const user = useUser();
+  const { user } = useUser();
   const config = useConfig();
 
   const [opened, toggleOpened] = useDisclosure(false);

frontend/src/hooks/user.hook.ts

@@ -1,7 +1,10 @@
 import { createContext, useContext } from "react";
-import { CurrentUser } from "../types/user.type";
+import { UserHook } from "../types/user.type";
 
-export const UserContext = createContext<CurrentUser | null>(null);
+export const UserContext = createContext<UserHook>({
+  user: null,
+  setUser: () => {},
+});
 
 const useUser = () => {
   return useContext(UserContext);

frontend/src/pages/_app.tsx

@@ -73,7 +73,7 @@ function App({ Component, pageProps }: AppProps) {
               <LoadingOverlay visible overlayOpacity={1} />
             ) : (
               <ConfigContext.Provider value={configVariables}>
-                <UserContext.Provider value={user} >
+                <UserContext.Provider value={{ user, setUser }}>
                   <LoadingOverlay visible={isLoading} overlayOpacity={1} />
                   <Header />
                   <Container>

frontend/src/pages/account/index.tsx

@@ -6,6 +6,7 @@ import {
   Paper,
   PasswordInput,
   Stack,
+  Tabs,
   Text,
   TextInput,
   Title,
@@ -13,14 +14,16 @@ import {
 import { useForm, yupResolver } from "@mantine/form";
 import { useModals } from "@mantine/modals";
 import { useRouter } from "next/router";
+import { Tb2Fa } from "react-icons/tb";
 import * as yup from "yup";
+import showEnableTotpModal from "../../components/account/showEnableTotpModal";
 import useUser from "../../hooks/user.hook";
 import authService from "../../services/auth.service";
 import userService from "../../services/user.service";
 import toast from "../../utils/toast.util";
 
 const Account = () => {
-  const user = useUser();
+  const { user, setUser } = useUser();
   const modals = useModals();
   const router = useRouter();
 
@@ -50,6 +53,36 @@ const Account = () => {
     ),
   });
 
+  const enableTotpForm = useForm({
+    initialValues: {
+      password: "",
+    },
+    validate: yupResolver(
+      yup.object().shape({
+        password: yup.string().min(8),
+      })
+    ),
+  });
+
+  const disableTotpForm = useForm({
+    initialValues: {
+      password: "",
+      code: "",
+    },
+    validate: yupResolver(
+      yup.object().shape({
+        password: yup.string().min(8),
+        code: yup
+          .string()
+          .min(6)
+          .max(6)
+          .matches(/^[0-9]+$/, { message: "Code must be a number" }),
+      })
+    ),
+  });
+
+  const refreshUser = async () => setUser(await userService.getCurrentUser());
+
   if (!user) {
     router.push("/");
     return;
@@ -117,31 +150,120 @@ const Account = () => {
           </Stack>
         </form>
       </Paper>
-      <Center mt={80}>
-        <Button
-          variant="light"
-          color="red"
-          onClick={() =>
-            modals.openConfirmModal({
-              title: "Account deletion",
-              children: (
-                <Text size="sm">
-                  Do you really want to delete your account including all your
-                  active shares?
-                </Text>
-              ),
 
-              labels: { confirm: "Delete", cancel: "Cancel" },
-              confirmProps: { color: "red" },
-              onConfirm: async () => {
-                await userService.removeCurrentUser();
-                window.location.reload();
-              },
-            })
-          }
-        >
-          Delete Account
-        </Button>
+      <Paper withBorder p="xl" mt="lg">
+        <Title order={5} mb="xs">
+          Security
+        </Title>
+
+        <Tabs defaultValue="totp">
+          <Tabs.List>
+            <Tabs.Tab value="totp" icon={<Tb2Fa size={14} />}>
+              TOTP
+            </Tabs.Tab>
+          </Tabs.List>
+
+          <Tabs.Panel value="totp" pt="xs">
+            {/* TODO: This is ugly, make it prettier */}
+            {/* If we have totp enabled, show different text */}
+            {user.totpVerified ? (
+              <>
+                <form
+                  onSubmit={disableTotpForm.onSubmit((values) => {
+                    authService
+                      .disableTOTP(values.code, values.password)
+                      .then(() => {
+                        toast.success("Successfully disabled TOTP");
+                        values.password = "";
+                        values.code = "";
+                        refreshUser();
+                      })
+                      .catch(toast.axiosError);
+                  })}
+                >
+                  <Stack>
+                    <PasswordInput
+                      description="Enter your current password to disable TOTP"
+                      label="Password"
+                      {...disableTotpForm.getInputProps("password")}
+                    />
+
+                    <TextInput
+                      variant="filled"
+                      label="Code"
+                      placeholder="******"
+                      {...disableTotpForm.getInputProps("code")}
+                    />
+
+                    <Group position="right">
+                      <Button color="red" type="submit">
+                        Disable
+                      </Button>
+                    </Group>
+                  </Stack>
+                </form>
+              </>
+            ) : (
+              <>
+                <form
+                  onSubmit={enableTotpForm.onSubmit((values) => {
+                    authService
+                      .enableTOTP(values.password)
+                      .then((result) => {
+                        showEnableTotpModal(modals, refreshUser, {
+                          qrCode: result.qrCode,
+                          secret: result.totpSecret,
+                          password: values.password,
+                        });
+                        values.password = "";
+                      })
+                      .catch(toast.axiosError);
+                  })}
+                >
+                  <Stack>
+                    <PasswordInput
+                      label="Password"
+                      description="Enter your current password to start enabling TOTP"
+                      {...enableTotpForm.getInputProps("password")}
+                    />
+                    <Group position="right">
+                      <Button type="submit">Start</Button>
+                    </Group>
+                  </Stack>
+                </form>
+              </>
+            )}
+          </Tabs.Panel>
+        </Tabs>
+      </Paper>
+
+      <Center mt={80}>
+        <Stack>
+          <Button
+            variant="light"
+            color="red"
+            onClick={() =>
+              modals.openConfirmModal({
+                title: "Account deletion",
+                children: (
+                  <Text size="sm">
+                    Do you really want to delete your account including all your
+                    active shares?
+                  </Text>
+                ),
+
+                labels: { confirm: "Delete", cancel: "Cancel" },
+                confirmProps: { color: "red" },
+                onConfirm: async () => {
+                  await userService.removeCurrentUser();
+                  window.location.reload();
+                },
+              })
+            }
+          >
+            Delete Account
+          </Button>
+        </Stack>
       </Center>
     </Container>
   );

frontend/src/pages/account/shares.tsx

@@ -28,7 +28,7 @@ const MyShares = () => {
   const modals = useModals();
   const clipboard = useClipboard();
   const router = useRouter();
-  const user = useUser();
+  const { user } = useUser();
 
   const [shares, setShares] = useState<MyShare[]>();
 

frontend/src/pages/admin/setup.tsx

@@ -10,7 +10,7 @@ import configService from "../../services/config.service";
 const Setup = () => {
   const router = useRouter();
   const config = useConfig();
-  const user = useUser();
+  const { user } = useUser();
 
   const [isLoading, setIsLoading] = useState(false);
 

frontend/src/pages/auth/signIn.tsx

@@ -4,7 +4,7 @@ import Meta from "../../components/Meta";
 import useUser from "../../hooks/user.hook";
 
 const SignIn = () => {
-  const user = useUser();
+  const { user } = useUser();
   const router = useRouter();
   if (user) {
     router.replace("/");

frontend/src/pages/auth/signUp.tsx

@@ -6,7 +6,7 @@ import useUser from "../../hooks/user.hook";
 
 const SignUp = () => {
   const config = useConfig();
-  const user = useUser();
+  const { user } = useUser();
   const router = useRouter();
   if (user) {
     router.replace("/");

frontend/src/pages/index.tsx

@@ -70,7 +70,7 @@ const useStyles = createStyles((theme) => ({
 
 export default function Home() {
   const config = useConfig();
-  const user = useUser();
+  const { user } = useUser();
 
   const { classes } = useStyles();
   const router = useRouter();

frontend/src/pages/upload.tsx

@@ -23,7 +23,7 @@ const Upload = () => {
   const router = useRouter();
   const modals = useModals();
 
-  const user = useUser();
+  const { user } = useUser();
   const config = useConfig();
   const [files, setFiles] = useState<FileUpload[]>([]);
   const [isUploading, setisUploading] = useState(false);

frontend/src/services/auth.service.ts

@@ -1,4 +1,4 @@
-import { getCookie, setCookies } from "cookies-next";
+import { getCookie, setCookie } from "cookies-next";
 import * as jose from "jose";
 import api from "./api.service";
 
@@ -11,8 +11,25 @@ const signIn = async (emailOrUsername: string, password: string) => {
     ...emailOrUsernameBody,
     password,
   });
-  setCookies("access_token", response.data.accessToken);
-  setCookies("refresh_token", response.data.refreshToken);
+  return response;
+};
+
+const signInTotp = async (
+  emailOrUsername: string,
+  password: string,
+  totp: string,
+  loginToken: string
+) => {
+  const emailOrUsernameBody = emailOrUsername.includes("@")
+    ? { email: emailOrUsername }
+    : { username: emailOrUsername };
+
+  const response = await api.post("auth/signIn/totp", {
+    ...emailOrUsernameBody,
+    password,
+    totp,
+    loginToken,
+  });
   return response;
 };
 
@@ -21,8 +38,8 @@ const signUp = async (email: string, username: string, password: string) => {
 };
 
 const signOut = () => {
-  setCookies("access_token", null);
-  setCookies("refresh_token", null);
+  setCookie("access_token", null);
+  setCookie("refresh_token", null);
   window.location.reload();
 };
 
@@ -37,7 +54,7 @@ const refreshAccessToken = async () => {
       const refreshToken = getCookie("refresh_token");
 
       const response = await api.post("auth/token", { refreshToken });
-      setCookies("access_token", response.data.accessToken);
+      setCookie("access_token", response.data.accessToken);
     }
   } catch {
     console.info("Refresh token invalid or expired");
@@ -48,10 +65,38 @@ const updatePassword = async (oldPassword: string, password: string) => {
   await api.patch("/auth/password", { oldPassword, password });
 };
 
+const enableTOTP = async (password: string) => {
+  const { data } = await api.post("/auth/totp/enable", { password });
+
+  return {
+    totpAuthUrl: data.totpAuthUrl,
+    totpSecret: data.totpSecret,
+    qrCode: data.qrCode,
+  };
+};
+
+const verifyTOTP = async (totpCode: string, password: string) => {
+  await api.post("/auth/totp/verify", {
+    code: totpCode,
+    password,
+  });
+};
+
+const disableTOTP = async (totpCode: string, password: string) => {
+  await api.post("/auth/totp/disable", {
+    code: totpCode,
+    password,
+  });
+};
+
 export default {
   signIn,
+  signInTotp,
   signUp,
   signOut,
   refreshAccessToken,
   updatePassword,
+  enableTOTP,
+  verifyTOTP,
+  disableTOTP,
 };

frontend/src/types/user.type.ts

@@ -3,6 +3,7 @@ type User = {
   username: string;
   email: string;
   isAdmin: boolean;
+  totpVerified: boolean;
 };
 
 export type CreateUser = {
@@ -26,4 +27,9 @@ export type UpdateCurrentUser = {
 
 export type CurrentUser = User & {};
 
+export type UserHook = {
+  user: CurrentUser | null;
+  setUser: (user: CurrentUser | null) => void;
+};
+
 export default User;