From 2ce40700bd23b77b35119a8a2091816429d7506c Mon Sep 17 00:00:00 2001
From: urec56 <mail@urec56.ru>
Date: Sat, 13 Jul 2024 20:16:22 +0400
Subject: [PATCH] =?UTF-8?q?=D0=98=D0=B7=D0=BC=D0=B5=D0=BD=D0=B5=D0=BD?=
 =?UTF-8?q?=D0=B8=D1=8F=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=D1=8F=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/tasks/tasks.py   |  2 +-
 app/users/router.py  | 42 ++++++++++++++----------------------------
 app/users/schemas.py |  4 ++--
 app/utils/auth.py    | 10 +++++-----
 4 files changed, 22 insertions(+), 36 deletions(-)

diff --git a/app/tasks/tasks.py b/app/tasks/tasks.py
index 5c6ed07..8fa4c34 100644
--- a/app/tasks/tasks.py
+++ b/app/tasks/tasks.py
@@ -25,7 +25,7 @@ def generate_confirmation_code(length=6) -> str:
 @celery.task
 def send_registration_confirmation_email(user_data: dict) -> None:
 	user_data = SConfirmationData.model_validate(user_data)
-	invitation_token = encode_confirmation_token(user_data)
+	invitation_token = encode_confirmation_token(user_data.confirmation_code)
 
 	confirmation_link = settings.INVITATION_LINK_HOST + "/api/users/email_verification/link/" + invitation_token
 
diff --git a/app/users/router.py b/app/users/router.py
index 97cdeb9..18b6b83 100644
--- a/app/users/router.py
+++ b/app/users/router.py
@@ -23,7 +23,6 @@ from app.users.schemas import (
 	SUserLogin,
 	SUserRegister,
 	SUserResponse,
-	SEmailVerification,
 	SUserAvatars,
 	SUserFilter,
 	SUser,
@@ -34,6 +33,7 @@ from app.users.schemas import (
 	SConfirmationData,
 	SUserPassword,
 	SGetUsersFilter,
+	SUserCode,
 )
 from app.tasks.tasks import (
 	send_registration_confirmation_email,
@@ -131,39 +131,25 @@ async def resend_email_verification(user: SUser = Depends(get_current_user)):
 	await RedisService.set_verification_code(redis=redis_session, user_id=user.id, verification_code=user_code)
 
 
-@router.get(
-	"/email_verification/link/{user_code}",
+@router.post(
+	"/email_verification/{user_code}",
 	status_code=status.HTTP_200_OK,
-	response_model=SEmailVerification,
+	response_model=None,
 )
-async def email_verification(user_code: str, uow=Depends(UnitOfWork)):
-	user_data = decode_confirmation_token(user_code)
+async def email_verification(user_code: SUserCode, user: SUser = Depends(get_current_user),uow=Depends(UnitOfWork)):
+	if len(user_code.user_code) == 6:
+		confirmation_code = user_code
+	else:
+		confirmation_code = decode_confirmation_token(user_code.user_code)
+
 	redis_session = get_redis_session()
+	verification_code = await RedisService.get_verification_code(redis=redis_session, user_id=user.id)
+
+	if verification_code != confirmation_code:
+		raise WrongCodeException
 	async with uow:
-		verification_code = await RedisService.get_verification_code(redis=redis_session, user_id=user_data.user_id)
-		if verification_code != user_data.confirmation_code:
-			raise WrongCodeException
-
-		await uow.user.change_data(user_id=user_data.user_id, role=settings.VERIFICATED_USER)
-		await uow.commit()
-		return {"email_verification": True}
-
-
-@router.get(
-	"/email_verification/code/{user_code}",
-	status_code=status.HTTP_200_OK,
-	response_model=SEmailVerification,
-)
-async def email_verification(user_code: str, user: SUser = Depends(get_current_user), uow=Depends(UnitOfWork)):
-	redis_session = get_redis_session()
-	async with uow:
-		verification_code = await RedisService.get_verification_code(redis=redis_session, user_id=user.id)
-		if verification_code != user_code:
-			raise WrongCodeException
-
 		await uow.user.change_data(user_id=user.id, role=settings.VERIFICATED_USER)
 		await uow.commit()
-		return {"email_verification": True}
 
 
 @router.post(
diff --git a/app/users/schemas.py b/app/users/schemas.py
index 84cc1ba..3875e0b 100644
--- a/app/users/schemas.py
+++ b/app/users/schemas.py
@@ -76,8 +76,8 @@ class STokenLogin(BaseModel):
 	authorization: str
 
 
-class SEmailVerification(BaseModel):
-	email_verification: bool
+class SUserCode(BaseModel):
+	user_code: str
 
 
 class SCreateInvitationLink(BaseModel):
diff --git a/app/utils/auth.py b/app/utils/auth.py
index 7f59a3b..828aa52 100644
--- a/app/utils/auth.py
+++ b/app/utils/auth.py
@@ -44,15 +44,15 @@ def decode_invitation_token(invitation_token: str) -> SInvitationData:
 	return SInvitationData.model_validate_json(user_data)
 
 
-def encode_confirmation_token(user_data: SConfirmationData) -> str:
-	invitation_token = cipher_suite.encrypt(user_data.model_dump_json().encode())
+def encode_confirmation_token(confirmation_code: str) -> str:
+	invitation_token = cipher_suite.encrypt(confirmation_code.encode())
 	return invitation_token.decode()
 
 
-def decode_confirmation_token(invitation_token: str) -> SConfirmationData:
+def decode_confirmation_token(invitation_token: str) -> str:
 	user_code = invitation_token.encode()
-	user_data = cipher_suite.decrypt(user_code)
-	return SConfirmationData.model_validate_json(user_data)
+	confirmation_code = cipher_suite.decrypt(user_code).decode()
+	return confirmation_code
 
 
 class AuthService: