Переделал подтверждение почты

app/utils/auth.py

@@ -1,12 +1,13 @@
 from datetime import datetime, timedelta, UTC
 
-from cryptography.fernet import Fernet
+from cryptography.fernet import Fernet, InvalidToken
 from jose import jwt
 from passlib.context import CryptContext
 
 from app.config import settings
-from app.users.exceptions import IncorrectAuthDataException, UserNotFoundException, UserMustConfirmEmailException
+from app.users.exceptions import IncorrectAuthDataException, UserNotFoundException, UserMustConfirmEmailException, \
-from app.chat.exceptions import UserDontHavePermissionException, ChatNotFoundException, UserCanNotReadThisChatException
+	WrongCodeException
+from app.chat.exceptions import ChatNotFoundException, UserCanNotReadThisChatException
 from app.utils.unit_of_work import UnitOfWork
 from app.users.schemas import SUser, SInvitationData
 
@@ -37,9 +38,12 @@ def encode_invitation_token(user_data: SInvitationData) -> str:
 
 
 def decode_invitation_token(invitation_token: str) -> SInvitationData:
+	try:
 		user_code = invitation_token.encode()
 		user_data = cipher_suite.decrypt(user_code)
 		return SInvitationData.model_validate_json(user_data)
+	except InvalidToken:
+		raise WrongCodeException
 
 
 def encode_confirmation_token(confirmation_code: str) -> str:
@@ -48,17 +52,16 @@ def encode_confirmation_token(confirmation_code: str) -> str:
 
 
 def decode_confirmation_token(invitation_token: str) -> str:
+	try:
 		user_code = invitation_token.encode()
 		confirmation_code = cipher_suite.decrypt(user_code).decode()
 		return confirmation_code
+	except InvalidToken:
+		raise WrongCodeException
 
 
 def get_confirmation_code(user_code: str) -> str:
-	if len(user_code) == 6:
+	return user_code if len(user_code) == 6 else decode_confirmation_token(user_code)
-		confirmation_code = user_code
-	else:
-		confirmation_code = decode_confirmation_token(user_code)
-	return confirmation_code
 
 
 class AuthService: