From 5e35c47b61dd60c99c84d8e2de15294b8973a647 Mon Sep 17 00:00:00 2001 From: wulan17 Date: Thu, 15 May 2025 23:16:52 +0700 Subject: [PATCH] pyrofork: Refactor Qr Code Signin Signed-off-by: wulan17 --- pyrogram/client.py | 6 +-- pyrogram/methods/auth/sign_in_qrcode.py | 48 +++++---------------- pyrogram/types/authorization/login_token.py | 9 ++++ 3 files changed, 22 insertions(+), 41 deletions(-) diff --git a/pyrogram/client.py b/pyrogram/client.py index e800969d..b467042a 100644 --- a/pyrogram/client.py +++ b/pyrogram/client.py @@ -39,6 +39,7 @@ import pyrogram from pyrogram import __version__, __license__ from pyrogram import enums from pyrogram import raw +from pyrogram import types from pyrogram import utils from pyrogram.crypto import aes from pyrogram.errors import CDNFileHashMismatch @@ -510,15 +511,14 @@ class Client(Methods): print(e.MESSAGE) self.password = None else: - if self.use_qrcode and isinstance(signed_in, raw.types.auth.LoginToken): + if self.use_qrcode and isinstance(signed_in, types.LoginToken): time_out = signed_in.expires - datetime.timestamp(datetime.now()) try: await asyncio.wait_for(self._wait_for_update_login_token(), timeout=time_out) except asyncio.TimeoutError: print("QR code expired, Requesting new Qr code...") continue - else: - break + break if isinstance(signed_in, User): return signed_in diff --git a/pyrogram/methods/auth/sign_in_qrcode.py b/pyrogram/methods/auth/sign_in_qrcode.py index ffca3e80..0b255788 100644 --- a/pyrogram/methods/auth/sign_in_qrcode.py +++ b/pyrogram/methods/auth/sign_in_qrcode.py @@ -28,14 +28,6 @@ from pyrogram.session import Session, Auth log = logging.getLogger(__name__) -QRCODE_AVAIL = False -try: - import qrcode - QRCODE_AVAIL = True -except ImportError: - QRCODE_AVAIL = False - - class SignInQrcode: async def sign_in_qrcode( self: "pyrogram.Client" @@ -54,7 +46,9 @@ class SignInQrcode: SessionPasswordNeeded: In case a password is needed to sign in. """ - if not QRCODE_AVAIL: + try: + import qrcode + except ImportError: raise ImportError("qrcode is missing! " "Please install it with `pip install qrcode`") r = await self.session.invoke( @@ -79,13 +73,13 @@ class SignInQrcode: print("Scan the QR code with your Telegram app.") qr.print_ascii() - return r - elif isinstance(r, raw.types.auth.LoginTokenSuccess): + return types.LoginToken._parse(r) + if isinstance(r, raw.types.auth.LoginTokenSuccess): await self.storage.user_id(r.authorization.user.id) await self.storage.is_bot(False) return types.User._parse(self, r.authorization.user) - elif isinstance(r, raw.types.auth.LoginTokenMigrateTo): + if isinstance(r, raw.types.auth.LoginTokenMigrateTo): # pylint: disable=access-member-before-definition await self.session.stop() @@ -107,33 +101,11 @@ class SignInQrcode: token=r.token ) ) - if isinstance(r, raw.types.auth.LoginToken): - base64_token = b64encode(r.token).decode("utf-8") - login_url = f"tg://login?token={base64_token}" - qr = qrcode.QRCode( - version=1, - error_correction=qrcode.constants.ERROR_CORRECT_L, - box_size=10, - border=4, - ) - qr.add_data(login_url) - qr.make(fit=True) - - print("Scan the QR code below with your Telegram app.") - qr.print_ascii() - - return types.LoginToken( - self, - r.token, - r.expires - ) - elif isinstance(r, raw.types.auth.LoginTokenSuccess): + if isinstance(r, raw.types.auth.LoginTokenSuccess): await self.storage.user_id(r.authorization.user.id) await self.storage.is_bot(False) return types.User._parse(self, r.authorization.user) - else: - raise pyrogram.exceptions.RPCError( - "Unknown response type from Telegram API" - ) - return r + raise pyrogram.exceptions.RPCError( + "Unknown response type from Telegram API" + ) diff --git a/pyrogram/types/authorization/login_token.py b/pyrogram/types/authorization/login_token.py index 6b3473a9..a869f539 100644 --- a/pyrogram/types/authorization/login_token.py +++ b/pyrogram/types/authorization/login_token.py @@ -18,6 +18,8 @@ from ..object import Object +from pyrogram import raw + class LoginToken(Object): """Contains info on a login token. @@ -35,3 +37,10 @@ class LoginToken(Object): self.token = token self.expires = expires + + @staticmethod + def _parse(login_token: "raw.base.LoginToken") -> "LoginToken": + return LoginToken( + token=login_token.token, + expires=login_token.expires, + )