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

---
 .../api/controller/v1/UserController.java     |  7 +++--
 .../api/model/user/input/DisableTFAInput.java | 28 +++++++++++++++++++
 .../cc/pulseapp/api/service/UserService.java  | 22 +++++++++++++--
 3 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 src/main/java/cc/pulseapp/api/model/user/input/DisableTFAInput.java

diff --git a/src/main/java/cc/pulseapp/api/controller/v1/UserController.java b/src/main/java/cc/pulseapp/api/controller/v1/UserController.java
index ebd91c1..f6c3f94 100644
--- a/src/main/java/cc/pulseapp/api/controller/v1/UserController.java
+++ b/src/main/java/cc/pulseapp/api/controller/v1/UserController.java
@@ -4,6 +4,7 @@ import cc.pulseapp.api.exception.impl.BadRequestException;
 import cc.pulseapp.api.model.user.User;
 import cc.pulseapp.api.model.user.UserDTO;
 import cc.pulseapp.api.model.user.input.CompleteOnboardingInput;
+import cc.pulseapp.api.model.user.input.DisableTFAInput;
 import cc.pulseapp.api.model.user.input.EnableTFAInput;
 import cc.pulseapp.api.model.user.input.UserExistsInput;
 import cc.pulseapp.api.model.user.response.UserSetupTFAResponse;
@@ -100,11 +101,13 @@ public final class UserController {
     /**
      * A POST endpoint to disable TFA for a useer.
      *
+     * @param input the input to process
      * @return the disabled response
+     * @throws BadRequestException if disabling fails
      */
     @PostMapping("/disable-tfa") @ResponseBody @NonNull
-    public ResponseEntity<Map<String, Object>> disableTwoFactor() {
-        userService.disableTwoFactor();
+    public ResponseEntity<Map<String, Object>> disableTwoFactor(DisableTFAInput input) throws BadRequestException {
+        userService.disableTwoFactor(input);
         return ResponseEntity.ok(Map.of("success", true));
     }
 
diff --git a/src/main/java/cc/pulseapp/api/model/user/input/DisableTFAInput.java b/src/main/java/cc/pulseapp/api/model/user/input/DisableTFAInput.java
new file mode 100644
index 0000000..3bcc51b
--- /dev/null
+++ b/src/main/java/cc/pulseapp/api/model/user/input/DisableTFAInput.java
@@ -0,0 +1,28 @@
+package cc.pulseapp.api.model.user.input;
+
+import cc.pulseapp.api.model.user.User;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.ToString;
+
+/**
+ * The input to disable TFA for a {@link User}.
+ *
+ * @author Braydon
+ */
+@AllArgsConstructor @Getter @ToString
+public final class DisableTFAInput {
+    /**
+     * The TFA pin to validate.
+     */
+    private final String pin;
+
+    /**
+     * Check if this input is valid.
+     *
+     * @return whether this input is valid
+     */
+    public boolean isValid() {
+        return pin != null && (pin.length() == 6);
+    }
+}
\ No newline at end of file
diff --git a/src/main/java/cc/pulseapp/api/service/UserService.java b/src/main/java/cc/pulseapp/api/service/UserService.java
index 0b1b100..85f853d 100644
--- a/src/main/java/cc/pulseapp/api/service/UserService.java
+++ b/src/main/java/cc/pulseapp/api/service/UserService.java
@@ -11,6 +11,7 @@ import cc.pulseapp.api.model.user.User;
 import cc.pulseapp.api.model.user.UserDTO;
 import cc.pulseapp.api.model.user.UserFlag;
 import cc.pulseapp.api.model.user.input.CompleteOnboardingInput;
+import cc.pulseapp.api.model.user.input.DisableTFAInput;
 import cc.pulseapp.api.model.user.input.EnableTFAInput;
 import cc.pulseapp.api.model.user.input.UserExistsInput;
 import cc.pulseapp.api.model.user.response.UserSetupTFAResponse;
@@ -197,7 +198,7 @@ public final class UserService {
             throw new BadRequestException(Error.TFA_SETUP_SECRET_MISMATCH);
         }
         if (!tfaService.getPin(secret).equals(input.getPin())) { // Ensure the pin is valid
-            throw new BadRequestException(Error.TFA_SETUP_PIN_INVALID);
+            throw new BadRequestException(Error.TFA_PIN_INVALID);
         }
         // Enable TFA for the user
         byte[] salt = HashUtils.generateSalt();
@@ -225,9 +226,22 @@ public final class UserService {
     /**
      * Disable two-factor auth for the
      * currently authenticated user.
+     *
+     * @param input the input to process
+     * @throws BadRequestException if disabling fails
      */
-    public void disableTwoFactor() {
+    public void disableTwoFactor(DisableTFAInput input) throws BadRequestException {
+        if (input == null || (!input.isValid())) { // Ensure the input was provided
+            throw new BadRequestException(Error.MALFORMED_DISABLE_TFA_INPUT);
+        }
         User user = authService.getAuthenticatedUser();
+        if (!user.hasFlag(UserFlag.TFA_ENABLED)) { // Ensure TFA is already on
+            throw new BadRequestException(Error.TFA_NOT_ENABLED);
+        }
+        if (!tfaService.getPin(user.getTfa().getSecret()).equals(input.getPin())) { // Ensure the pin is valid
+            throw new BadRequestException(Error.TFA_PIN_INVALID);
+        }
+        // Disable TFA for the user
         user.setTfa(null);
         user.removeFlag(UserFlag.TFA_ENABLED);
         userRepository.save(user);
@@ -247,11 +261,13 @@ public final class UserService {
         MALFORMED_USER_EXISTS_INPUT,
         MALFORMED_ONBOARDING_INPUT,
         MALFORMED_ENABLE_TFA_INPUT,
+        MALFORMED_DISABLE_TFA_INPUT,
         ORGANIZATION_SLUG_INVALID,
         ALREADY_ONBOARDED,
         TFA_ALREADY_ENABLED,
+        TFA_NOT_ENABLED,
         TFA_SETUP_TIMED_OUT,
         TFA_SETUP_SECRET_MISMATCH,
-        TFA_SETUP_PIN_INVALID,
+        TFA_PIN_INVALID,
     }
 }
\ No newline at end of file