feat(api,frontend): add qr code login
This commit is contained in:
@@ -1,16 +1,19 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import secrets
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from pyrogram.errors import SessionPasswordNeeded
|
||||
from pyrogram.errors import AuthTokenExpired, SessionPasswordNeeded
|
||||
from pyrogram.qrlogin import QRLogin
|
||||
from pyrogram.types import User
|
||||
|
||||
from userbot import PyroClient
|
||||
from utils.env import env
|
||||
|
||||
LOGIN_TTL_SECONDS = 900
|
||||
QR_POLL_SECONDS = 25
|
||||
PENDING_DIRNAME = "pending"
|
||||
|
||||
|
||||
@@ -18,12 +21,27 @@ 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
|
||||
phone: str
|
||||
phone_code_hash: str
|
||||
started_at: float
|
||||
phone: str = ""
|
||||
phone_code_hash: str = ""
|
||||
qr: QrState | None = None
|
||||
watcher: asyncio.Task[None] | None = None
|
||||
|
||||
|
||||
def _sessions_dir() -> Path:
|
||||
@@ -38,6 +56,29 @@ def _pending_dir() -> Path:
|
||||
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] = {}
|
||||
@@ -49,21 +90,53 @@ class LoginManager:
|
||||
raise LoginError(msg)
|
||||
return login
|
||||
|
||||
async def start(self, phone: str) -> str:
|
||||
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, phone, sent.phone_code_hash, time.monotonic()
|
||||
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:
|
||||
@@ -81,6 +154,7 @@ class LoginManager:
|
||||
|
||||
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")
|
||||
@@ -88,6 +162,7 @@ class LoginManager:
|
||||
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:
|
||||
|
||||
@@ -37,7 +37,8 @@ class PasswordRequest(BaseModel):
|
||||
|
||||
class LoginState(BaseModel):
|
||||
login_id: str
|
||||
stage: Literal["code", "password", "done"]
|
||||
stage: Literal["code", "qr", "password", "done"]
|
||||
qr_url: str | None = None
|
||||
account: AccountView | None = None
|
||||
|
||||
|
||||
@@ -79,6 +80,24 @@ async def start_login(body: PhoneRequest) -> LoginState:
|
||||
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]
|
||||
|
||||
@@ -17,7 +17,7 @@ class PyroClient(Client):
|
||||
api_hash="b18441a1ff607e10a989891a5462e627",
|
||||
device_model="Desktop",
|
||||
system_version="Windows 11 x64",
|
||||
app_version="6.2.4 x64",
|
||||
app_version="6.7.8 x64",
|
||||
lang_pack="tdesktop",
|
||||
client_platform=enums.ClientPlatform.DESKTOP,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user