feat: invite code

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:
-            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))
+        user: Optional[User] = 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')
-        if not user:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Пользователь заблокирован или удален')
+        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))