From 26ed666855818e72b2d78010a0b1adf6c8535d36 Mon Sep 17 00:00:00 2001 From: urec56 Date: Tue, 16 Jul 2024 18:54:36 +0400 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=B5=D1=80=D0=B5=D0=B4=D0=B5=D0=BB?= =?UTF-8?q?=D0=B0=D0=BB=20=D0=BF=D0=BE=D0=B4=D1=82=D0=B2=D0=B5=D1=80=D0=B6?= =?UTF-8?q?=D0=B4=D0=B5=D0=BD=D0=B8=D0=B5=20=D0=BF=D0=BE=D1=87=D1=82=D1=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/utils/auth.py | 31 +++++++++++++++++-------------- 1 file changed, 17 insertions(+), 14 deletions(-) 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: