Fix QRLogin

Author: KurimuzonAkuma

pyrogram/client.py

@@ -502,7 +502,7 @@ async def authorize(self) -> User:
                     print("Password hint: {}".format(await self.get_password_hint()))
 
                     if not self.password:
-                        self.password = await ainput("Enter password (empty to recover): ", hide=self.hide_password, loop=self.loop)
+                        self.password = await ainput("Enter 2FA password (empty to recover): ", hide=self.hide_password, loop=self.loop)
 
                     try:
                         if not self.password:
@@ -557,12 +557,30 @@ async def authorize(self) -> User:
 
         return signed_up
 
-    async def authorize_qr(self, except_ids: List[int] = []) -> User:
+    async def authorize_qr(self, except_ids: List[int] = []) -> "User":
         from qrcode import QRCode
+
         qr_login = QRLogin(self, except_ids)
+        await qr_login.recreate()
+
+        qr = QRCode(version=1)
+        signed_in = None
 
         while True:
             try:
+                print(
+                    "\x1b[2J\n"
+                    f"Welcome to Pyrogram (version {__version__})\n"
+                    "Pyrogram is free software and comes with ABSOLUTELY NO WARRANTY. Licensed\n"
+                    f"under the terms of the {__license__}.\n"
+                    "Scan the QR code below to login\n"
+                    "Settings -> Privacy and Security -> Active Sessions -> Scan QR Code.",
+                    flush=True
+                )
+
+                qr.clear()
+                qr.add_data(qr_login.url)
+                qr.print_ascii(tty=True)
                 log.info("Waiting for QR code being scanned.")
 
                 signed_in = await qr_login.wait()
@@ -573,21 +591,42 @@ async def authorize_qr(self, except_ids: List[int] = []) -> User:
             except asyncio.TimeoutError:
                 log.info("Recreating QR code.")
                 await qr_login.recreate()
-                print("\x1b[2J")
-                print(f"Welcome to Pyrogram (version {__version__})")
-                print(f"Pyrogram is free software and comes with ABSOLUTELY NO WARRANTY. Licensed\n"
-                      f"under the terms of the {__license__}.\n")
-                print("Scan the QR code below to login")
-                print("Settings -> Privacy and Security -> Active Sessions -> Scan QR Code.\n")
-
-                qrcode = QRCode(version=1)
-                qrcode.add_data(qr_login.url)
-                qrcode.print_ascii(invert=True)
-            except SessionPasswordNeeded:
-                print(f"Password hint: {await self.get_password_hint()}")
-                return await self.check_password(
-                    await ainput("Enter 2FA password: ", hide=self.hide_password, loop=self.loop)
-                )
+            except SessionPasswordNeeded as e:
+                print(e.MESSAGE)
+
+                while True:
+                    print("Password hint: {}".format(await self.get_password_hint()))
+
+                    if not self.password:
+                        self.password = await ainput("Enter 2FA password (empty to recover): ", hide=self.hide_password, loop=self.loop)
+
+                    try:
+                        if not self.password:
+                            confirm = await ainput("Confirm password recovery (y/N): ", loop=self.loop)
+
+                            if confirm.lower() == "y":
+                                email_pattern = await self.send_recovery_code()
+                                print(f"The recovery code has been sent to {email_pattern}")
+
+                                while True:
+                                    recovery_code = await ainput("Enter recovery code: ", loop=self.loop)
+
+                                    try:
+                                        return await self.recover_password(recovery_code)
+                                    except BadRequest as e:
+                                        print(e.MESSAGE)
+                                    except Exception as e:
+                                        log.exception(e)
+                                        raise
+                            else:
+                                self.password = None
+                        else:
+                            return await self.check_password(self.password)
+                    except BadRequest as e:
+                        print(e.MESSAGE)
+                        self.password = None
+            else:
+                break
 
     def set_parse_mode(self, parse_mode: Optional["enums.ParseMode"]):
         """Set the parse mode to be used globally by the client.

pyrogram/qrlogin.py

@@ -24,22 +24,54 @@
 
 import pyrogram
 from pyrogram import filters, handlers, raw, types
-from pyrogram.methods.messages.inline_session import get_session
+from pyrogram.session import Session
+from pyrogram.session.auth import Auth
 
 log = logging.getLogger(__name__)
 
+
+async def get_session(client: "pyrogram.Client", dc_id: int) -> Session:
+    if dc_id == await client.storage.dc_id():
+        return client.session
+
+    async with client.media_sessions_lock:
+        if client.media_sessions.get(dc_id):
+            return client.media_sessions[dc_id]
+
+        session = client.media_sessions[dc_id] = Session(
+            client,
+            dc_id,
+            await Auth(client, dc_id, await client.storage.test_mode()).create(),
+            await client.storage.test_mode(),
+            is_media=False,
+        )
+
+        await session.start()
+
+        return session
+
+
 class QRLogin:
     def __init__(self, client, except_ids: List[int] = []):
-        self.client = client
-        self.request = raw.functions.auth.ExportLoginToken(
-            api_id=client.api_id,
-            api_hash=client.api_hash,
-            except_ids=except_ids
-        )
-        self.r = None
+        self.client: "pyrogram.Client" = client
+        self.except_ids: List[int] = except_ids
+        self.r: "raw.base.auth.LoginToken" = None
 
     async def recreate(self):
-        self.r = await self.client.invoke(self.request)
+        self.r = await self.client.invoke(
+            raw.functions.auth.ExportLoginToken(
+                api_id=self.client.api_id,
+                api_hash=self.client.api_hash,
+                except_ids=self.except_ids,
+            )
+        )
+
+        if isinstance(self.r, raw.types.auth.LoginTokenMigrateTo):
+            await self.client.storage.dc_id(self.r.dc_id)
+            self.client.session = await get_session(self.client, self.r.dc_id)
+            self.r = await self.client.invoke(
+                raw.functions.auth.ImportLoginToken(token=self.r.token)
+            )
 
         return self.r
 
@@ -58,9 +90,7 @@ async def raw_handler(client, update, users, chats):
         handler = self.client.add_handler(
             handlers.RawUpdateHandler(
                 raw_handler,
-                filters=filters.create(
-                    lambda _, __, u: isinstance(u, raw.types.UpdateLoginToken)
-                )
+                filters=filters.create(lambda _, __, u: isinstance(u, raw.types.UpdateLoginToken)),
             )
         )
 
@@ -74,14 +104,6 @@ async def raw_handler(client, update, users, chats):
 
         await self.recreate()
 
-        if isinstance(self.r, raw.types.auth.LoginTokenMigrateTo):
-            session = await get_session(self.client, self.r.dc_id)
-            self.r = await session.invoke(
-                raw.functions.auth.ImportLoginToken(
-                    token=self.token
-                )
-            )
-
         if isinstance(self.r, raw.types.auth.LoginTokenSuccess):
             user = types.User._parse(self.client, self.r.authorization.user)
 
@@ -90,7 +112,7 @@ async def raw_handler(client, update, users, chats):
 
             return user
 
-        raise TypeError('Unexpected login token response: {}'.format(self.r))
+        raise TypeError("Unexpected login token response: {}".format(self.r))
 
     @property
     def url(self) -> str: