feat: invite code

constants.py

@@ -25,8 +25,8 @@ allowed_telegram_ids = [
     502869937,  # Sasha
     1006239222,  # Sasha Serbin
     7326211785,
-    6427522679, # Ekaterina Manager
+    6427522679,  # Ekaterina Manager
-    5734685107 #Kristina v2
+    5734685107  # Kristina v2
 ]
 
 MONTHS = (
@@ -53,3 +53,6 @@ DEAL_STATUS_STR = [
     'Завершено',
     'Отменено',
 ]
+
+# 30 minutes in seconds
+INVITE_CODE_EXPIRY = 30 * 60

models/auth.py

@@ -1,3 +1,4 @@
+import datetime
 from typing import TYPE_CHECKING, Optional
 
 from sqlalchemy import BigInteger, Table, ForeignKey, Column
@@ -32,7 +33,6 @@ user_pay_rate = Table(
 )
 
 
-
 class Permission(BaseModel):
     __tablename__ = 'permissions'
     key: Mapped[str] = mapped_column(primary_key=True)
@@ -83,7 +83,7 @@ class User(BaseModel):
     is_blocked: Mapped[bool] = mapped_column(nullable=False, server_default='0')
     is_deleted: Mapped[bool] = mapped_column(nullable=False, server_default='0')
 
-    role_key: Mapped[int] = mapped_column(ForeignKey('roles.key'), server_default=UserRole.user)
+    role_key: Mapped[str] = mapped_column(ForeignKey('roles.key'), server_default=UserRole.user)
     role: Mapped["Role"] = relationship(
         'Role',
         lazy='joined'
@@ -139,6 +139,19 @@ class User(BaseModel):
     )
 
 
+class InviteCode(BaseModel):
+    __tablename__ = 'invite_codes'
+    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])
+
+    created_at: Mapped[datetime.datetime] = mapped_column(nullable=False)
+    created_by_id: Mapped[int] = mapped_column(ForeignKey('users.id'))
+    created_by: Mapped["User"] = relationship('User', foreign_keys=[created_by_id])
+
+
 class Position(BaseModel):
     __tablename__ = 'positions'
     key: Mapped[str] = mapped_column(primary_key=True)

routers/user.py

@@ -1,6 +1,6 @@
 from fastapi import APIRouter, Depends, UploadFile
 
-from backend.dependecies import SessionDependency
+from backend.dependecies import SessionDependency, CurrentUserDependency
 from schemas.user import *
 from services.auth import authorized_user
 from services.user import UserService
@@ -83,3 +83,15 @@ async def upload_passport_image(
 ):
     file_bytes = upload_file.file.read()
     return await UserService(session).upload_passport_image(user_id, file_bytes)
+
+
+@user_router.post(
+    '/generate-invite-code',
+    response_model=GenerateInviteCodeResponse,
+    operation_id='generate_invite_code'
+)
+async def generate_invite_code(
+        session: SessionDependency,
+        user: CurrentUserDependency
+):
+    return await UserService(session).generate_invite_code(user)

schemas/auth.py

@@ -9,6 +9,7 @@ class AuthLoginRequest(CustomModelSnake):
     hash: str
     id: int
     photo_url: Optional[str] = None
+    invite_code: Optional[str] = None
 
 
 class AuthLoginResponse(BaseSchema):

schemas/user.py

@@ -108,4 +108,7 @@ class GetManagersResponse(BaseSchema):
 class UploadPassportImageResponse(OkMessageSchema):
     image_url: str | None = None
 
+
+class GenerateInviteCodeResponse(OkMessageSchema):
+    invite_code: Optional[str] = None
 # endregion

services/auth.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from typing import Union, Annotated
 
 from fastapi import Depends, HTTPException
@@ -11,7 +12,7 @@ import backend.config
 import constants
 from backend.session import get_session
 from enums.user import UserRole
-from models import User
+from models import User, InviteCode
 from schemas.auth import *
 from services.base import BaseService
 
@@ -62,20 +63,37 @@ class AuthService(BaseService):
         return jwt.encode(payload, backend.config.SECRET_KEY, algorithm=algorithm)
 
     async def authenticate(self, request: AuthLoginRequest):
-        if request.id not in constants.allowed_telegram_ids:
+        user: Optional[User] = await self.session.scalar(select(User).where(User.telegram_id == request.id))
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Invalid credentials')
-
-        user: Union[User, None] = await self.session.scalar(select(User).where(User.telegram_id == request.id))
         if user and (user.is_deleted or user.is_blocked):
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Invalid credentials')
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Пользователь заблокирован или удален')
-        if not user:
+        if not user and request.invite_code:
+            invite_code = await self.session.scalar(
+                select(
+                    InviteCode
+                )
+                .where(
+                    InviteCode.code == request.invite_code,
+                    InviteCode.is_activated == False
+                )
+            )
+            if not invite_code:
+                raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Неверный код приглашения')
+            # check if code is expired
+            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)
+            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='Пользователь не найден')
         payload = {
             'sub': str(user.id),
             'role': user.role_key,

services/user.py

@@ -1,9 +1,15 @@
+import datetime
+import random
+import string
+import time
+
 from sqlalchemy import select, update, delete, insert, and_
 from sqlalchemy.orm import selectinload
 
 from backend import config
 from external.s3_uploader.uploader import S3Uploader
-from models import User, user_position, user_pay_rate, PassportImage, DepartmentSection, UserDepartmentSection
+from models import User, user_position, user_pay_rate, PassportImage, DepartmentSection, UserDepartmentSection, \
+    InviteCode
 from services.base import BaseService
 from schemas.user import *
 
@@ -159,3 +165,41 @@ class UserService(BaseService):
             )
         except Exception as e:
             return UploadPassportImageResponse(ok=False, message=str(e))
+
+    @staticmethod
+    def _generate_invite_code(length=10):
+        timestamp = str(int(time.time() * 1000))[-6:]
+        chars = string.ascii_letters + string.digits
+        random_part = ''.join(random.choice(chars) for _ in range(length - len(timestamp))).upper()
+        code = list(timestamp + random_part)
+        random.shuffle(code)
+        return ''.join(code)
+
+    async def generate_invite_code(self, user: User) -> GenerateInviteCodeResponse:
+        MAX_ATTEMPTS = 5
+        try:
+            if not user.is_admin:
+                return GenerateInviteCodeResponse(ok=False,
+                                                  message="Сгенерировать код приглашения может только администратор")
+            code_in_database = True
+            attempt = 0
+            invite_code = ""
+            while code_in_database and attempt < MAX_ATTEMPTS:
+                invite_code = self._generate_invite_code()
+                stmt = select(InviteCode).where(InviteCode.code == invite_code, InviteCode.is_activated == False)
+                user_with_code = await self.session.scalar(stmt)
+                code_in_database = bool(user_with_code)
+                attempt += 1
+            if code_in_database or not invite_code:
+                return GenerateInviteCodeResponse(ok=False, message="Не удалось сгенерировать уникальный код")
+            new_invite_code = InviteCode(
+                code=invite_code,
+                created_at=datetime.datetime.now(),
+                created_by_id=user.id
+            )
+            self.session.add(new_invite_code)
+            await self.session.commit()
+            return GenerateInviteCodeResponse(ok=True, message="Код приглашения успешно создан",
+                                              invite_code=invite_code)
+        except Exception as e:
+            return GenerateInviteCodeResponse(ok=False, message=str(e))