impl tfa setup

pom.xml

@@ -81,9 +81,9 @@
             <scope>provided</scope>
         </dependency>
         <dependency>
-            <groupId>com.relops</groupId>
+            <groupId>com.github.ben-manes.caffeine</groupId>
-            <artifactId>snowflake</artifactId>
+            <artifactId>caffeine</artifactId>
-            <version>1.1</version>
+            <version>3.1.8</version>
             <scope>compile</scope>
         </dependency>
         <dependency>
@@ -92,6 +92,18 @@
             <version>7.2.0</version>
             <scope>compile</scope>
         </dependency>
+        <dependency>
+            <groupId>com.relops</groupId>
+            <artifactId>snowflake</artifactId>
+            <version>1.1</version>
+            <scope>compile</scope>
+        </dependency>
+        <dependency>
+            <groupId>de.taimos</groupId>
+            <artifactId>totp</artifactId>
+            <version>1.0</version>
+            <scope>compile</scope>
+        </dependency>
 
         <!-- Unirest -->
         <dependency>

src/main/java/cc/pulseapp/api/controller/v1/UserController.java

@@ -4,6 +4,8 @@ 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.EnableTFAInput;
+import cc.pulseapp.api.model.user.response.UserSetupTFAResponse;
 import cc.pulseapp.api.service.UserService;
 import lombok.NonNull;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -11,6 +13,7 @@ import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.*;
 
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -68,4 +71,21 @@ public final class UserController {
         userService.completeOnboarding(input);
         return ResponseEntity.ok(Map.of("success", true));
     }
+
+    /**
+     * A POST endpoint to start
+     * setting up TFA for a user.
+     *
+     * @return the setup response
+     * @throws BadRequestException if the setup fails
+     */
+    @PostMapping("/setup-tfa") @ResponseBody @NonNull
+    public ResponseEntity<UserSetupTFAResponse> setupTwoFactor() throws BadRequestException {
+        return ResponseEntity.ok(userService.setupTwoFactor());
+    }
+
+    @PostMapping("/enable-tfa") @ResponseBody @NonNull
+    public ResponseEntity<List<String>> enableTwoFactor(EnableTFAInput input) throws BadRequestException {
+        return ResponseEntity.ok(userService.enableTwoFactor(input));
+    }
 }

src/main/java/cc/pulseapp/api/model/user/TFAProfile.java

@@ -0,0 +1,32 @@
+package cc.pulseapp.api.model.user;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.ToString;
+
+import java.util.List;
+
+/**
+ * The two-factor authentication
+ * profile of a {@link User}.
+ *
+ * @author Braydon
+ */
+@AllArgsConstructor @Getter @ToString
+public final class TFAProfile {
+    /**
+     * The TFA secret of the user.
+     */
+    @NonNull private final String secret;
+
+    /**
+     * The sale for the user's backup codes.
+     */
+    @NonNull private final String backupCodesSalt;
+
+    /**
+     * The (encrypted) backup codes of the user.
+     */
+    @NonNull private final List<String> backupCodes;
+}

src/main/java/cc/pulseapp/api/model/user/User.java

