Skip to content

Commit

Permalink
Add QR Login support
Browse files Browse the repository at this point in the history
Thanks to @Lonami
Co-authored-by: Lonami <Lonami@totufals@hotmail.com>
  • Loading branch information
KurimuzonAkuma committed Dec 17, 2024
1 parent a6dbd7f commit e22ec80
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 12 deletions.
37 changes: 35 additions & 2 deletions pyrogram/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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__}"
Expand Down Expand Up @@ -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__()

Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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.
Expand Down
14 changes: 12 additions & 2 deletions pyrogram/methods/utilities/start.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.

import logging
from typing import List

import pyrogram
from pyrogram import raw
Expand All @@ -26,7 +27,8 @@

class Start:
async def start(
self: "pyrogram.Client"
self: "pyrogram.Client",
except_ids: List[int] = [],
):
"""Start the client.
Expand Down Expand Up @@ -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
Expand Down
85 changes: 77 additions & 8 deletions pyrogram/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,24 +16,24 @@
# You should have received a copy of the GNU Lesser General Public License
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.

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
import hashlib
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):
Expand Down Expand Up @@ -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')}"

0 comments on commit e22ec80

Please sign in to comment.