homelab/userbot/panel/backend/app/auth_service.py

from __future__ import annotations

import asyncio
import secrets
from contextlib import suppress
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta

from pyrogram import Client
from pyrogram.errors import (
    ApiIdInvalid,
    FloodWait,
    PasswordHashInvalid,
    PhoneCodeExpired,
    PhoneCodeInvalid,
    PhoneNumberInvalid,
    RPCError,
    SessionPasswordNeeded,
    Unauthorized,
)

from .errors import PanelError
from .models import AccountBase, PhoneStart, StringSessionStart


@dataclass
class AuthFlow:
    flow_id: str
    account: PhoneStart
    client: Client
    phone_code_hash: str
    expires_at: datetime


@dataclass
class AuthorizedAccount:
    account: AccountBase
    session_string: str


class TelegramAuthService:
    def __init__(self, ttl_seconds: int = 600, max_flows: int = 50) -> None:
        self.ttl = timedelta(seconds=ttl_seconds)
        self.max_flows = max_flows
        self.flows: dict[str, AuthFlow] = {}
        self._lock = asyncio.Lock()

    async def start_phone(self, account: PhoneStart) -> str:
        await self._cleanup_expired()
        async with self._lock:
            if len(self.flows) >= self.max_flows:
                raise PanelError(
                    429,
                    "Too many pending authorization flows; try again later",
                )
        flow_id = secrets.token_urlsafe(24)
        telegram = Client(
            f"auth-{flow_id}",
            api_id=account.api_id,
            api_hash=account.api_hash,
            in_memory=True,
            no_updates=True,
        )
        try:
            await telegram.connect()
            sent = await telegram.send_code(account.phone)
        except Exception as exc:
            await self._disconnect(telegram)
            raise self._translate(exc) from exc

        flow = AuthFlow(
            flow_id=flow_id,
            account=account,
            client=telegram,
            phone_code_hash=sent.phone_code_hash,
            expires_at=datetime.now(UTC) + self.ttl,
        )
        async with self._lock:
            self.flows[flow_id] = flow
        return flow_id

    async def submit_code(
        self,
        flow_id: str,
        code: str,
    ) -> AuthorizedAccount | None:
        flow = await self._get(flow_id)
        try:
            await flow.client.sign_in(
                flow.account.phone,
                flow.phone_code_hash,
                code,
            )
        except SessionPasswordNeeded:
            return None
        except PhoneCodeExpired as exc:
            await self._discard(flow)
            raise self._translate(exc) from exc
        except Exception as exc:
            raise self._translate(exc) from exc
        return await self._finish(flow)

    async def submit_password(
        self,
        flow_id: str,
        password: str,
    ) -> AuthorizedAccount:
        flow = await self._get(flow_id)
        try:
            await flow.client.check_password(password)
        except Exception as exc:
            raise self._translate(exc) from exc
        return await self._finish(flow)

    async def validate_string_session(
        self,
        account: StringSessionStart,
    ) -> AuthorizedAccount:
        telegram = Client(
            f"validate-{secrets.token_urlsafe(12)}",
            api_id=account.api_id,
            api_hash=account.api_hash,
            session_string=account.session_string,
            in_memory=True,
            no_updates=True,
        )
        try:
            await telegram.connect()
            await telegram.get_me()
            session_string = await telegram.export_session_string()
        except Exception as exc:
            raise self._translate(exc, session=True) from exc
        finally:
            await self._disconnect(telegram)
        return AuthorizedAccount(account=account, session_string=session_string)

    async def close(self) -> None:
        async with self._lock:
            flows = list(self.flows.values())
            self.flows.clear()
        await asyncio.gather(
            *(self._disconnect(flow.client) for flow in flows),
            return_exceptions=True,
        )

    async def _get(self, flow_id: str) -> AuthFlow:
        await self._cleanup_expired()
        async with self._lock:
            flow = self.flows.get(flow_id)
        if flow is None:
            raise PanelError(410, "Authorization flow expired; start again")
        return flow

    async def _finish(self, flow: AuthFlow) -> AuthorizedAccount:
        try:
            await flow.client.get_me()
            session_string = await flow.client.export_session_string()
        except Exception as exc:
            raise self._translate(exc) from exc
        finally:
            async with self._lock:
                self.flows.pop(flow.flow_id, None)
            await self._disconnect(flow.client)
        return AuthorizedAccount(account=flow.account, session_string=session_string)

    async def _discard(self, flow: AuthFlow) -> None:
        async with self._lock:
            self.flows.pop(flow.flow_id, None)
        await self._disconnect(flow.client)

    async def _cleanup_expired(self) -> None:
        now = datetime.now(UTC)
        async with self._lock:
            expired = [
                self.flows.pop(flow_id)
                for flow_id, flow in list(self.flows.items())
                if flow.expires_at <= now
            ]
        if expired:
            await asyncio.gather(
                *(self._disconnect(flow.client) for flow in expired),
                return_exceptions=True,
            )

    @staticmethod
    async def _disconnect(telegram: Client) -> None:
        with suppress(Exception):
            if telegram.is_connected:
                await telegram.disconnect()

    @staticmethod
    def _translate(exc: Exception, *, session: bool = False) -> PanelError:
        if isinstance(exc, FloodWait):
            return PanelError(429, f"Telegram rate limit; retry in {exc.value} seconds")
        if isinstance(exc, ApiIdInvalid):
            return PanelError(422, "Telegram API ID or API Hash is invalid")
        if isinstance(exc, PhoneNumberInvalid):
            return PanelError(422, "Phone number is invalid")
        if isinstance(exc, PhoneCodeInvalid):
            return PanelError(422, "Telegram code is invalid")
        if isinstance(exc, PhoneCodeExpired):
            return PanelError(410, "Telegram code expired; start again")
        if isinstance(exc, PasswordHashInvalid):
            return PanelError(422, "2FA password is invalid")
        if session and isinstance(exc, (Unauthorized, RPCError)):
            return PanelError(422, "StringSession is invalid or expired")
        if isinstance(exc, RPCError):
            return PanelError(422, "Telegram rejected the authorization request")
        return PanelError(503, "Telegram authorization is unavailable")