151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
from collections.abc import Coroutine
|
||
from typing import Any, Literal
|
||
|
||
import asyncpg
|
||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||
from fastapi import APIRouter, HTTPException, status
|
||
from pydantic import BaseModel
|
||
from pyrogram.errors import FloodWait, RPCError
|
||
from pyrogram.types import User
|
||
|
||
from api.login import LoginError, login_manager
|
||
from userbot import DEVICE_MODEL_LIMIT
|
||
from utils.read import accounts
|
||
from utils.read.models import AccountView
|
||
|
||
router = APIRouter(prefix="/api", tags=["accounts"], route_class=DishkaRoute)
|
||
|
||
_ERROR_MESSAGES = {
|
||
"PhoneNumberInvalid": "Неверный номер телефона",
|
||
"PhoneNumberBanned": "Номер заблокирован в Telegram",
|
||
"PhoneCodeInvalid": "Неверный код",
|
||
"PhoneCodeExpired": "Код истёк, запросите новый",
|
||
"PasswordHashInvalid": "Неверный пароль",
|
||
}
|
||
|
||
|
||
class PhoneRequest(BaseModel):
|
||
phone: str
|
||
|
||
|
||
class CodeRequest(BaseModel):
|
||
code: str
|
||
|
||
|
||
class PasswordRequest(BaseModel):
|
||
password: str
|
||
|
||
|
||
class DeviceModelRequest(BaseModel):
|
||
device_model: str
|
||
|
||
|
||
class LoginState(BaseModel):
|
||
login_id: str
|
||
stage: Literal["code", "qr", "password", "done"]
|
||
qr_url: str | None = None
|
||
account: AccountView | None = None
|
||
|
||
|
||
async def _guard[T](coro: Coroutine[Any, Any, T]) -> T:
|
||
try:
|
||
return await coro
|
||
except LoginError as exc:
|
||
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
|
||
except FloodWait as exc:
|
||
raise HTTPException(
|
||
status.HTTP_429_TOO_MANY_REQUESTS,
|
||
f"Слишком много попыток, подождите {exc.value} с",
|
||
) from exc
|
||
except RPCError as exc:
|
||
detail = _ERROR_MESSAGES.get(type(exc).__name__, str(exc))
|
||
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail) from exc
|
||
|
||
|
||
async def _complete(pool: asyncpg.Pool, login_id: str, user: User) -> LoginState:
|
||
session_name = str(user.id)
|
||
account_id = await accounts.sync_account(pool, user, session_name)
|
||
await login_manager.finalize(login_id, session_name)
|
||
await accounts.notify_accounts_changed(pool)
|
||
return LoginState(
|
||
login_id=login_id,
|
||
stage="done",
|
||
account=await accounts.get_account(pool, account_id),
|
||
)
|
||
|
||
|
||
@router.get("/accounts")
|
||
async def list_accounts(pool: FromDishka[asyncpg.Pool]) -> list[AccountView]:
|
||
return await accounts.list_accounts(pool)
|
||
|
||
|
||
@router.post("/accounts/login")
|
||
async def start_login(body: PhoneRequest) -> LoginState:
|
||
login_id = await _guard(login_manager.start(body.phone))
|
||
return LoginState(login_id=login_id, stage="code")
|
||
|
||
|
||
@router.post("/accounts/login/qr")
|
||
async def start_qr_login() -> LoginState:
|
||
login_id, url = await _guard(login_manager.start_qr())
|
||
return LoginState(login_id=login_id, stage="qr", qr_url=url)
|
||
|
||
|
||
@router.get("/accounts/login/{login_id}/qr")
|
||
async def poll_qr_login(login_id: str, pool: FromDishka[asyncpg.Pool]) -> LoginState:
|
||
state = await _guard(login_manager.wait_qr(login_id))
|
||
if state.user is not None:
|
||
return await _complete(pool, login_id, state.user)
|
||
if state.password_needed:
|
||
return LoginState(login_id=login_id, stage="password")
|
||
if state.error is not None:
|
||
raise HTTPException(status.HTTP_400_BAD_REQUEST, state.error)
|
||
return LoginState(login_id=login_id, stage="qr", qr_url=state.url)
|
||
|
||
|
||
@router.post("/accounts/login/{login_id}/code")
|
||
async def submit_code(
|
||
login_id: str, body: CodeRequest, pool: FromDishka[asyncpg.Pool]
|
||
) -> LoginState:
|
||
user = await _guard(login_manager.submit_code(login_id, body.code))
|
||
if user is None:
|
||
return LoginState(login_id=login_id, stage="password")
|
||
return await _complete(pool, login_id, user)
|
||
|
||
|
||
@router.post("/accounts/login/{login_id}/password")
|
||
async def submit_password(
|
||
login_id: str, body: PasswordRequest, pool: FromDishka[asyncpg.Pool]
|
||
) -> LoginState:
|
||
user = await _guard(login_manager.submit_password(login_id, body.password))
|
||
return await _complete(pool, login_id, user)
|
||
|
||
|
||
@router.delete("/accounts/login/{login_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||
async def cancel_login(login_id: str) -> None:
|
||
await login_manager.cancel(login_id)
|
||
|
||
|
||
@router.patch("/accounts/{account_id}/device")
|
||
async def rename_device(
|
||
account_id: int, body: DeviceModelRequest, pool: FromDishka[asyncpg.Pool]
|
||
) -> AccountView:
|
||
device_model = " ".join(body.device_model.split())
|
||
if not device_model or len(device_model) > DEVICE_MODEL_LIMIT:
|
||
raise HTTPException(
|
||
status.HTTP_400_BAD_REQUEST,
|
||
f"Имя устройства: от 1 до {DEVICE_MODEL_LIMIT} символов",
|
||
)
|
||
account = await accounts.set_device_model(pool, account_id, device_model)
|
||
if account is None:
|
||
raise HTTPException(status.HTTP_404_NOT_FOUND, "Аккаунт не найден")
|
||
await accounts.notify_accounts_changed(pool)
|
||
return account
|
||
|
||
|
||
@router.delete("/accounts/{account_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||
async def logout_account(account_id: int, pool: FromDishka[asyncpg.Pool]) -> None:
|
||
if await accounts.deactivate_account(pool, account_id) is None:
|
||
raise HTTPException(status.HTTP_404_NOT_FOUND, "Аккаунт не найден")
|
||
await accounts.notify_accounts_changed(pool)
|