Files
ai-char-bot/repos/user_repo.py

121 lines
4.3 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
from datetime import datetime, timedelta
from enum import Enum
from typing import Optional
from aiogram.types import User
from motor.motor_asyncio import AsyncIOMotorClient
from utils.security import get_password_hash
class UserStatus:
ALLOWED = "allowed"
DENIED = "denied"
PENDING = "pending"
NONE = "none" # Пользователя нет в базе
class UsersRepo:
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
self.collection = client[db_name]["users"]
async def get_user(self, user_id: int):
user = await self.collection.find_one({"user_id": user_id})
if user:
user["id"] = str(user["_id"])
return user
async def get_user_by_username(self, username: str):
user = await self.collection.find_one({"username": username})
if user:
user["id"] = str(user["_id"])
return user
async def create_user(self, username: str, password: str, full_name: Optional[str] = None):
"""Создает нового пользователя с username/паролем"""
existing = await self.get_user_by_username(username)
if existing:
raise ValueError("User with this username already exists")
user_doc = {
"username": username,
"hashed_password": get_password_hash(password),
"full_name": full_name,
"status": UserStatus.PENDING, # По умолчанию PENDING
"created_at": datetime.now(),
"is_email_user": False, # Теперь это просто "обычный" юзер, не телеграм (хотя поле можно переименовать)
"is_web_user": True,
"is_admin": False,
"project_ids": [],
"current_project_id": None
}
result = await self.collection.insert_one(user_doc)
user = await self.collection.find_one({"_id": result.inserted_id})
if user:
user["id"] = str(user["_id"])
return user
async def get_pending_users(self):
"""Возвращает список пользователей со статусом PENDING"""
cursor = self.collection.find({"status": UserStatus.PENDING})
users = await cursor.to_list(length=100)
for user in users:
user["id"] = str(user["_id"])
return users
async def approve_user(self, username: str):
await self.collection.update_one(
{"username": username},
{"$set": {"status": UserStatus.ALLOWED}}
)
async def deny_user(self, username: str):
await self.collection.update_one(
{"username": username},
{"$set": {"status": UserStatus.DENIED}}
)
async def create_or_update_request(self, user: User):
"""
Обновляет дату последнего запроса и ставит статус PENDING.
Сохраняет всю инфу о юзере (для Telegram пользователей).
"""
now = datetime.now()
data = {
"user_id": user.id,
"username": user.username,
"full_name": user.full_name,
"status": UserStatus.PENDING,
"last_request_date": now,
"is_email_user": False
}
await self.collection.update_one(
{"user_id": user.id},
{"$set": data},
upsert=True
)
async def set_status(self, user_id: int, status: str):
"""Меняет статус (разрешен/запрещен)"""
await self.collection.update_one(
{"user_id": user_id},
{"$set": {"status": status}}
)
async def can_request_access(self, user_id: int) -> bool:
"""
Проверяет, можно ли отправить запрос (прошло ли 24 часа).
Возвращает True, если пользователя нет или прошло > 24ч.
"""
user = await self.get_user(user_id)
if not user:
return True
last_date = user.get("last_request_date")
if not last_date:
return True
# Проверка на 24 часа
if datetime.now() - last_date > timedelta(hours=24):
return True
return False