From ade31f8989649053c87c8cfb78d00f7f53850edb Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Sun, 24 Apr 2022 11:56:07 +0200 Subject: [PATCH] Update the session string format --- pyrogram/client.py | 67 ++++++++++++++++-------------- pyrogram/storage/memory_storage.py | 33 +++++++++++---- pyrogram/storage/storage.py | 31 +++++++------- 3 files changed, 78 insertions(+), 53 deletions(-) diff --git a/pyrogram/client.py b/pyrogram/client.py index 18e5253b..6b2ac9a1 100644 --- a/pyrogram/client.py +++ b/pyrogram/client.py @@ -49,7 +49,7 @@ from pyrogram.errors import ( from pyrogram.handlers.handler import Handler from pyrogram.methods import Methods from pyrogram.session import Auth, Session -from pyrogram.storage import Storage, FileStorage, MemoryStorage +from pyrogram.storage import FileStorage, MemoryStorage from pyrogram.types import User, TermsOfService from pyrogram.utils import ainput from .dispatcher import Dispatcher @@ -65,14 +65,8 @@ class Client(Methods): """Pyrogram Client, the main means for interacting with Telegram. Parameters: - session_name (``str``): - Pass a string of your choice to give a name to the client session, e.g.: "*my_account*". This name will be - used to save a file on disk that stores details needed to reconnect without asking again for credentials. - Alternatively, if you don't want a file to be saved on disk, pass the special name ``":memory:"`` to start - an in-memory session that will be discarded as soon as you stop the Client. In order to reconnect again - using a memory storage without having to login again, you can use - :meth:`~pyrogram.Client.export_session_string` before stopping the client to get a session string you can - pass here as argument. + name (``str``): + Pass a string of your choice to give a name to the client, e.g.: "my_account". api_id (``int`` | ``str``, *optional*): The *api_id* part of your Telegram API key, as integer. @@ -116,6 +110,17 @@ class Client(Methods): Pass your Bot API token to create a bot session, e.g.: "123456:ABC-DEF1234ghIkl-zyx57W2v1u123ew11" Only applicable for new sessions. + session_string (``str``, *optional*): + Pass a session string to load the session in-memory. + Implies ``in_memory=True``. + + in_memory (``bool``, *optional*): + Pass True to start an in-memory session that will be discarded as soon as the client stops. + In order to reconnect again using an in-memory session without having to login again, you can use + :meth:`~pyrogram.Client.export_session_string` before stopping the client to get a session string you can + pass to the ``session_string`` parameter. + Defaults to False. + phone_number (``str``, *optional*): Pass your phone number as string (with your Country Code prefix included) to avoid entering it manually. Only applicable for new sessions. @@ -187,7 +192,7 @@ class Client(Methods): def __init__( self, - session_name: Union[str, Storage], + name: str, api_id: int = None, api_hash: str = None, app_version: str = APP_VERSION, @@ -198,6 +203,8 @@ class Client(Methods): proxy: dict = None, test_mode: bool = False, bot_token: str = None, + session_string: str = None, + in_memory: bool = None, phone_number: str = None, phone_code: str = None, password: str = None, @@ -212,7 +219,7 @@ class Client(Methods): ): super().__init__() - self.session_name = session_name + self.name = name self.api_id = api_id self.api_hash = api_hash self.app_version = app_version @@ -223,6 +230,8 @@ class Client(Methods): self.proxy = proxy self.test_mode = test_mode self.bot_token = bot_token + self.session_string = session_string + self.in_memory = in_memory self.phone_number = phone_number self.phone_code = phone_code self.password = password @@ -237,16 +246,12 @@ class Client(Methods): self.executor = ThreadPoolExecutor(self.workers, thread_name_prefix="Handler") - if isinstance(session_name, str): - if session_name == ":memory:" or len(session_name) >= MemoryStorage.SESSION_STRING_SIZE: - session_name = re.sub(r"[\n\s]+", "", session_name) - self.storage = MemoryStorage(session_name) - else: - self.storage = FileStorage(session_name, self.workdir) - elif isinstance(session_name, Storage): - self.storage = session_name + if self.session_string: + self.storage = MemoryStorage(self.name, self.session_string) + elif self.in_memory: + self.storage = MemoryStorage(self.name) else: - raise ValueError("Unknown storage engine") + self.storage = FileStorage(self.name, self.workdir) self.dispatcher = Dispatcher(self) @@ -638,7 +643,7 @@ class Client(Methods): self.add_handler(handler, group) log.info('[{}] [LOAD] {}("{}") in group {} from "{}"'.format( - self.session_name, type(handler).__name__, name, group, module_path)) + self.name, type(handler).__name__, name, group, module_path)) count += 1 except Exception: @@ -651,11 +656,11 @@ class Client(Methods): try: module = import_module(module_path) except ImportError: - log.warning(f'[{self.session_name}] [LOAD] Ignoring non-existent module "{module_path}"') + log.warning(f'[{self.name}] [LOAD] Ignoring non-existent module "{module_path}"') continue if "__path__" in dir(module): - log.warning(f'[{self.session_name}] [LOAD] Ignoring namespace "{module_path}"') + log.warning(f'[{self.name}] [LOAD] Ignoring namespace "{module_path}"') continue if handlers is None: @@ -670,13 +675,13 @@ class Client(Methods): self.add_handler(handler, group) log.info('[{}] [LOAD] {}("{}") in group {} from "{}"'.format( - self.session_name, type(handler).__name__, name, group, module_path)) + self.name, type(handler).__name__, name, group, module_path)) count += 1 except Exception: if warn_non_existent_functions: log.warning('[{}] [LOAD] Ignoring non-existent function "{}" from "{}"'.format( - self.session_name, name, module_path)) + self.name, name, module_path)) if exclude: for path, handlers in exclude: @@ -686,11 +691,11 @@ class Client(Methods): try: module = import_module(module_path) except ImportError: - log.warning(f'[{self.session_name}] [UNLOAD] Ignoring non-existent module "{module_path}"') + log.warning(f'[{self.name}] [UNLOAD] Ignoring non-existent module "{module_path}"') continue if "__path__" in dir(module): - log.warning(f'[{self.session_name}] [UNLOAD] Ignoring namespace "{module_path}"') + log.warning(f'[{self.name}] [UNLOAD] Ignoring namespace "{module_path}"') continue if handlers is None: @@ -705,19 +710,19 @@ class Client(Methods): self.remove_handler(handler, group) log.info('[{}] [UNLOAD] {}("{}") from group {} in "{}"'.format( - self.session_name, type(handler).__name__, name, group, module_path)) + self.name, type(handler).__name__, name, group, module_path)) count -= 1 except Exception: if warn_non_existent_functions: log.warning('[{}] [UNLOAD] Ignoring non-existent function "{}" from "{}"'.format( - self.session_name, name, module_path)) + self.name, name, module_path)) if count > 0: log.info('[{}] Successfully loaded {} plugin{} from "{}"'.format( - self.session_name, count, "s" if count > 1 else "", root)) + self.name, count, "s" if count > 1 else "", root)) else: - log.warning(f'[{self.session_name}] No plugin loaded from "{root}"') + log.warning(f'[{self.name}] No plugin loaded from "{root}"') async def handle_download(self, packet): file_id, directory, file_name, in_memory, file_size, progress, progress_args = packet diff --git a/pyrogram/storage/memory_storage.py b/pyrogram/storage/memory_storage.py index 1035356d..2c01f447 100644 --- a/pyrogram/storage/memory_storage.py +++ b/pyrogram/storage/memory_storage.py @@ -27,23 +27,42 @@ log = logging.getLogger(__name__) class MemoryStorage(SQLiteStorage): - def __init__(self, name: str): + def __init__(self, name: str, session_string: str = None): super().__init__(name) + self.session_string = session_string + async def open(self): self.conn = sqlite3.connect(":memory:", check_same_thread=False) self.create() - if self.name != ":memory:": - dc_id, test_mode, auth_key, user_id, is_bot = struct.unpack( - (self.SESSION_STRING_FORMAT if len(self.name) == MemoryStorage.SESSION_STRING_SIZE else - self.SESSION_STRING_FORMAT_64), - base64.urlsafe_b64decode( - self.name + "=" * (-len(self.name) % 4) + if self.session_string: + # Old format + if len(self.session_string) in [self.SESSION_STRING_SIZE, self.SESSION_STRING_SIZE_64]: + dc_id, test_mode, auth_key, user_id, is_bot = struct.unpack( + (self.OLD_SESSION_STRING_FORMAT + if len(self.session_string) == self.SESSION_STRING_SIZE else + self.OLD_SESSION_STRING_FORMAT_64), + base64.urlsafe_b64decode(self.session_string + "=" * (-len(self.session_string) % 4)) ) + + await self.dc_id(dc_id) + await self.test_mode(test_mode) + await self.auth_key(auth_key) + await self.user_id(user_id) + await self.is_bot(is_bot) + await self.date(0) + + log.warning("You are using an old session string format. Use export_session_string to update") + return + + dc_id, api_id, test_mode, auth_key, user_id, is_bot = struct.unpack( + self.SESSION_STRING_FORMAT, + base64.urlsafe_b64decode(self.session_string + "=" * (-len(self.session_string) % 4)) ) await self.dc_id(dc_id) + await self.api_id(api_id) await self.test_mode(test_mode) await self.auth_key(auth_key) await self.user_id(user_id) diff --git a/pyrogram/storage/storage.py b/pyrogram/storage/storage.py index 8daaae7e..de397718 100644 --- a/pyrogram/storage/storage.py +++ b/pyrogram/storage/storage.py @@ -17,18 +17,19 @@ # along with Pyrogram. If not, see . import base64 +import lzma import struct from typing import List, Tuple -from pyrogram import utils - class Storage: - SESSION_STRING_FORMAT = ">B?256sI?" - SESSION_STRING_FORMAT_64 = ">B?256sQ?" + OLD_SESSION_STRING_FORMAT = ">B?256sI?" + OLD_SESSION_STRING_FORMAT_64 = ">B?256sQ?" SESSION_STRING_SIZE = 351 SESSION_STRING_SIZE_64 = 356 + SESSION_STRING_FORMAT = ">BI?256sQ?" + def __init__(self, name: str): self.name = name @@ -78,14 +79,14 @@ class Storage: raise NotImplementedError async def export_session_string(self): - user_id = await self.user_id() - return base64.urlsafe_b64encode( - struct.pack( - self.SESSION_STRING_FORMAT if user_id < utils.MAX_USER_ID_OLD else self.SESSION_STRING_FORMAT_64, - await self.dc_id(), - await self.test_mode(), - await self.auth_key(), - user_id, - await self.is_bot() - ) - ).decode().rstrip("=") + packed = struct.pack( + self.SESSION_STRING_FORMAT, + await self.dc_id(), + await self.api_id(), + await self.test_mode(), + await self.auth_key(), + await self.user_id(), + await self.is_bot() + ) + + return base64.urlsafe_b64encode(packed).decode().rstrip("=")