diff --git a/app/utils/auth.py b/app/utils/auth.py index 1addc5c..be16c47 100644 --- a/app/utils/auth.py +++ b/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.chat.exceptions import UserDontHavePermissionException, ChatNotFoundException, UserCanNotReadThisChatException +from app.users.exceptions import IncorrectAuthDataException, UserNotFoundException, UserMustConfirmEmailException, \ + 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: - user_code = invitation_token.encode() - user_data = cipher_suite.decrypt(user_code) - return SInvitationData.model_validate_json(user_data) + 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: - user_code = invitation_token.encode() - confirmation_code = cipher_suite.decrypt(user_code).decode() - return confirmation_code + 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: - confirmation_code = user_code - else: - confirmation_code = decode_confirmation_token(user_code) - return confirmation_code + return user_code if len(user_code) == 6 else decode_confirmation_token(user_code) class AuthService: