feat(userbot): add Kubernetes control panel
Manage Telegram instances through Kubernetes with legacy adoption for forust and anna. Build and deploy the panel image alongside the runtime.
This commit is contained in:
@@ -0,0 +1,202 @@
|
||||
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) -> None:
|
||||
self.ttl = timedelta(seconds=ttl_seconds)
|
||||
self.flows: dict[str, AuthFlow] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def start_phone(self, account: PhoneStart) -> str:
|
||||
await self._cleanup_expired()
|
||||
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")
|
||||
Reference in New Issue
Block a user