From f864f66e624133470bdcbed761d3c1c0880eeac8 Mon Sep 17 00:00:00 2001 From: Max Sandholm Date: Thu, 26 Jan 2023 23:43:44 +0200 Subject: [PATCH 1/3] Add websocket for QR login to provisioning API --- mautrix_telegram/web/provisioning/__init__.py | 48 ++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/mautrix_telegram/web/provisioning/__init__.py b/mautrix_telegram/web/provisioning/__init__.py index 21759ccf..c52101d6 100644 --- a/mautrix_telegram/web/provisioning/__init__.py +++ b/mautrix_telegram/web/provisioning/__init__.py @@ -21,6 +21,8 @@ import json import logging from aiohttp import web +from telethon.errors import SessionPasswordNeededError +from telethon.tl.custom import QRLogin from telethon.tl.functions.messages import GetAllStickersRequest from telethon.tl.types import ChannelForbidden, ChatForbidden, TypeChat, User as TLUser from telethon.utils import get_peer_id, resolve_id @@ -78,6 +80,7 @@ class ProvisioningAPI(AuthAPI): self.app.router.add_route("POST", f"{user_prefix}/retry_takeout", self.retry_takeout) self.app.router.add_route("POST", f"{user_prefix}/logout", self.logout) + self.app.router.add_route("GET", f"{user_prefix}/login/qr", self.login_qr) self.app.router.add_route("POST", f"{user_prefix}/login/bot_token", self.send_bot_token) self.app.router.add_route("POST", f"{user_prefix}/login/request_code", self.request_code) self.app.router.add_route("POST", f"{user_prefix}/login/send_code", self.send_code) @@ -528,6 +531,36 @@ class ProvisioningAPI(AuthAPI): user.takeout_retry_immediate.set() return web.json_response({}, status=200) + async def login_qr(self, request: web.Request) -> web.Response: + _, user, err = await self.get_user_request_info(request, websocket=True) + if err is not None: + return err + + await user.ensure_started(even_if_no_session=True) + qr_login = QRLogin(user.client, ignored_ids=[]) + + ws = web.WebSocketResponse(protocols=["net.maunium.telegram.login"]) + await ws.prepare(request) + + user_info = None + try: + await qr_login.recreate() + await ws.send_json({"code": qr_login.url, "timeout": 30}) + user_info = await qr_login.wait() + except asyncio.TimeoutError: + await ws.send_json({"success": False, "error": "timeout"}) + await ws.close() + return ws + except SessionPasswordNeededError: + await ws.send_json({"success": False, "error": "password-needed"}) + await ws.close() + return ws + + await self.postprocess_login(user, user_info) + await ws.send_json({"success": True}) + await ws.close() + return ws + async def send_bot_token(self, request: web.Request) -> web.Response: data, user, err = await self.get_user_request_info(request) if err is not None: @@ -653,6 +686,15 @@ class ProvisioningAPI(AuthAPI): ) return None + def check_websocket_authorization(self, request: web.Request) -> web.Response | None: + auth_parts = request.headers.get("Sec-WebSocket-Protocol").split(",") + for part in auth_parts: + if part.strip() == f"net.maunium.telegram.auth-{self.secret}": + return None + return self.get_error_response( + error="Shared secret is not valid.", errcode="shared_secret_invalid", status=401 + ) + @staticmethod async def get_data(request: web.Request) -> dict | None: try: @@ -707,8 +749,12 @@ class ProvisioningAPI(AuthAPI): expect_logged_in: bool | None = False, require_puppeting: bool = False, want_data: bool = True, + websocket: bool = False, ) -> tuple[dict | None, User | None, web.Response | None]: - err = self.check_authorization(request) + if not websocket: + err = self.check_authorization(request) + else: + err = self.check_websocket_authorization(request) if err is not None: return None, None, err From 67f75796faa5d961e3eae49606ae2996b8bfa7e6 Mon Sep 17 00:00:00 2001 From: Max Sandholm Date: Fri, 27 Jan 2023 17:37:48 +0200 Subject: [PATCH 2/3] Correct retry and timeout for QR websocket --- mautrix_telegram/web/provisioning/__init__.py | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/mautrix_telegram/web/provisioning/__init__.py b/mautrix_telegram/web/provisioning/__init__.py index c52101d6..293a23c8 100644 --- a/mautrix_telegram/web/provisioning/__init__.py +++ b/mautrix_telegram/web/provisioning/__init__.py @@ -14,6 +14,7 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . from __future__ import annotations +import datetime from typing import TYPE_CHECKING, Awaitable, Callable import asyncio @@ -542,19 +543,33 @@ class ProvisioningAPI(AuthAPI): ws = web.WebSocketResponse(protocols=["net.maunium.telegram.login"]) await ws.prepare(request) + retries = 0 user_info = None - try: - await qr_login.recreate() - await ws.send_json({"code": qr_login.url, "timeout": 30}) - user_info = await qr_login.wait() - except asyncio.TimeoutError: + while retries < 4: + try: + await qr_login.recreate() + await ws.send_json( + { + "code": qr_login.url, + "timeout": int( + ( + qr_login.expires - datetime.datetime.now(tz=datetime.timezone.utc) + ).total_seconds() + ), + } + ) + user_info = await qr_login.wait() + break + except asyncio.TimeoutError: + retries += 1 + except SessionPasswordNeededError: + await ws.send_json({"success": False, "error": "password-needed"}) + await ws.close() + return ws + else: await ws.send_json({"success": False, "error": "timeout"}) await ws.close() return ws - except SessionPasswordNeededError: - await ws.send_json({"success": False, "error": "password-needed"}) - await ws.close() - return ws await self.postprocess_login(user, user_info) await ws.send_json({"success": True}) From ad2b49928ae1a4395d8cd22a3565b50c032c494a Mon Sep 17 00:00:00 2001 From: Max Sandholm Date: Fri, 27 Jan 2023 17:40:12 +0200 Subject: [PATCH 3/3] Sort imports --- mautrix_telegram/web/provisioning/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mautrix_telegram/web/provisioning/__init__.py b/mautrix_telegram/web/provisioning/__init__.py index 293a23c8..7320b176 100644 --- a/mautrix_telegram/web/provisioning/__init__.py +++ b/mautrix_telegram/web/provisioning/__init__.py @@ -14,10 +14,10 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . from __future__ import annotations -import datetime from typing import TYPE_CHECKING, Awaitable, Callable import asyncio +import datetime import json import logging