Add mfa setup/validation/deletion

foxocorp · Dec 5, 2024 · 6fd3a3b · 6fd3a3b
1 parent f8037c6
commit 6fd3a3b
Show file tree

Hide file tree

Showing 9 changed files with 126 additions and 80 deletions.
diff --git a/foxogram-api/src/main/java/su/foxogram/controllers/UsersController.java b/foxogram-api/src/main/java/su/foxogram/controllers/UsersController.java
@@ -8,12 +8,15 @@
 import su.foxogram.constants.AttributesConstants;
 import su.foxogram.dtos.request.UserDeleteDTO;
 import su.foxogram.dtos.request.UserEditDTO;
+import su.foxogram.dtos.response.MFAKeyDTO;
 import su.foxogram.dtos.response.OkDTO;
 import su.foxogram.dtos.response.UserDTO;
 import su.foxogram.exceptions.*;
 import su.foxogram.models.User;
 import su.foxogram.services.UsersService;
 
+import java.security.NoSuchAlgorithmException;
+
 @Slf4j
 @RestController
 @RequestMapping(value = APIConstants.USERS, produces = "application/json")
@@ -49,16 +52,41 @@ public OkDTO deleteUser(@RequestAttribute(value = AttributesConstants.USER) User
 		String password = body.getPassword();
 		log.info("USER deletion requested ({}, {}) request", user.getId(), user.getEmail());
 
-		usersService.requestUserDelete(user, password, accessToken);
+		usersService.requestUserDelete(user, password);
 
 		return new OkDTO(true);
 	}
 
-	@PostMapping("/@me/delete/confirm/{code}")
-	public OkDTO deleteUserConfirm(@RequestAttribute(value = AttributesConstants.USER) User user, @RequestAttribute(value = AttributesConstants.ACCESS_TOKEN) String accessToken, @PathVariable String code, HttpServletRequest request) throws CodeIsInvalidException, CodeExpiredException {
+	@PostMapping("/@me/delete/confirm/")
+	public OkDTO deleteUserConfirm(@RequestAttribute(value = AttributesConstants.USER) User user, @RequestAttribute(value = AttributesConstants.ACCESS_TOKEN) String accessToken, HttpServletRequest request) throws CodeIsInvalidException, CodeExpiredException {
 		log.info("USER deletion confirm ({}, {}) request", user.getId(), user.getEmail());
 
-		usersService.confirmUserDelete(user, code);
+		usersService.confirmUserDelete(user);
+
+		return new OkDTO(true);
+	}
+
+	@PostMapping("/@me/mfa/")
+	public MFAKeyDTO setupMFA(@RequestAttribute(value = AttributesConstants.USER) User user) throws NoSuchAlgorithmException, MFAIsAlreadySetException {
+		log.info("USER mfa setup ({}, {}) request", user.getId(), user.getEmail());
+
+		String key = usersService.setupMFA(user);
+
+		return new MFAKeyDTO(key);
+	}
+
+	@DeleteMapping("/@me/mfa/")
+	public OkDTO deleteMFA(@RequestAttribute(value = AttributesConstants.USER) User user) throws NoSuchAlgorithmException, MFAIsAlreadySetException, MFAIsNotSetException {
+		log.info("USER mfa delete ({}, {}) request", user.getId(), user.getEmail());
+
+		usersService.deleteMFA(user);
+
+		return new OkDTO(true);
+	}
+
+	@PostMapping("/@me/mfa/setup/validate")
+	public OkDTO mfaValidate(@RequestAttribute(value = AttributesConstants.USER) User user) {
+		log.info("USER mfa validation ({}, {}) request", user.getId(), user.getEmail());
 
 		return new OkDTO(true);
 	}

diff --git a/foxogram-api/src/main/java/su/foxogram/services/UsersService.java b/foxogram-api/src/main/java/su/foxogram/services/UsersService.java
@@ -9,13 +9,15 @@
 import su.foxogram.constants.UserConstants;
 import su.foxogram.dtos.request.UserEditDTO;
 import su.foxogram.exceptions.*;
-import su.foxogram.models.Code;
 import su.foxogram.models.User;
 import su.foxogram.repositories.CodeRepository;
 import su.foxogram.repositories.UserRepository;
 import su.foxogram.util.CodeGenerator;
 import su.foxogram.util.Encryptor;
+import su.foxogram.util.Totp;
 
+import java.security.NoSuchAlgorithmException;
+import java.util.Base64;
 import java.util.Optional;
 
 @Slf4j
@@ -60,39 +62,36 @@ public User editUser(User user, UserEditDTO body) throws UserCredentialsDuplicat
 		return user;
 	}
 
