beavergram/backend/src/api/login.py

import asyncio
import contextlib
import secrets
import time
from dataclasses import dataclass, field
from pathlib import Path

from pyrogram.errors import AuthTokenExpired, SessionPasswordNeeded
from pyrogram.qrlogin import QRLogin
from pyrogram.types import User

from userbot.modules.client import PyroClient
from utils.env import env

LOGIN_TTL_SECONDS = 900
QR_POLL_SECONDS = 25
PENDING_DIRNAME = "pending"


class LoginError(Exception):
    pass


@dataclass
class QrState:
    url: str
    user: User | None = None
    password_needed: bool = False
    error: str | None = None
    changed: asyncio.Event = field(default_factory=asyncio.Event)

    @property
    def settled(self) -> bool:
        return self.user is not None or self.password_needed or self.error is not None


@dataclass
class PendingLogin:
    client: PyroClient
    started_at: float
    phone: str = ""
    phone_code_hash: str = ""
    qr: QrState | None = None
    watcher: asyncio.Task[None] | None = None


def _sessions_dir() -> Path:
    path = Path(env.tg.sessions_dir)
    path.mkdir(parents=True, exist_ok=True)
    return path


def _pending_dir() -> Path:
    path = _sessions_dir() / PENDING_DIRNAME
    path.mkdir(parents=True, exist_ok=True)
    return path


async def _stop_watcher(login: PendingLogin) -> None:
    if login.watcher is None or login.watcher.done():
        return
    login.watcher.cancel()
    with contextlib.suppress(asyncio.CancelledError, Exception):
        await login.watcher


async def _watch_qr(qr: QRLogin, state: QrState) -> None:
    while not state.settled:
        try:
            try:
                state.user = await qr.wait()
            except (TimeoutError, AuthTokenExpired):
                await qr.recreate()
                state.url = qr.url
        except SessionPasswordNeeded:
            state.password_needed = True
        except Exception as exc:
            state.error = str(exc)
        state.changed.set()


class LoginManager:
    def __init__(self) -> None:
        self._logins: dict[str, PendingLogin] = {}

    def _get(self, login_id: str) -> PendingLogin:
        login = self._logins.get(login_id)
        if login is None:
            msg = "Сессия входа истекла, начните заново"
            raise LoginError(msg)
        return login

    async def _connect(self) -> tuple[str, PyroClient]:
        await self._sweep()
        login_id = secrets.token_hex(8)
        client = PyroClient(login_id, workdir=str(_pending_dir()), load_handlers=False)
        await client.connect()
        return login_id, client

    async def start(self, phone: str) -> str:
        login_id, client = await self._connect()
        try:
            sent = await client.send_code(phone)
        except Exception:
            await self._discard(login_id, client)
            raise
        self._logins[login_id] = PendingLogin(
            client, time.monotonic(), phone=phone, phone_code_hash=sent.phone_code_hash
        )
        return login_id

    async def start_qr(self) -> tuple[str, str]:
        login_id, client = await self._connect()
        qr = QRLogin(client)
        try:
            await qr.recreate()
        except Exception:
            await self._discard(login_id, client)
            raise
        state = QrState(qr.url)
        self._logins[login_id] = PendingLogin(
            client,
            time.monotonic(),
            qr=state,
            watcher=asyncio.create_task(_watch_qr(qr, state)),
        )
        return login_id, state.url

    async def wait_qr(self, login_id: str) -> QrState:
        login = self._get(login_id)
        if login.qr is None:
            msg = "Этот вход начат по номеру телефона"
            raise LoginError(msg)
        with contextlib.suppress(TimeoutError):
            async with asyncio.timeout(QR_POLL_SECONDS):
                await login.qr.changed.wait()
        login.qr.changed.clear()
        return login.qr

    async def submit_code(self, login_id: str, code: str) -> User | None:
        login = self._get(login_id)
        try:
            user = await login.client.sign_in(login.phone, login.phone_code_hash, code)
        except SessionPasswordNeeded:
            return None
        if not isinstance(user, User):
            await self.cancel(login_id)
            msg = "Этот номер не зарегистрирован в Telegram"
            raise LoginError(msg)
        return user

    async def submit_password(self, login_id: str, password: str) -> User:
        return await self._get(login_id).client.check_password(password)

    async def finalize(self, login_id: str, session_name: str) -> None:
        login = self._logins.pop(login_id)
        await _stop_watcher(login)
        await login.client.disconnect()
        source = _pending_dir() / f"{login_id}.session"
        source.replace(_sessions_dir() / f"{session_name}.session")

    async def cancel(self, login_id: str) -> None:
        login = self._logins.pop(login_id, None)
        if login is not None:
            await _stop_watcher(login)
            await self._discard(login_id, login.client)

    async def _discard(self, login_id: str, client: PyroClient) -> None:
        with contextlib.suppress(Exception):
            await client.disconnect()
        (_pending_dir() / f"{login_id}.session").unlink(missing_ok=True)

    async def _sweep(self) -> None:
        now = time.monotonic()
        for login_id, login in list(self._logins.items()):
            if now - login.started_at > LOGIN_TTL_SECONDS:
                await self.cancel(login_id)
        for path in _pending_dir().glob("*.session"):
            if path.stem not in self._logins:
                path.unlink(missing_ok=True)


login_manager = LoginManager()