diff --git a/constants.py b/constants.py index ca00049..f1f3da7 100644 --- a/constants.py +++ b/constants.py @@ -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 diff --git a/models/auth.py b/models/auth.py index 7f2b735..ea018c4 100644 --- a/models/auth.py +++ b/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) diff --git a/routers/user.py b/routers/user.py index afec672..8105c20 100644 --- a/routers/user.py +++ b/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) diff --git a/schemas/auth.py b/schemas/auth.py index 0badf3a..91e12ef 100644 --- a/schemas/auth.py +++ b/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): diff --git a/schemas/user.py b/schemas/user.py index 87b4ec9..6ee38ae 100644 --- a/schemas/user.py +++ b/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 diff --git a/services/auth.py b/services/auth.py index fab325c..db0a925 100644 --- a/services/auth.py +++ b/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, diff --git a/services/user.py b/services/user.py index e594ed7..d43f02c 100644 --- a/services/user.py +++ b/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))