feat: invite code

This commit is contained in:
2025-03-05 18:39:54 +03:00
parent c75ff66f27
commit 6c2698679d
7 changed files with 107 additions and 13 deletions

View File

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

View File

@@ -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)

View File

@@ -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)

View File

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

View File

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

View File

@@ -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,

View File

@@ -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))