Update the session string format

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``):
+        name (``str``):
-            Pass a string of your choice to give a name to the client session, e.g.: "*my_account*". This name will be
+            Pass a string of your choice to give a name to the client, e.g.: "my_account".
-            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.
 
         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 self.session_string:
-            if session_name == ":memory:" or len(session_name) >= MemoryStorage.SESSION_STRING_SIZE:
+            self.storage = MemoryStorage(self.name, self.session_string)
-                session_name = re.sub(r"[\n\s]+", "", session_name)
+        elif self.in_memory:
-                self.storage = MemoryStorage(session_name)
+            self.storage = MemoryStorage(self.name)
         else:
-                self.storage = FileStorage(session_name, self.workdir)
+            self.storage = FileStorage(self.name, self.workdir)
-        elif isinstance(session_name, Storage):
-            self.storage = session_name
-        else:
-            raise ValueError("Unknown storage engine")
 
         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

pyrogram/storage/memory_storage.py

@@ -27,20 +27,23 @@ 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:":
+        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.SESSION_STRING_FORMAT if len(self.name) == MemoryStorage.SESSION_STRING_SIZE else
+                    (self.OLD_SESSION_STRING_FORMAT
-                 self.SESSION_STRING_FORMAT_64),
+                     if len(self.session_string) == self.SESSION_STRING_SIZE else
-                base64.urlsafe_b64decode(
+                     self.OLD_SESSION_STRING_FORMAT_64),
-                    self.name + "=" * (-len(self.name) % 4)
+                    base64.urlsafe_b64decode(self.session_string + "=" * (-len(self.session_string) % 4))
-                )
                 )
 
                 await self.dc_id(dc_id)
@@ -50,5 +53,21 @@ class MemoryStorage(SQLiteStorage):
                 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)
+            await self.is_bot(is_bot)
+            await self.date(0)
+
     async def delete(self):
         pass

pyrogram/storage/storage.py

@@ -17,18 +17,19 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import base64
+import lzma
 import struct
 from typing import List, Tuple
 
-from pyrogram import utils
-
 
 class Storage:
-    SESSION_STRING_FORMAT = ">B?256sI?"
+    OLD_SESSION_STRING_FORMAT = ">B?256sI?"
-    SESSION_STRING_FORMAT_64 = ">B?256sQ?"
+    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()
+        packed = struct.pack(
-        return base64.urlsafe_b64encode(
+            self.SESSION_STRING_FORMAT,
-            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.api_id(),
             await self.test_mode(),
             await self.auth_key(),
-                user_id,
+            await self.user_id(),
             await self.is_bot()
         )
-        ).decode().rstrip("=")
+
+        return base64.urlsafe_b64encode(packed).decode().rstrip("=")