@@ -34,7 +34,7 @@ public final class User {
     @Indexed @NonNull private final String username;
 
     /**
-     * The password for this user.
+     * The (encrypted) password for this user.
      */
     @NonNull private final String password;
 
@@ -53,6 +53,12 @@ public final class User {
      */
     @NonNull private final UserTier tier;
 
+    /**
+     * The TFA profile of this
+     * user, present if TFA is enabled.
+     */
+    private TFAProfile tfa;
+
     /**
      * The flags for this user.
      */

src/main/java/cc/pulseapp/api/model/user/UserFlag.java

@@ -16,6 +16,11 @@ public enum UserFlag {
      */
     COMPLETED_ONBOARDING,
 
+    /**
+     * The user has two-factor auth enabled.
+     */
+    TFA_ENABLED,
+
     /**
      * The user is an administrator.
      */

src/main/java/cc/pulseapp/api/model/user/input/EnableTFAInput.java

@@ -0,0 +1,34 @@
+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 enable TFA for a {@link User}.
+ *
+ * @author Braydon
+ */
+@AllArgsConstructor @Getter @ToString
+public final class EnableTFAInput {
+    /**
+     * The TFA secret to use.
+     */
+    private final String secret;
+
+    /**
+     * 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 secret != null && (!secret.isBlank())
+                && pin != null && (pin.length() == 6);
+    }
+}

src/main/java/cc/pulseapp/api/model/user/response/UserAuthResponse.java

@@ -1,7 +1,7 @@
 package cc.pulseapp.api.model.user.response;
 
-import cc.pulseapp.api.model.user.session.Session;
 import cc.pulseapp.api.model.user.UserDTO;
+import cc.pulseapp.api.model.user.session.Session;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.NonNull;

src/main/java/cc/pulseapp/api/model/user/response/UserSetupTFAResponse.java

@@ -0,0 +1,27 @@
+package cc.pulseapp.api.model.user.response;
+
+import cc.pulseapp.api.model.user.User;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.ToString;
+
+/**
+ * The response for when a {@link User}
+ * initializes the setup of two-factor
+ * authentication.
+ *
+ * @author Braydon
+ */
+@AllArgsConstructor @Getter @ToString
+public final class UserSetupTFAResponse {
+    /**
+     * The TFA secret.
+     */
+    @NonNull private final String secret;
+
+    /**
+     * The URL to the QR code.
+     */
+    @NonNull private final String qrCodeUrl;
+}

src/main/java/cc/pulseapp/api/model/user/session/Session.java

@@ -25,7 +25,7 @@ public final class Session {
     /**
      * The snowflake of the user this session is for.
      */
-    @JsonIgnore private final long userSnowflake;
+    @Indexed @JsonIgnore private final long userSnowflake;
 
     /**
      * The access token for the user.

src/main/java/cc/pulseapp/api/repository/OrganizationRepository.java

@@ -3,6 +3,7 @@ package cc.pulseapp.api.repository;
 import cc.pulseapp.api.model.org.Organization;
 import lombok.NonNull;
 import org.springframework.data.mongodb.repository.MongoRepository;
+import org.springframework.stereotype.Repository;
 
 import java.util.List;
 
@@ -11,6 +12,7 @@ import java.util.List;
  *
  * @author Braydon
  */
+@Repository
 public interface OrganizationRepository extends MongoRepository<Organization, Long> {
     /**
      * Find an organization by its name.

src/main/java/cc/pulseapp/api/repository/SessionRepository.java

@@ -3,12 +3,16 @@ package cc.pulseapp.api.repository;
 import cc.pulseapp.api.model.user.session.Session;
 import lombok.NonNull;
 import org.springframework.data.repository.CrudRepository;
+import org.springframework.stereotype.Repository;
+
+import java.util.List;
 
 /**
  * The repository for {@link Session}'s.
  *
  * @author Braydon
  */
+@Repository
 public interface SessionRepository extends CrudRepository<Session, String> {
     /**
      * Find a session by the access token.
@@ -17,4 +21,12 @@ public interface SessionRepository extends CrudRepository<Session, String> {
      * @return the session, null if none
      */
     Session findByAccessToken(@NonNull String accessToken);
+
+    /**
+     * Get all sessions for a user.
+     *
+     * @param userSnowflake the user's snowflake
+     * @return the sessions
+     */
+    List<Session> findAllByUserSnowflake(long userSnowflake);
 }

src/main/java/cc/pulseapp/api/repository/StatusPageRepository.java

@@ -3,6 +3,7 @@ package cc.pulseapp.api.repository;
 import cc.pulseapp.api.model.page.StatusPage;
 import lombok.NonNull;
 import org.springframework.data.mongodb.repository.MongoRepository;
+import org.springframework.stereotype.Repository;
 
 import java.util.List;
 
@@ -11,6 +12,7 @@ import java.util.List;
  *
  * @author Braydon
  */
+@Repository
 public interface StatusPageRepository extends MongoRepository<StatusPage, Long> {
     /**
      * Find a status page by its name.

src/main/java/cc/pulseapp/api/repository/UserRepository.java

@@ -3,12 +3,14 @@ package cc.pulseapp.api.repository;
 import cc.pulseapp.api.model.user.User;
 import lombok.NonNull;
 import org.springframework.data.mongodb.repository.MongoRepository;
+import org.springframework.stereotype.Repository;
 
 /**
  * The repository for interacting with {@link User}'s.
  *
  * @author Braydon
  */
+@Repository
 public interface UserRepository extends MongoRepository<User, Long> {
     /**
      * Find a user by their email.

src/main/java/cc/pulseapp/api/service/AuthService.java

@@ -1,13 +1,15 @@
 package cc.pulseapp.api.service;
 
 import cc.pulseapp.api.common.HashUtils;
-import cc.pulseapp.api.common.RequestUtils;
 import cc.pulseapp.api.common.StringUtils;
 import cc.pulseapp.api.exception.impl.BadRequestException;
 import cc.pulseapp.api.exception.impl.ResourceNotFoundException;
 import cc.pulseapp.api.model.Feature;
 import cc.pulseapp.api.model.IGenericResponse;
-import cc.pulseapp.api.model.user.*;
+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.UserTier;
 import cc.pulseapp.api.model.user.input.UserLoginInput;
 import cc.pulseapp.api.model.user.input.UserRegistrationInput;
 import cc.pulseapp.api.model.user.response.UserAuthResponse;
@@ -86,7 +88,7 @@ public final class AuthService {
         User user = userRepository.save(new User(
                 snowflakeService.generateSnowflake(), input.getEmail(), input.getUsername().toLowerCase(),
                 HashUtils.hash(salt, input.getPassword()), Base64.getEncoder().encodeToString(salt),
-                null, UserTier.FREE, 0, now
+                null, UserTier.FREE, null, 0, now
         ));
         return new UserAuthResponse(generateSession(request, user), UserDTO.asDTO(user, now));
     }

src/main/java/cc/pulseapp/api/service/TFAService.java

@@ -0,0 +1,65 @@
+package cc.pulseapp.api.service;
+
+import de.taimos.totp.TOTP;
+import lombok.NonNull;
+import org.apache.commons.codec.binary.Base32;
+import org.apache.commons.codec.binary.Hex;
+import org.springframework.stereotype.Service;
+
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
+import java.security.SecureRandom;
+
+/**
+ * This service is responsible for
+ * two factor authentication.
+ *
+ * @author Braydon
+ */
+@Service
+public final class TFAService {
+    private static final Base32 BASE_32 = new Base32();
+    private static final String APP_ISSUER = "Pulse App";
+
+    /**
+     * Generate a secret key.
+     *
+     * @return the secret key
+     */
+    @NonNull
+    public String generateSecretKey() {
+        SecureRandom random = new SecureRandom();
+        byte[] bytes = new byte[15];
+        random.nextBytes(bytes);
+        return BASE_32.encodeToString(bytes);
+    }
+
+    /**
+     * Generate a QR code URL for a user.
+     *
+     * @param username the user's username
+     * @param secret   the user's tfa secret
+     * @return the qr code url
+     */
+    @NonNull
+    public String generateQrCodeUrl(@NonNull String username, @NonNull String secret) {
+        return "otpauth://totp/"
+                + URLEncoder.encode(APP_ISSUER + ":" + username, StandardCharsets.UTF_8).replace("+", "%20")
+                + "?issuer=" + URLEncoder.encode(APP_ISSUER, StandardCharsets.UTF_8).replace("+", "%20")
+                + "&secret=" + URLEncoder.encode(secret, StandardCharsets.UTF_8).replace("+", "%20");
+    }
+
+    /**
+     * Get the current 6-digit pin
+     * from the given secret key.
+     *
+     * @param secretKey the secret key
+     * @return the 6-digit pin
+     */
+    @NonNull
+    public String getPin(@NonNull String secretKey) {
+        byte[] bytes = BASE_32.decode(secretKey);
+        String hexKey = Hex.encodeHexString(bytes);
+        return TOTP.getOTP(hexKey);
+    }
+}

src/main/java/cc/pulseapp/api/service/UserService.java

@@ -1,19 +1,30 @@
 package cc.pulseapp.api.service;
 
+import cc.pulseapp.api.common.HashUtils;
 import cc.pulseapp.api.common.StringUtils;
 import cc.pulseapp.api.exception.impl.BadRequestException;
 import cc.pulseapp.api.model.IGenericResponse;
 import cc.pulseapp.api.model.org.Organization;
+import cc.pulseapp.api.model.user.TFAProfile;
 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.EnableTFAInput;
+import cc.pulseapp.api.model.user.response.UserSetupTFAResponse;
+import cc.pulseapp.api.repository.SessionRepository;
 import cc.pulseapp.api.repository.UserRepository;
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import lombok.NonNull;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 
+import java.util.ArrayList;
+import java.util.Base64;
 import java.util.Date;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 /**
  * @author Braydon
@@ -40,20 +51,44 @@ public final class UserService {
      */
     @NonNull private final StatusPageService statusPageService;
 
+    /**
+     * Thw two-factor auth service to use.
+     */
+    @NonNull private final TFAService tfaService;
+
     /**
      * The user repository to use.
      */
     @NonNull private final UserRepository userRepository;
 
+    /**
+     * The session repository to use.
+     */
+    @NonNull private final SessionRepository sessionRepository;
+
+    /**
+     * A map of users who are setting up two-factor auth.
+     * <p>
+     * The key kis the user's snowflake
+     * and the value is the secret.
+     * </p>
+     */
+    private final Cache<Long, String> settingUpTfa = Caffeine.newBuilder()
+            .expireAfterWrite(5L, TimeUnit.MINUTES)
+            .build();
+
     @Autowired
     public UserService(@NonNull AuthService authService, @NonNull SnowflakeService snowflakeService,
                        @NonNull OrganizationService orgService, @NonNull StatusPageService statusPageService,
-                       @NonNull UserRepository userRepository) {
+                       @NonNull TFAService tfaService, @NonNull UserRepository userRepository,
+                       @NonNull SessionRepository sessionRepository) {
         this.authService = authService;
         this.snowflakeService = snowflakeService;
         this.orgService = orgService;
         this.statusPageService = statusPageService;
+        this.tfaService = tfaService;
         this.userRepository = userRepository;
+        this.sessionRepository = sessionRepository;
     }
 
     /**
@@ -103,12 +138,74 @@ public final class UserService {
         userRepository.save(user);
     }
 
+    /**
+     * Start setting up TFA for a user.
+     *
+     * @return the setup response
+     * @throws BadRequestException if the setup fails
+     */
+    @NonNull
+    public UserSetupTFAResponse setupTwoFactor() throws BadRequestException {
+        User user = authService.getAuthenticatedUser();
+        if (user.hasFlag(UserFlag.TFA_ENABLED)) { // Ensure TFA isn't already on
+            throw new BadRequestException(Error.TFA_ALREADY_ENABLED);
+        }
+        String secret = tfaService.generateSecretKey();
+        settingUpTfa.put(user.getSnowflake(), secret); // Store temporarily
+        return new UserSetupTFAResponse(secret, tfaService.generateQrCodeUrl(user.getUsername(), secret));
+    }
+
+    @NonNull
+    public List<String> enableTwoFactor(EnableTFAInput input) throws BadRequestException {
+        if (input == null || (!input.isValid())) { // Ensure the input was provided
+            throw new BadRequestException(Error.MALFORMED_ENABLE_TFA_INPUT);
+        }
+        User user = authService.getAuthenticatedUser();
+        if (user.hasFlag(UserFlag.TFA_ENABLED)) { // Ensure TFA isn't already on
+            throw new BadRequestException(Error.TFA_ALREADY_ENABLED);
+        }
+        String secret = settingUpTfa.getIfPresent(user.getSnowflake()); // Get the setup TFA secret
+        if (secret == null) { // No secret, creation session timed out
+            throw new BadRequestException(Error.TFA_SETUP_TIMED_OUT);
+        }
+        if (!secret.equals(input.getSecret())) { // Ensure the original and received secrets are the same
+            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);
+        }
+        // Enable TFA for the user
+        byte[] salt = HashUtils.generateSalt();
+        List<String> originalBackupCodes = new ArrayList<>();
+        for (int i = 0; i < 8; i++) { // Generate 8 backup codes
+            originalBackupCodes.add(StringUtils.generateRandom(6, false, true, false));
+        }
+
+        // Encrypt the stored backup codes
+        List<String> storedBackupCodes = originalBackupCodes.stream()
+                .map(backupCode -> HashUtils.hash(salt, backupCode))
+                .toList();
+        user.setTfa(new TFAProfile(secret, Base64.getEncoder().encodeToString(salt), storedBackupCodes));
+        user.addFlag(UserFlag.TFA_ENABLED);
+        userRepository.save(user);
+
+        // And finally invalidate all of the sessions for the user
+        sessionRepository.deleteAll(sessionRepository.findAllByUserSnowflake(user.getSnowflake()));
+
+        return originalBackupCodes;
+    }
+
     /**
      * User errors.
      */
     private enum Error implements IGenericResponse {
         MALFORMED_ONBOARDING_INPUT,
+        MALFORMED_ENABLE_TFA_INPUT,
         ORGANIZATION_SLUG_INVALID,
         ALREADY_ONBOARDED,
+        TFA_ALREADY_ENABLED,
+        TFA_SETUP_TIMED_OUT,
+        TFA_SETUP_SECRET_MISMATCH,
+        TFA_SETUP_PIN_INVALID,
     }
 }