From 3d9a3cef731f2c25111fda58676fc5e43ac428c7 Mon Sep 17 00:00:00 2001
From: admin <fakc445@gmail.com>
Date: Wed, 5 Mar 2025 19:34:16 +0300
Subject: [PATCH] feat: invite code

---
 models/auth.py   |  4 ++--
 routers/user.py  |  3 ++-
 schemas/user.py  |  4 ++++
 services/auth.py | 13 ++++++-------
 services/user.py |  5 +++--
 5 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/models/auth.py b/models/auth.py
index ea018c4..64244bb 100644
--- a/models/auth.py
+++ b/models/auth.py
@@ -144,8 +144,8 @@ class InviteCode(BaseModel):
     code: Mapped[str] = mapped_column(primary_key=True)
 
     is_activated: Mapped[bool] = mapped_column(nullable=False, default=False)
-    activated_by_id: Mapped[Optional[int]] = mapped_column(ForeignKey('users.id'))
-    activated_by: Mapped["User"] = relationship('User', foreign_keys=[activated_by_id])
+    user_id: Mapped[int] = mapped_column(ForeignKey('users.id'))
+    user: Mapped["User"] = relationship('User', foreign_keys=[user_id])
 
     created_at: Mapped[datetime.datetime] = mapped_column(nullable=False)
     created_by_id: Mapped[int] = mapped_column(ForeignKey('users.id'))
diff --git a/routers/user.py b/routers/user.py
index 8105c20..2994447 100644
--- a/routers/user.py
+++ b/routers/user.py
@@ -92,6 +92,7 @@ async def upload_passport_image(
 )
 async def generate_invite_code(
         session: SessionDependency,
+        request:GenerateInviteCodeRequest,
         user: CurrentUserDependency
 ):
-    return await UserService(session).generate_invite_code(user)
+    return await UserService(session).generate_invite_code(user, request)
diff --git a/schemas/user.py b/schemas/user.py
index 6ee38ae..1350702 100644
--- a/schemas/user.py
+++ b/schemas/user.py
@@ -82,6 +82,10 @@ class CreateUserRequest(BaseSchema):
     data: UserCreate
 
 
+class GenerateInviteCodeRequest(BaseSchema):
+    user_id: int
+
+
 # endregion
 
 # region Responses
diff --git a/services/auth.py b/services/auth.py
index db0a925..cd5e602 100644
--- a/services/auth.py
+++ b/services/auth.py
@@ -38,6 +38,8 @@ async def get_current_user(
         user = await session.get(User, user_id)
         if not user:
             raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Invalid credentials')
+        if user.is_deleted or user.is_blocked:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Пользователь заблокирован или удален')
         return user
     except JWTError as e:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Invalid token')
@@ -82,15 +84,12 @@ class AuthService(BaseService):
             delta = datetime.now() - invite_code.created_at
             if delta.seconds >= constants.INVITE_CODE_EXPIRY:
                 raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Код приглашения устарел')
-            user = User(
-                telegram_id=request.id,
-                is_admin=False,
-                role_key=UserRole.user
-            )
-            self.session.add(user)
+            user = await self.session.get(User, invite_code.user_id)
+            if not user:
+                raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Пользователь не найден')
+            user.telegram_id = request.id
             await self.session.flush()
             invite_code.is_activated = True
-            invite_code.activated_by_id = user.id
             await self.session.commit()
         if not user:
             raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Пользователь не найден')
diff --git a/services/user.py b/services/user.py
index d43f02c..a06d6d5 100644
--- a/services/user.py
+++ b/services/user.py
@@ -175,7 +175,7 @@ class UserService(BaseService):
         random.shuffle(code)
         return ''.join(code)
 
-    async def generate_invite_code(self, user: User) -> GenerateInviteCodeResponse:
+    async def generate_invite_code(self, user: User, request: GenerateInviteCodeRequest) -> GenerateInviteCodeResponse:
         MAX_ATTEMPTS = 5
         try:
             if not user.is_admin:
@@ -195,7 +195,8 @@ class UserService(BaseService):
             new_invite_code = InviteCode(
                 code=invite_code,
                 created_at=datetime.datetime.now(),
-                created_by_id=user.id
+                created_by_id=user.id,
+                user_id=request.user_id
             )
             self.session.add(new_invite_code)
             await self.session.commit()