diff --git a/pyrogram/client.py b/pyrogram/client.py index beab6ae212..bdb0987db2 100644 --- a/pyrogram/client.py +++ b/pyrogram/client.py @@ -46,7 +46,7 @@ VolumeLocNotFound, ChannelPrivate, BadRequest, AuthBytesInvalid, FloodWait, FloodPremiumWait, - ChannelInvalid, PersistentTimestampInvalid, PersistentTimestampOutdated + ChannelInvalid, PersistentTimestampInvalid, PersistentTimestampOutdated, ) from pyrogram.handlers.handler import Handler from pyrogram.methods import Methods @@ -209,6 +209,10 @@ class Client(Methods): init_connection_params (:obj:`~pyrogram.raw.base.JSONValue`, *optional*): Additional initConnection parameters. For now, only the tz_offset field is supported, for specifying timezone offset in seconds. + + use_qr_login (``bool``, *optional*): + Use QR Code for login to your account. + Defaults to False. """ APP_VERSION = f"Pyrogram {__version__}" @@ -269,7 +273,8 @@ def __init__( client_platform: "enums.ClientPlatform" = enums.ClientPlatform.OTHER, init_connection_params: Optional["raw.base.JSONValue"] = None, connection_factory: Type[Connection] = Connection, - protocol_factory: Type[TCP] = TCPAbridged + protocol_factory: Type[TCP] = TCPAbridged, + use_qr_login: bool = False, ): super().__init__() @@ -306,6 +311,7 @@ def __init__( self.init_connection_params = init_connection_params self.connection_factory = connection_factory self.protocol_factory = protocol_factory + self.use_qr_login = use_qr_login self.executor = ThreadPoolExecutor(self.workers, thread_name_prefix="Handler") @@ -507,6 +513,33 @@ async def authorize(self) -> User: return signed_up + async def authorize_qr(self, except_ids: List[int] = []) -> Optional[User]: + import qrcode + qr_login = utils.QRLogin(self, except_ids) + + while True: + try: + if await qr_login.wait(): + break + except asyncio.TimeoutError: + 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\n") + print("Scan the QR code below to login") + print("Settings -> Privacy and Security -> Active Sessions -> Scan QR Code.\n") + + qrcode = qrcode.QRCode(version=1) + qrcode.add_data(qr_login.url) + qrcode.print_ascii(invert=True) + except SessionPasswordNeeded: + print(f"Password hint: {await self.client.get_password_hint()}") + await self.client.check_password( + await ainput("Enter 2FA password: ", hide=self.client.hide_password) + ) + continue + def set_parse_mode(self, parse_mode: Optional["enums.ParseMode"]): """Set the parse mode to be used globally by the client. diff --git a/pyrogram/methods/utilities/start.py b/pyrogram/methods/utilities/start.py index 922ea72bc2..6d1b00bcab 100644 --- a/pyrogram/methods/utilities/start.py +++ b/pyrogram/methods/utilities/start.py @@ -17,6 +17,7 @@ # along with Pyrogram. If not, see . import logging +from typing import List import pyrogram from pyrogram import raw @@ -26,7 +27,8 @@ class Start: async def start( - self: "pyrogram.Client" + self: "pyrogram.Client", + except_ids: List[int] = [], ): """Start the client. @@ -59,7 +61,15 @@ async def main(): try: if not is_authorized: - await self.authorize() + if self.use_qr_login: + try: + import qrcode + await self.authorize_qr(except_ids=except_ids) + except ImportError: + log.warning("qrcode package not found, falling back to authorization prompt") + await self.authorize() + else: + await self.authorize() if self.takeout and not await self.storage.is_bot(): self.takeout_id = (await self.invoke(raw.functions.account.InitTakeoutSession())).id diff --git a/pyrogram/utils.py b/pyrogram/utils.py index 040b35010e..dba935707a 100644 --- a/pyrogram/utils.py +++ b/pyrogram/utils.py @@ -16,11 +16,6 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . -from concurrent.futures.thread import ThreadPoolExecutor -from datetime import datetime, timezone -from getpass import getpass -from io import BytesIO -from typing import Union, List, Dict, Optional import asyncio import base64 import functools @@ -28,12 +23,17 @@ import os import re import struct +from concurrent.futures.thread import ThreadPoolExecutor +from datetime import datetime, timezone +from getpass import getpass +from io import BytesIO +from typing import Dict, List, Optional, Union import pyrogram -from pyrogram import raw, enums -from pyrogram import types +from pyrogram import Client, enums, filters, handlers, raw, types +from pyrogram.file_id import DOCUMENT_TYPES, PHOTO_TYPES, FileId, FileType +from pyrogram.methods.messages.inline_session import get_session from pyrogram.types.messages_and_media.message import Str -from pyrogram.file_id import FileId, FileType, PHOTO_TYPES, DOCUMENT_TYPES async def ainput(prompt: str = "", *, hide: bool = False): @@ -574,3 +574,72 @@ def from_inline_bytes(data: bytes, file_name: str = None) -> BytesIO: b.name = file_name or f"photo_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.jpg" return b + + +class QRLogin: + def __init__(self, client: 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 + + async def recreate(self): + self.r = await self.client.invoke(self.request) + + return self.r + + async def wait(self, timeout: float = None) -> Optional["types.User"]: + if timeout is None: + if not self.r: + raise asyncio.TimeoutError + + timeout = self.r.expires - int(datetime.datetime.now().timestamp()) + + event = asyncio.Event() + + async def raw_handler(client: Client, update, users, chats): + event.set() + + await self.client.dispatcher.start() + + handler = self.client.add_handler( + handlers.RawUpdateHandler( + raw_handler, + filters=filters.create( + lambda _, __, u: isinstance(u, raw.types.UpdateLoginToken) + ) + ) + ) + + try: + await asyncio.wait_for(event.wait(), timeout=timeout) + finally: + self.client.remove_handler(*handler) + await self.client.dispatcher.stop() + + 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) + + await self.client.storage.user_id(user.id) + await self.client.storage.is_bot(False) + + return user + + raise TypeError('Unexpected login token response: {}'.format(self.r)) + + @property + def url(self) -> str: + return f"tg://login?token={base64.urlsafe_b64encode(self.r.token).decode('utf-8')}"