-	public void requestUserDelete(User user, String password, String accessToken) throws UserCredentialsIsInvalidException, CodeIsInvalidException {
+	public void requestUserDelete(User user, String password) throws UserCredentialsIsInvalidException, CodeIsInvalidException {
 		if (!Encryptor.verifyPassword(password, user.getPassword()))
 			throw new UserCredentialsIsInvalidException();
 
 		sendEmail(user, EmailConstants.Type.ACCOUNT_DELETE);
 	}
 
-	public void confirmUserDelete(User user, String pathCode) throws CodeIsInvalidException, CodeExpiredException {
-		Code code = validateCode(pathCode);
-
-		deleteUserAndCode(user, code);
+	public void confirmUserDelete(User user) {
+		deleteUser(user);
 	}
 
-	private Code validateCode(String pathCode) throws CodeIsInvalidException, CodeExpiredException {
-		Code code = codeRepository.findByValue(pathCode);
+	public String setupMFA(User user) throws NoSuchAlgorithmException, MFAIsAlreadySetException {
+		if (user.hasFlag(UserConstants.Flags.MFA_ENABLED)) throw new MFAIsAlreadySetException();
 
-		if (code == null)
-			throw new CodeIsInvalidException();
+		String encodedKey = Base64.getEncoder().encodeToString(Totp.generateKey().getEncoded());
 
-		if (code.expiresAt <= System.currentTimeMillis())
-			throw new CodeExpiredException();
+		user.addFlag(UserConstants.Flags.MFA_ENABLED);
+		user.addFlag(UserConstants.Flags.AWAITING_CONFIRMATION);
+		user.setKey(encodedKey);
 
-		return code;
-	}
+		userRepository.save(user);
 
-	private void deleteUserAndCode(User user, Code code) {
-		deleteUser(user);
-		deleteVerificationCode(code);
+		return encodedKey;
 	}
 
-	private void deleteVerificationCode(Code code) {
-		codeRepository.delete(code);
-		log.info("CODE record deleted ({}, {}) successfully", code.getUserId(), code.getValue());
+	public void deleteMFA(User user) throws MFAIsNotSetException {
+		if (!user.hasFlag(UserConstants.Flags.MFA_ENABLED)) throw new MFAIsNotSetException();
+
+		user.removeFlag(UserConstants.Flags.MFA_ENABLED);
+		user.setKey(null);
 	}
 
 	private void deleteUser(User user) {

diff --git a/foxogram-common/src/main/java/su/foxogram/constants/ExceptionsConstants.java b/foxogram-common/src/main/java/su/foxogram/constants/ExceptionsConstants.java
@@ -16,7 +16,9 @@ public enum Codes {
 		CODE_EXPIRED(502),
 		NEED_TO_WAIT_BEFORE_RESEND(503),
 		TOTP_KEY_IS_INVALID(601),
-		MFA_IS_INVALID(602);
+		MFA_IS_INVALID(602),
+		MFA_IS_ALREADY_SET(603),
+		MFA_IS_NOT_SET(604);
 
 		private final int code;
 

diff --git a/foxogram-common/src/main/java/su/foxogram/dtos/response/MFAKeyDTO.java b/foxogram-common/src/main/java/su/foxogram/dtos/response/MFAKeyDTO.java
@@ -0,0 +1,14 @@
+package su.foxogram.dtos.response;
+
+import lombok.Getter;
+import lombok.Setter;
+
+@Getter
+@Setter
+public class MFAKeyDTO {
+	private String key;
+
+	public MFAKeyDTO(String key) {
+		this.key = key;
+	}
+}
diff --git a/foxogram-common/src/main/java/su/foxogram/exceptions/MFAIsAlreadySetException.java b/foxogram-common/src/main/java/su/foxogram/exceptions/MFAIsAlreadySetException.java
@@ -0,0 +1,13 @@
+package su.foxogram.exceptions;
+
+import org.springframework.http.HttpStatus;
+import org.springframework.web.bind.annotation.ResponseStatus;
+import su.foxogram.constants.ExceptionsConstants;
+
+@ResponseStatus(HttpStatus.FORBIDDEN)
+public class MFAIsAlreadySetException extends BaseException {
+
+	public MFAIsAlreadySetException() {
+		super("MFA is already set.", MFAIsAlreadySetException.class.getAnnotation(ResponseStatus.class).value(), ExceptionsConstants.Codes.MFA_IS_ALREADY_SET.getValue());
+	}
+}
diff --git a/foxogram-common/src/main/java/su/foxogram/exceptions/MFAIsNotSetException.java b/foxogram-common/src/main/java/su/foxogram/exceptions/MFAIsNotSetException.java
@@ -0,0 +1,13 @@
+package su.foxogram.exceptions;
+
+import org.springframework.http.HttpStatus;
+import org.springframework.web.bind.annotation.ResponseStatus;
+import su.foxogram.constants.ExceptionsConstants;
+
+@ResponseStatus(HttpStatus.FORBIDDEN)
+public class MFAIsNotSetException extends BaseException {
+
+	public MFAIsNotSetException() {
+		super("MFA is not set.", MFAIsNotSetException.class.getAnnotation(ResponseStatus.class).value(), ExceptionsConstants.Codes.MFA_IS_NOT_SET.getValue());
+	}
+}
diff --git a/foxogram-common/src/main/java/su/foxogram/interceptors/AuthenticationInterceptor.java b/foxogram-common/src/main/java/su/foxogram/interceptors/AuthenticationInterceptor.java
@@ -8,20 +8,21 @@
 import org.springframework.stereotype.Component;
 import org.springframework.web.servlet.HandlerInterceptor;
 import su.foxogram.constants.AttributesConstants;
-import su.foxogram.exceptions.UserEmailNotVerifiedException;
-import su.foxogram.exceptions.UserUnauthorizedException;
+import su.foxogram.constants.UserConstants;
+import su.foxogram.exceptions.*;
 import su.foxogram.models.User;
 import su.foxogram.services.AuthenticationService;
+import su.foxogram.util.Totp;
 
 import java.util.Set;
 
 @Component
 public class AuthenticationInterceptor implements HandlerInterceptor {
 
 	private static final Set<String> ALLOWED_PATHS = Set.of(
-			"email/verify",
-			"users/@me",
-			"email/resend"
+			"/v1/auth/email/verify",
+			"/v1/users/@me/**",
+			"/v1/auth/email/resend"
 	);
 
 	final AuthenticationService authenticationService;
@@ -32,21 +33,35 @@ public AuthenticationInterceptor(AuthenticationService authenticationService) {
 	}
 
 	@Override
-	public boolean preHandle(HttpServletRequest request, @NotNull HttpServletResponse response, @NotNull Object handler) throws UserUnauthorizedException, UserEmailNotVerifiedException {
+	public boolean preHandle(HttpServletRequest request, @NotNull HttpServletResponse response, @NotNull Object handler) throws UserUnauthorizedException, UserEmailNotVerifiedException, MFAIsInvalidException, TOTPKeyIsInvalidException, CodeExpiredException, CodeIsInvalidException {
 		String requestURI = request.getRequestURI();
-
-		boolean checkIfEmailVerified = ALLOWED_PATHS.stream().anyMatch(requestURI::contains);
+		boolean MFAValidationRequired = ALLOWED_PATHS.stream().anyMatch(requestURI::contains);
 
 		String accessToken = request.getHeader(HttpHeaders.AUTHORIZATION);
 
 		if (accessToken == null || !accessToken.startsWith("Bearer "))
 			throw new UserUnauthorizedException();
 
-		User user = authenticationService.getUser(accessToken, checkIfEmailVerified);
+		User user = authenticationService.getUser(accessToken);
 
-		request.setAttribute(AttributesConstants.USER, user);
-		request.setAttribute(AttributesConstants.ACCESS_TOKEN, accessToken);
+		if (MFAValidationRequired && user.hasFlag(UserConstants.Flags.AWAITING_CONFIRMATION)) {
+			validateMFA(user, request);
+		}
 
 		return true;
 	}
+
+	private void validateMFA(User user, HttpServletRequest request) throws MFAIsInvalidException, TOTPKeyIsInvalidException, CodeExpiredException, CodeIsInvalidException {
+		String code = request.getHeader("Code");
+
+		if (!user.hasFlag(UserConstants.Flags.MFA_ENABLED)) {
+			boolean MFAVerified = authenticationService.validateCode(code) != null;
+
+			request.setAttribute(AttributesConstants.MFA_VERIFIED, MFAVerified);
+		} else {
+			boolean MFAVerified = Totp.validate(user.getKey(), code);
+
+			request.setAttribute(AttributesConstants.MFA_VERIFIED, MFAVerified);
+		}
+	}
 }
diff --git a/foxogram-common/src/main/java/su/foxogram/interceptors/MFAInterceptor.java b/foxogram-common/src/main/java/su/foxogram/interceptors/MFAInterceptor.java
diff --git a/foxogram-common/src/main/java/su/foxogram/services/AuthenticationService.java b/foxogram-common/src/main/java/su/foxogram/services/AuthenticationService.java
@@ -17,10 +17,8 @@
 import su.foxogram.structures.Snowflake;
 import su.foxogram.util.CodeGenerator;
 import su.foxogram.util.Encryptor;
-import su.foxogram.util.Totp;
 
 import java.security.NoSuchAlgorithmException;
-import java.util.Base64;
 
 @Slf4j
 @Service
@@ -44,19 +42,14 @@ public AuthenticationService(UserRepository userRepository, CodeRepository codeR
 		this.apiConfig = apiConfig;
 	}
 
-	public User getUser(String header, boolean checkIfEmailVerified) throws UserUnauthorizedException, UserEmailNotVerifiedException {
-		return validate(header.substring(7), checkIfEmailVerified);
+	public User getUser(String header) throws UserUnauthorizedException, UserEmailNotVerifiedException {
+		return validate(header.substring(7));
 	}
 
-	public User validate(String token, boolean checkIfEmailVerified) throws UserUnauthorizedException, UserEmailNotVerifiedException {
+	public User validate(String token) throws UserUnauthorizedException, UserEmailNotVerifiedException {
 		String userId = jwtService.validate(token).getId();
 
-		User user = userRepository.findById(userId).orElseThrow(UserUnauthorizedException::new);
-
-		if (user.hasFlag(UserConstants.Flags.AWAITING_CONFIRMATION) && !checkIfEmailVerified)
-			throw new UserEmailNotVerifiedException();
-
-		return user;
+		return userRepository.findById(userId).orElseThrow(UserUnauthorizedException::new);
 	}
 
 	public String userSignUp(String username, String email, String password) throws UserCredentialsDuplicateException, NoSuchAlgorithmException {
@@ -79,7 +72,7 @@ private User createUser(String username, String email, String password) throws N
 		String avatar = new Avatar("").getId();
 		long flags = UserConstants.Flags.AWAITING_CONFIRMATION.getBit();
 		int type = UserConstants.Type.USER.getType();
-		String key = Base64.getEncoder().encodeToString(Totp.generateKey().getEncoded());
+		String key = null;
 
 		return new User(id, avatar, null, username, email, Encryptor.hashPassword(password), flags, type, deletion, key);
 	}
@@ -135,7 +128,7 @@ public void verifyEmail(User user, String pathCode) throws CodeIsInvalidExceptio
 		deleteVerificationCode(code);
 	}
 
-	private Code validateCode(String pathCode) throws CodeIsInvalidException, CodeExpiredException {
+	public Code validateCode(String pathCode) throws CodeIsInvalidException, CodeExpiredException {
 		Code code = codeRepository.findByValue(pathCode);
 
 		if (code == null)