Merge pull request #220 from bakatrouble/session_storage

Implement extendable session storage, move to SQLite as main storage (WIP)

MANIFEST.in

@@ -1,7 +1,7 @@
 ## Include
 include README.md COPYING COPYING.lesser NOTICE requirements.txt
 recursive-include compiler *.py *.tl *.tsv *.txt
-recursive-include pyrogram mime.types
+recursive-include pyrogram mime.types schema.sql
 
 ## Exclude
 prune pyrogram/api/errors/exceptions

docs/source/api/methods.rst

@@ -32,6 +32,7 @@ Utilities
     - :meth:`~Client.add_handler`
     - :meth:`~Client.remove_handler`
     - :meth:`~Client.stop_transmission`
+    - :meth:`~Client.export_session_string`
 
 Messages
 ^^^^^^^^
@@ -186,6 +187,7 @@ Details
 .. automethod:: Client.add_handler()
 .. automethod:: Client.remove_handler()
 .. automethod:: Client.stop_transmission()
+.. automethod:: Client.export_session_string()
 
 ..  Messages
 .. automethod:: Client.send_message()

docs/source/glossary.rst

@@ -58,7 +58,7 @@ Terms
         Pyrogram --- to automate some behaviours, like sending messages or reacting to text commands or any other event.
 
     Session
-        Also known as *login session*, is a strictly personal piece of information created and held by both parties
+        Also known as *login session*, is a strictly personal piece of data created and held by both parties
         (client and server) which is used to grant permission into a single account without having to start a new
         authorization process from scratch.
 

docs/source/index.rst

@@ -130,6 +130,7 @@ Meta
     topics/auto-auth
     topics/session-settings
     topics/tgcrypto
+    topics/storage-engines
     topics/text-formatting
     topics/serialize
     topics/proxy

docs/source/topics/storage-engines.rst

@@ -0,0 +1,95 @@
+Storage Engines
+===============
+
+Every time you login to Telegram, some personal piece of data are created and held by both parties (the client, Pyrogram
+and the server, Telegram). This session data is uniquely bound to your own account, indefinitely (until you logout or
+decide to manually terminate it) and is used to authorize a client to execute API calls on behalf of your identity.
+
+Persisting Sessions
+-------------------
+
+In order to make a client reconnect successfully between restarts, that is, without having to start a new
+authorization process from scratch each time, Pyrogram needs to store the generated session data somewhere.
+
+Other useful data being stored is peers' cache. In short, peers are all those entities you can chat with, such as users
+or bots, basic groups, but also channels and supergroups. Because of how Telegram works, a unique pair of **id** and
+**access_hash** is needed to contact a peer. This, plus other useful info such as the peer type, is what is stored
+inside a session storage.
+
+So, if you ever wondered how is Pyrogram able to contact peers just by asking for their ids, it's because of this very
+reason: the peer *id* is looked up in the internal database and the available *access_hash* is retrieved, which is then
+used to correctly invoke API methods.
+
+Different Storage Engines
+-------------------------
+
+Let's now talk about how Pyrogram actually stores all the relevant data. Pyrogram offers two different types of storage
+engines: a **File Storage** and a **Memory Storage**. These engines are well integrated in the library and require a
+minimal effort to set up. Here's how they work:
+
+File Storage
+^^^^^^^^^^^^
+
+This is the most common storage engine. It is implemented by using **SQLite**, which will store the session and peers
+details. The database will be saved to disk as a single portable file and is designed to efficiently save and retrieve
+peers whenever they are needed.
+
+To use this type of engine, simply pass any name of your choice to the ``session_name`` parameter of the
+:obj:`~pyrogram.Client` constructor, as usual:
+
+.. code-block:: python
+
+    from pyrogram import Client
+
+    with Client("my_account") as app:
+        print(app.get_me())
+
+Once you successfully log in (either with a user or a bot identity), a session file will be created and saved to disk as
+``my_account.session``. Any subsequent client restart will make Pyrogram search for a file named that way and the
+session database will be automatically loaded.
+
+Memory Storage
+^^^^^^^^^^^^^^
+
+In case you don't want to have any session file saved on disk, you can use an in-memory storage by passing the special
+session name "**:memory:**" to the ``session_name`` parameter of the :obj:`~pyrogram.Client` constructor:
+
+.. code-block:: python
+
+    from pyrogram import Client
+
+    with Client(":memory:") as app:
+        print(app.get_me())
+
+This database is still backed by SQLite, but exists purely in memory. However, once you stop a client, the entire
+database is discarded and the session details used for logging in again will be lost forever.
+
+Session Strings
+---------------
+
+Session strings are useful when you want to run authorized Pyrogram clients on platforms like
+`Heroku <https://www.heroku.com/>`_, where their ephemeral filesystems makes it much harder for a file-based storage
+engine to properly work as intended.
+
+In case you want to use an in-memory storage, but also want to keep access to the session you created, call
+:meth:`~pyrogram.Client.export_session_string` anytime before stopping the client...
+
+.. code-block:: python
+
+    from pyrogram import Client
+
+    with Client(":memory:") as app:
+        print(app.export_session_string())
+
+...and save the resulting string somewhere. You can use this string as session name the next time you want to login
+using the same session; the storage used will still be completely in-memory:
+
+.. code-block:: python
+
+    from pyrogram import Client
+
+    session_string = "...ZnUIFD8jsjXTb8g_vpxx48k1zkov9sapD-tzjz-S4WZv70M..."
+
+    with Client(session_string) as app:
+        print(app.get_me())
+

pyrogram/client/client.py

@@ -16,8 +16,6 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-import base64
-import json
 import logging
 import math
 import mimetypes
@@ -52,6 +50,7 @@ from pyrogram.errors import (
 from pyrogram.session import Auth, Session
 from .ext import utils, Syncer, BaseClient, Dispatcher
 from .methods import Methods
+from .storage import Storage, FileStorage, MemoryStorage
 
 log = logging.getLogger(__name__)
 
@@ -61,8 +60,13 @@ class Client(Methods, BaseClient):
 
     Parameters:
         session_name (``str``):
-            Name to uniquely identify a session of either a User or a Bot, e.g.: "my_account". This name will be used
+            Pass a string of your choice to give a name to the client session, e.g.: "*my_account*". This name will be
-            to save a file to disk that stores details needed for reconnecting without asking again for credentials.
+            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``, *optional*):
             The *api_id* part of your Telegram API Key, as integer. E.g.: 12345
@@ -176,7 +180,7 @@ class Client(Methods, BaseClient):
 
     def __init__(
         self,
-        session_name: str,
+        session_name: Union[str, Storage],
         api_id: Union[int, str] = None,
         api_hash: str = None,
         app_version: str = None,
@@ -223,12 +227,23 @@ class Client(Methods, BaseClient):
         self.first_name = first_name
         self.last_name = last_name
         self.workers = workers
-        self.workdir = workdir
+        self.workdir = Path(workdir)
-        self.config_file = config_file
+        self.config_file = Path(config_file)
         self.plugins = plugins
         self.no_updates = no_updates
         self.takeout = takeout
 
+        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
+        else:
+            raise ValueError("Unknown storage engine")
+
         self.dispatcher = Dispatcher(self, workers)
 
     def __enter__(self):
@@ -263,50 +278,32 @@ class Client(Methods, BaseClient):
         if self.is_started:
             raise ConnectionError("Client has already been started")
 
-        if self.BOT_TOKEN_RE.match(self.session_name):
-            self.is_bot = True
-            self.bot_token = self.session_name
-            self.session_name = self.session_name.split(":")[0]
-            log.warning('\nWARNING: You are using a bot token as session name!\n'
-                        'This usage will be deprecated soon. Please use a session file name to load '
-                        'an existing session and the bot_token argument to create new sessions.\n'
-                        'More info: https://docs.pyrogram.org/intro/auth#bot-authorization\n')
-
         self.load_config()
         self.load_session()
         self.load_plugins()
 
-        self.session = Session(
+        self.session = Session(self, self.storage.dc_id, self.storage.auth_key)
-            self,
-            self.dc_id,
-            self.auth_key
-        )
 
         self.session.start()
         self.is_started = True
 
         try:
-            if self.user_id is None:
+            if self.storage.user_id is None:
                 if self.bot_token is None:
-                    self.is_bot = False
+                    self.storage.is_bot = False
                     self.authorize_user()
                 else:
-                    self.is_bot = True
+                    self.storage.is_bot = True
                     self.authorize_bot()
 
-                self.save_session()
+            if not self.storage.is_bot:
-
-            if not self.is_bot:
                 if self.takeout:
                     self.takeout_id = self.send(functions.account.InitTakeoutSession()).id
                     log.warning("Takeout session {} initiated".format(self.takeout_id))
 
                 now = time.time()
 
-                if abs(now - self.date) > Client.OFFLINE_SLEEP:
+                if abs(now - self.storage.date) > Client.OFFLINE_SLEEP:
-                    self.peers_by_username = {}
-                    self.peers_by_phone = {}
-
                     self.get_initial_dialogs()
                     self.get_contacts()
                 else:
@@ -505,19 +502,15 @@ class Client(Methods, BaseClient):
         except UserMigrate as e:
             self.session.stop()
 
-            self.dc_id = e.x
+            self.storage.dc_id = e.x
-            self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create()
+            self.storage.auth_key = Auth(self, self.storage.dc_id).create()
-
+            self.session = Session(self, self.storage.dc_id, self.storage.auth_key)
-            self.session = Session(
-                self,
-                self.dc_id,
-                self.auth_key
-            )
 
             self.session.start()
+
             self.authorize_bot()
         else:
-            self.user_id = r.user.id
+            self.storage.user_id = r.user.id
 
             print("Logged in successfully as @{}".format(r.user.username))
 
@@ -558,20 +551,10 @@ class Client(Methods, BaseClient):
             except (PhoneMigrate, NetworkMigrate) as e:
                 self.session.stop()
 
-                self.dc_id = e.x
+                self.storage.dc_id = e.x
+                self.storage.auth_key = Auth(self, self.storage.dc_id).create()
 
-                self.auth_key = Auth(
+                self.session = Session(self, self.storage.dc_id, self.storage.auth_key)
-                    self.dc_id,
-                    self.test_mode,
-                    self.ipv6,
-                    self._proxy
-                ).create()
-
-                self.session = Session(
-                    self,
-                    self.dc_id,
-                    self.auth_key
-                )
 
                 self.session.start()
             except (PhoneNumberInvalid, PhoneNumberBanned) as e:
@@ -751,13 +734,13 @@ class Client(Methods, BaseClient):
             )
 
         self.password = None
-        self.user_id = r.user.id
+        self.storage.user_id = r.user.id
 
         print("Logged in successfully as {}".format(r.user.first_name))
 
     def fetch_peers(
         self,
-        entities: List[
+        peers: List[
             Union[
                 types.User,
                 types.Chat, types.ChatForbidden,
@@ -766,64 +749,57 @@ class Client(Methods, BaseClient):
         ]
     ) -> bool:
         is_min = False
+        parsed_peers = []
 
-        for entity in entities:
+        for peer in peers:
-            if isinstance(entity, types.User):
+            username = None
-                user_id = entity.id
+            phone_number = None
 
-                access_hash = entity.access_hash
+            if isinstance(peer, types.User):
+                peer_id = peer.id
+                access_hash = peer.access_hash
+
+                username = peer.username
+                phone_number = peer.phone
+
+                if peer.bot:
+                    peer_type = "bot"
+                else:
+                    peer_type = "user"
 
                 if access_hash is None:
                     is_min = True
                     continue
 
-                username = entity.username
-                phone = entity.phone
-
-                input_peer = types.InputPeerUser(
-                    user_id=user_id,
-                    access_hash=access_hash
-                )
-
-                self.peers_by_id[user_id] = input_peer
-
                 if username is not None:
-                    self.peers_by_username[username.lower()] = input_peer
+                    username = username.lower()
+            elif isinstance(peer, (types.Chat, types.ChatForbidden)):
+                peer_id = -peer.id
+                access_hash = 0
+                peer_type = "group"
+            elif isinstance(peer, (types.Channel, types.ChannelForbidden)):
+                peer_id = int("-100" + str(peer.id))
+                access_hash = peer.access_hash
 
-                if phone is not None:
+                username = getattr(peer, "username", None)
-                    self.peers_by_phone[phone] = input_peer
 
-            if isinstance(entity, (types.Chat, types.ChatForbidden)):
+                if peer.broadcast:
-                chat_id = entity.id
+                    peer_type = "channel"
-                peer_id = -chat_id
+                else:
-
+                    peer_type = "supergroup"
-                input_peer = types.InputPeerChat(
-                    chat_id=chat_id
-                )
-
-                self.peers_by_id[peer_id] = input_peer
-
-            if isinstance(entity, (types.Channel, types.ChannelForbidden)):
-                channel_id = entity.id
-                peer_id = int("-100" + str(channel_id))
-
-                access_hash = entity.access_hash
 
                 if access_hash is None:
                     is_min = True
                     continue
 
-                username = getattr(entity, "username", None)
-
-                input_peer = types.InputPeerChannel(
-                    channel_id=channel_id,
-                    access_hash=access_hash
-                )
-
-                self.peers_by_id[peer_id] = input_peer
-
                 if username is not None:
-                    self.peers_by_username[username.lower()] = input_peer
+                    username = username.lower()
+            else:
+                continue
+
+            parsed_peers.append((peer_id, access_hash, peer_type, username, phone_number))
+
+        self.storage.update_peers(parsed_peers)
 
         return is_min
 
@@ -1084,36 +1060,23 @@ class Client(Methods, BaseClient):
                 self.plugins = None
 
     def load_session(self):
-        try:
+        self.storage.open()
-            with open(os.path.join(self.workdir, "{}.session".format(self.session_name)), encoding="utf-8") as f:
-                s = json.load(f)
-        except FileNotFoundError:
-            self.dc_id = 1
-            self.date = 0
-            self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create()
-        else:
-            self.dc_id = s["dc_id"]
-            self.test_mode = s["test_mode"]
-            self.auth_key = base64.b64decode("".join(s["auth_key"]))
-            self.user_id = s["user_id"]
-            self.date = s.get("date", 0)
-            # TODO: replace default with False once token session name will be deprecated
-            self.is_bot = s.get("is_bot", self.is_bot)
 
-            for k, v in s.get("peers_by_id", {}).items():
+        session_empty = any([
-                self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v)
+            self.storage.test_mode is None,
+            self.storage.auth_key is None,
+            self.storage.user_id is None,
+            self.storage.is_bot is None
+        ])
 
-            for k, v in s.get("peers_by_username", {}).items():
+        if session_empty:
-                peer = self.peers_by_id.get(v, None)
+            self.storage.dc_id = 1
+            self.storage.date = 0
 
-                if peer:
+            self.storage.test_mode = self.test_mode
-                    self.peers_by_username[k] = peer
+            self.storage.auth_key = Auth(self, self.storage.dc_id).create()
-
+            self.storage.user_id = None
-            for k, v in s.get("peers_by_phone", {}).items():
+            self.storage.is_bot = None
-                peer = self.peers_by_id.get(v, None)
-
-                if peer:
-                    self.peers_by_phone[k] = peer
 
     def load_plugins(self):
         if self.plugins:
@@ -1237,26 +1200,6 @@ class Client(Methods, BaseClient):
                 log.warning('[{}] No plugin loaded from "{}"'.format(
                     self.session_name, root))
 
-    def save_session(self):
-        auth_key = base64.b64encode(self.auth_key).decode()
-        auth_key = [auth_key[i: i + 43] for i in range(0, len(auth_key), 43)]
-
-        os.makedirs(self.workdir, exist_ok=True)
-
-        with open(os.path.join(self.workdir, "{}.session".format(self.session_name)), "w", encoding="utf-8") as f:
-            json.dump(
-                dict(
-                    dc_id=self.dc_id,
-                    test_mode=self.test_mode,
-                    auth_key=auth_key,
-                    user_id=self.user_id,
-                    date=self.date,
-                    is_bot=self.is_bot,
-                ),
-                f,
-                indent=4
-            )
-
     def get_initial_dialogs_chunk(self, offset_date: int = 0):
         while True:
             try:
@@ -1274,7 +1217,7 @@ class Client(Methods, BaseClient):
                 log.warning("get_dialogs flood: waiting {} seconds".format(e.x))
                 time.sleep(e.x)
             else:
-                log.info("Total peers: {}".format(len(self.peers_by_id)))
+                log.info("Total peers: {}".format(self.storage.peers_count))
                 return r
 
     def get_initial_dialogs(self):
@@ -1312,7 +1255,7 @@ class Client(Methods, BaseClient):
             KeyError: In case the peer doesn't exist in the internal database.
         """
         try:
-            return self.peers_by_id[peer_id]
+            return self.storage.get_peer_by_id(peer_id)
         except KeyError:
             if type(peer_id) is str:
                 if peer_id in ("self", "me"):
@@ -1323,17 +1266,19 @@ class Client(Methods, BaseClient):
                 try:
                     int(peer_id)
                 except ValueError:
-                    if peer_id not in self.peers_by_username:
+                    try:
+                        return self.storage.get_peer_by_username(peer_id)
+                    except KeyError:
                         self.send(
                             functions.contacts.ResolveUsername(
                                 username=peer_id
                             )
                         )
 
-                    return self.peers_by_username[peer_id]
+                        return self.storage.get_peer_by_username(peer_id)
                 else:
                     try:
-                        return self.peers_by_phone[peer_id]
+                        return self.storage.get_peer_by_phone_number(peer_id)
                     except KeyError:
                         raise PeerIdInvalid
 
@@ -1341,7 +1286,10 @@ class Client(Methods, BaseClient):
                 self.fetch_peers(
                     self.send(
                         functions.users.GetUsers(
-                            id=[types.InputUser(user_id=peer_id, access_hash=0)]
+                            id=[types.InputUser(
+                                user_id=peer_id,
+                                access_hash=0
+                            )]
                         )
                     )
                 )
@@ -1349,7 +1297,10 @@ class Client(Methods, BaseClient):
                 if str(peer_id).startswith("-100"):
                     self.send(
                         functions.channels.GetChannels(
-                            id=[types.InputChannel(channel_id=int(str(peer_id)[4:]), access_hash=0)]
+                            id=[types.InputChannel(
+                                channel_id=int(str(peer_id)[4:]),
+                                access_hash=0
+                            )]
                         )
                     )
                 else:
@@ -1360,7 +1311,7 @@ class Client(Methods, BaseClient):
                     )
 
             try:
-                return self.peers_by_id[peer_id]
+                return self.storage.get_peer_by_id(peer_id)
             except KeyError:
                 raise PeerIdInvalid
 
@@ -1435,7 +1386,7 @@ class Client(Methods, BaseClient):
         file_id = file_id or self.rnd_id()
         md5_sum = md5() if not is_big and not is_missing_part else None
 
-        session = Session(self, self.dc_id, self.auth_key, is_media=True)
+        session = Session(self, self.storage.dc_id, self.storage.auth_key, is_media=True)
         session.start()
 
         try:
@@ -1521,19 +1472,14 @@ class Client(Methods, BaseClient):
             session = self.media_sessions.get(dc_id, None)
 
             if session is None:
-                if dc_id != self.dc_id:
+                if dc_id != self.storage.dc_id:
                     exported_auth = self.send(
                         functions.auth.ExportAuthorization(
                             dc_id=dc_id
                         )
                     )
 
-                    session = Session(
+                    session = Session(self, dc_id, Auth(self, dc_id).create(), is_media=True)
-                        self,
-                        dc_id,
-                        Auth(dc_id, self.test_mode, self.ipv6, self._proxy).create(),
-                        is_media=True
-                    )
 
                     session.start()
 
@@ -1546,12 +1492,7 @@ class Client(Methods, BaseClient):
                         )
                     )
                 else:
-                    session = Session(
+                    session = Session(self, dc_id, self.storage.auth_key, is_media=True)
-                        self,
-                        dc_id,
-                        self.auth_key,
-                        is_media=True
-                    )
 
                     session.start()
 
@@ -1636,13 +1577,7 @@ class Client(Methods, BaseClient):
                     cdn_session = self.media_sessions.get(r.dc_id, None)
 
                     if cdn_session is None:
-                        cdn_session = Session(
+                        cdn_session = Session(self, r.dc_id, Auth(self, r.dc_id).create(), is_media=True, is_cdn=True)
-                            self,
-                            r.dc_id,
-                            Auth(r.dc_id, self.test_mode, self.ipv6, self._proxy).create(),
-                            is_media=True,
-                            is_cdn=True
-                        )
 
                         cdn_session.start()
 
@@ -1738,3 +1673,11 @@ class Client(Methods, BaseClient):
 
         if extensions:
             return extensions.split(" ")[0]
+
+    def export_session_string(self):
+        """Export the current session as serialized string.
+
+        Returns:
+            ``str``: The session serialized into a printable, url-safe string.
+        """
+        return self.storage.export_session_string()

pyrogram/client/ext/base_client.py

@@ -88,18 +88,10 @@ class BaseClient:
             mime_types_to_extensions[mime_type] = " ".join(extensions)
 
     def __init__(self):
-        self.is_bot = None
+        self.storage = None
-        self.dc_id = None
-        self.auth_key = None
-        self.user_id = None
-        self.date = None
 
         self.rnd_id = MsgId
 
-        self.peers_by_id = {}
-        self.peers_by_username = {}
-        self.peers_by_phone = {}
-
         self.markdown = Markdown(self)
         self.html = HTML(self)
 

pyrogram/client/ext/syncer.py

@@ -16,16 +16,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-import base64
-import json
 import logging
-import os
-import shutil
 import time
 from threading import Thread, Event, Lock
 
-from . import utils
-
 log = logging.getLogger(__name__)
 
 
@@ -81,48 +75,13 @@ class Syncer:
 
     @classmethod
     def sync(cls, client):
-        temporary = os.path.join(client.workdir, "{}.sync".format(client.session_name))
-        persistent = os.path.join(client.workdir, "{}.session".format(client.session_name))
-
         try:
-            auth_key = base64.b64encode(client.auth_key).decode()
+            start = time.time()
-            auth_key = [auth_key[i: i + 43] for i in range(0, len(auth_key), 43)]
+            client.storage.save()
-
-            data = dict(
-                dc_id=client.dc_id,
-                test_mode=client.test_mode,
-                auth_key=auth_key,
-                user_id=client.user_id,
-                date=int(time.time()),
-                is_bot=bool(client.is_bot),
-                peers_by_id={
-                    k: getattr(v, "access_hash", None)
-                    for k, v in client.peers_by_id.copy().items()
-                },
-                peers_by_username={
-                    k: utils.get_peer_id(v)
-                    for k, v in client.peers_by_username.copy().items()
-                },
-                peers_by_phone={
-                    k: utils.get_peer_id(v)
-                    for k, v in client.peers_by_phone.copy().items()
-                }
-            )
-
-            os.makedirs(client.workdir, exist_ok=True)
-
-            with open(temporary, "w", encoding="utf-8") as f:
-                json.dump(data, f, indent=4)
-
-                f.flush()
-                os.fsync(f.fileno())
         except Exception as e:
             log.critical(e, exc_info=True)
         else:
-            shutil.move(temporary, persistent)
+            log.info('Synced "{}" in {:.6} ms'.format(
-            log.info("Synced {}".format(client.session_name))
+                client.storage.name,
-        finally:
+                (time.time() - start) * 1000
-            try:
+            ))
-                os.remove(temporary)
-            except OSError:
-                pass

pyrogram/client/ext/utils.py

@@ -18,16 +18,16 @@
 
 import base64
 import struct
-from base64 import b64decode, b64encode
 from typing import Union, List
 
 import pyrogram
+
 from . import BaseClient
 from ...api import types
 
 
 def decode(s: str) -> bytes:
-    s = b64decode(s + "=" * (-len(s) % 4), "-_")
+    s = base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))
     r = b""
 
     assert s[-1] == 2
@@ -59,7 +59,7 @@ def encode(s: bytes) -> str:
 
             r += bytes([i])
 
-    return b64encode(r, b"-_").decode().rstrip("=")
+    return base64.urlsafe_b64encode(r).decode().rstrip("=")
 
 
 def get_peer_id(input_peer) -> int:

pyrogram/client/methods/contacts/get_contacts.py

@@ -46,5 +46,4 @@ class GetContacts(BaseClient):
                 log.warning("get_contacts flood: waiting {} seconds".format(e.x))
                 time.sleep(e.x)
             else:
-                log.info("Total contacts: {}".format(len(self.peers_by_phone)))
                 return pyrogram.List(pyrogram.User._parse(self, user) for user in contacts.users)

pyrogram/client/storage/__init__.py

@@ -0,0 +1,21 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# 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 .memory_storage import MemoryStorage
+from .file_storage import FileStorage
+from .storage import Storage

pyrogram/client/storage/file_storage.py

@@ -0,0 +1,102 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+import base64
+import json
+import logging
+import os
+import sqlite3
+from pathlib import Path
+from sqlite3 import DatabaseError
+from threading import Lock
+from typing import Union
+
+from .memory_storage import MemoryStorage
+
+log = logging.getLogger(__name__)
+
+
+class FileStorage(MemoryStorage):
+    FILE_EXTENSION = ".session"
+
+    def __init__(self, name: str, workdir: Path):
+        super().__init__(name)
+
+        self.workdir = workdir
+        self.database = workdir / (self.name + self.FILE_EXTENSION)
+        self.conn = None  # type: sqlite3.Connection
+        self.lock = Lock()
+
+    # noinspection PyAttributeOutsideInit
+    def migrate_from_json(self, path: Union[str, Path]):
+        log.warning("JSON session storage detected! Pyrogram will now convert it into an SQLite session storage...")
+
+        with open(path, encoding="utf-8") as f:
+            json_session = json.load(f)
+
+        os.remove(path)
+
+        self.open()
+
+        self.dc_id = json_session["dc_id"]
+        self.test_mode = json_session["test_mode"]
+        self.auth_key = base64.b64decode("".join(json_session["auth_key"]))
+        self.user_id = json_session["user_id"]
+        self.date = json_session.get("date", 0)
+        self.is_bot = json_session.get("is_bot", False)
+
+        peers_by_id = json_session.get("peers_by_id", {})
+        peers_by_phone = json_session.get("peers_by_phone", {})
+
+        peers = {}
+
+        for k, v in peers_by_id.items():
+            if v is None:
+                type_ = "group"
+            elif k.startswith("-100"):
+                type_ = "channel"
+            else:
+                type_ = "user"
+
+            peers[int(k)] = [int(k), int(v) if v is not None else None, type_, None, None]
+
+        for k, v in peers_by_phone.items():
+            peers[v][4] = k
+
+        # noinspection PyTypeChecker
+        self.update_peers(peers.values())
+
+        log.warning("Done! The session has been successfully converted from JSON to SQLite storage")
+
+    def open(self):
+        database_exists = os.path.isfile(self.database)
+
+        self.conn = sqlite3.connect(
+            str(self.database),
+            timeout=1,
+            check_same_thread=False
+        )
+
+        try:
+            if not database_exists:
+                self.create()
+
+            with self.conn:
+                self.conn.execute("VACUUM")
+        except DatabaseError:
+            self.migrate_from_json(self.database)

pyrogram/client/storage/memory_storage.py

@@ -0,0 +1,241 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+import base64
+import inspect
+import logging
+import sqlite3
+import struct
+import time
+from pathlib import Path
+from threading import Lock
+from typing import List, Tuple
+
+from pyrogram.api import types
+from pyrogram.client.storage.storage import Storage
+
+log = logging.getLogger(__name__)
+
+
+class MemoryStorage(Storage):
+    SCHEMA_VERSION = 1
+    USERNAME_TTL = 8 * 60 * 60
+    SESSION_STRING_FMT = ">B?256sI?"
+    SESSION_STRING_SIZE = 351
+
+    def __init__(self, name: str):
+        super().__init__(name)
+
+        self.conn = None  # type: sqlite3.Connection
+        self.lock = Lock()
+
+    def create(self):
+        with self.lock, self.conn:
+            with open(Path(__file__).parent / "schema.sql", "r") as schema:
+                self.conn.executescript(schema.read())
+
+            self.conn.execute(
+                "INSERT INTO version VALUES (?)",
+                (self.SCHEMA_VERSION,)
+            )
+
+            self.conn.execute(
+                "INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?)",
+                (1, None, None, 0, None, None)
+            )
+
+    def _import_session_string(self, string_session: str):
+        decoded = base64.urlsafe_b64decode(string_session + "=" * (-len(string_session) % 4))
+        return struct.unpack(self.SESSION_STRING_FMT, decoded)
+
+    def export_session_string(self):
+        packed = struct.pack(
+            self.SESSION_STRING_FMT,
+            self.dc_id,
+            self.test_mode,
+            self.auth_key,
+            self.user_id,
+            self.is_bot
+        )
+
+        return base64.urlsafe_b64encode(packed).decode().rstrip("=")
+
+    # noinspection PyAttributeOutsideInit
+    def open(self):
+        self.conn = sqlite3.connect(":memory:", check_same_thread=False)
+        self.create()
+
+        if self.name != ":memory:":
+            imported_session_string = self._import_session_string(self.name)
+
+            self.dc_id, self.test_mode, self.auth_key, self.user_id, self.is_bot = imported_session_string
+            self.date = 0
+
+        self.name = ":memory:" + str(self.user_id or "<unknown>")
+
+    # noinspection PyAttributeOutsideInit
+    def save(self):
+        self.date = int(time.time())
+
+        with self.lock:
+            self.conn.commit()
+
+    def close(self):
+        with self.lock:
+            self.conn.close()
+
+    def update_peers(self, peers: List[Tuple[int, int, str, str, str]]):
+        with self.lock:
+            self.conn.executemany(
+                "REPLACE INTO peers (id, access_hash, type, username, phone_number)"
+                "VALUES (?, ?, ?, ?, ?)",
+                peers
+            )
+
+    def clear_peers(self):
+        with self.lock, self.conn:
+            self.conn.execute(
+                "DELETE FROM peers"
+            )
+
+    @staticmethod
+    def _get_input_peer(peer_id: int, access_hash: int, peer_type: str):
+        if peer_type in ["user", "bot"]:
+            return types.InputPeerUser(
+                user_id=peer_id,
+                access_hash=access_hash
+            )
+
+        if peer_type == "group":
+            return types.InputPeerChat(
+                chat_id=-peer_id
+            )
+
+        if peer_type in ["channel", "supergroup"]:
+            return types.InputPeerChannel(
+                channel_id=int(str(peer_id)[4:]),
+                access_hash=access_hash
+            )
+
+        raise ValueError("Invalid peer type")
+
+    def get_peer_by_id(self, peer_id: int):
+        r = self.conn.execute(
+            "SELECT id, access_hash, type FROM peers WHERE id = ?",
+            (peer_id,)
+        ).fetchone()
+
+        if r is None:
+            raise KeyError("ID not found")
+
+        return self._get_input_peer(*r)
+
+    def get_peer_by_username(self, username: str):
+        r = self.conn.execute(
+            "SELECT id, access_hash, type, last_update_on FROM peers WHERE username = ?",
+            (username,)
+        ).fetchone()
+
+        if r is None:
+            raise KeyError("Username not found")
+
+        if abs(time.time() - r[3]) > self.USERNAME_TTL:
+            raise KeyError("Username expired")
+
+        return self._get_input_peer(*r[:3])
+
+    def get_peer_by_phone_number(self, phone_number: str):
+        r = self.conn.execute(
+            "SELECT id, access_hash, type FROM peers WHERE phone_number = ?",
+            (phone_number,)
+        ).fetchone()
+
+        if r is None:
+            raise KeyError("Phone number not found")
+
+        return self._get_input_peer(*r)
+
+    @property
+    def peers_count(self):
+        return self.conn.execute(
+            "SELECT COUNT(*) FROM peers"
+        ).fetchone()[0]
+
+    def _get(self):
+        attr = inspect.stack()[1].function
+
+        return self.conn.execute(
+            "SELECT {} FROM sessions".format(attr)
+        ).fetchone()[0]
+
+    def _set(self, value):
+        attr = inspect.stack()[1].function
+
+        with self.lock, self.conn:
+            self.conn.execute(
+                "UPDATE sessions SET {} = ?".format(attr),
+                (value,)
+            )
+
+    @property
+    def dc_id(self):
+        return self._get()
+
+    @dc_id.setter
+    def dc_id(self, value):
+        self._set(value)
+
+    @property
+    def test_mode(self):
+        return self._get()
+
+    @test_mode.setter
+    def test_mode(self, value):
+        self._set(value)
+
+    @property
+    def auth_key(self):
+        return self._get()
+
+    @auth_key.setter
+    def auth_key(self, value):
+        self._set(value)
+
+    @property
+    def date(self):
+        return self._get()
+
+    @date.setter
+    def date(self, value):
+        self._set(value)
+
+    @property
+    def user_id(self):
+        return self._get()
+
+    @user_id.setter
+    def user_id(self, value):
+        self._set(value)
+
+    @property
+    def is_bot(self):
+        return self._get()
+
+    @is_bot.setter
+    def is_bot(self, value):
+        self._set(value)

pyrogram/client/storage/schema.sql

@@ -0,0 +1,34 @@
+CREATE TABLE sessions (
+    dc_id     INTEGER PRIMARY KEY,
+    test_mode INTEGER,
+    auth_key  BLOB,
+    date      INTEGER NOT NULL,
+    user_id   INTEGER,
+    is_bot    INTEGER
+);
+
+CREATE TABLE peers (
+    id             INTEGER PRIMARY KEY,
+    access_hash    INTEGER,
+    type           INTEGER NOT NULL,
+    username       TEXT,
+    phone_number   TEXT,
+    last_update_on INTEGER NOT NULL DEFAULT (CAST(STRFTIME('%s', 'now') AS INTEGER))
+);
+
+CREATE TABLE version (
+    number INTEGER PRIMARY KEY
+);
+
+CREATE INDEX idx_peers_id ON peers (id);
+CREATE INDEX idx_peers_username ON peers (username);
+CREATE INDEX idx_peers_phone_number ON peers (phone_number);
+
+CREATE TRIGGER trg_peers_last_update_on
+    AFTER UPDATE
+    ON peers
+    BEGIN
+        UPDATE peers
+        SET last_update_on = CAST(STRFTIME('%s', 'now') AS INTEGER)
+        WHERE id = NEW.id;
+    END;

pyrogram/client/storage/storage.py

@@ -0,0 +1,98 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+
+class Storage:
+    def __init__(self, name: str):
+        self.name = name
+
+    def open(self):
+        raise NotImplementedError
+
+    def save(self):
+        raise NotImplementedError
+
+    def close(self):
+        raise NotImplementedError
+
+    def update_peers(self, peers):
+        raise NotImplementedError
+
+    def get_peer_by_id(self, peer_id):
+        raise NotImplementedError
+
+    def get_peer_by_username(self, username):
+        raise NotImplementedError
+
+    def get_peer_by_phone_number(self, phone_number):
+        raise NotImplementedError
+
+    def export_session_string(self):
+        raise NotImplementedError
+
+    @property
+    def peers_count(self):
+        raise NotImplementedError
+
+    @property
+    def dc_id(self):
+        raise NotImplementedError
+
+    @dc_id.setter
+    def dc_id(self, value):
+        raise NotImplementedError
+
+    @property
+    def test_mode(self):
+        raise NotImplementedError
+
+    @test_mode.setter
+    def test_mode(self, value):
+        raise NotImplementedError
+
+    @property
+    def auth_key(self):
+        raise NotImplementedError
+
+    @auth_key.setter
+    def auth_key(self, value):
+        raise NotImplementedError
+
+    @property
+    def date(self):
+        raise NotImplementedError
+
+    @date.setter
+    def date(self, value):
+        raise NotImplementedError
+
+    @property
+    def user_id(self):
+        raise NotImplementedError
+
+    @user_id.setter
+    def user_id(self, value):
+        raise NotImplementedError
+
+    @property
+    def is_bot(self):
+        raise NotImplementedError
+
+    @is_bot.setter
+    def is_bot(self, value):
+        raise NotImplementedError

pyrogram/session/auth.py

@@ -22,10 +22,12 @@ from hashlib import sha1
 from io import BytesIO
 from os import urandom
 
+import pyrogram
 from pyrogram.api import functions, types
 from pyrogram.api.core import TLObject, Long, Int
 from pyrogram.connection import Connection
 from pyrogram.crypto import AES, RSA, Prime
+
 from .internals import MsgId
 
 log = logging.getLogger(__name__)
@@ -34,11 +36,11 @@ log = logging.getLogger(__name__)
 class Auth:
     MAX_RETRIES = 5
 
-    def __init__(self, dc_id: int, test_mode: bool, ipv6: bool, proxy: dict):
+    def __init__(self, client: "pyrogram.Client", dc_id: int):
         self.dc_id = dc_id
-        self.test_mode = test_mode
+        self.test_mode = client.storage.test_mode
-        self.ipv6 = ipv6
+        self.ipv6 = client.ipv6
-        self.proxy = proxy
+        self.proxy = client.proxy
 
         self.connection = None
 

pyrogram/session/session.py

@@ -34,6 +34,7 @@ from pyrogram.api.core import Message, TLObject, MsgContainer, Long, FutureSalt,
 from pyrogram.connection import Connection
 from pyrogram.crypto import AES, KDF
 from pyrogram.errors import RPCError, InternalServerError, AuthKeyDuplicated
+
 from .internals import MsgId, MsgFactory
 
 log = logging.getLogger(__name__)
@@ -70,12 +71,14 @@ class Session:
         64: "[64] invalid container"
     }
 
-    def __init__(self,
+    def __init__(
-                 client: pyrogram,
+        self,
-                 dc_id: int,
+        client: pyrogram,
-                 auth_key: bytes,
+        dc_id: int,
-                 is_media: bool = False,
+        auth_key: bytes,
-                 is_cdn: bool = False):
+        is_media: bool = False,
+        is_cdn: bool = False
+    ):
         if not Session.notice_displayed:
             print("Pyrogram v{}, {}".format(__version__, __copyright__))
             print("Licensed under the terms of the " + __license__, end="\n\n")
@@ -113,7 +116,12 @@ class Session:
 
     def start(self):
         while True:
-            self.connection = Connection(self.dc_id, self.client.test_mode, self.client.ipv6, self.client.proxy)
+            self.connection = Connection(
+                self.dc_id,
+                self.client.storage.test_mode,
+                self.client.ipv6,
+                self.client.proxy
+            )
 
             try:
                 self.connection.connect()

setup.py

@@ -168,7 +168,8 @@ setup(
     python_requires="~=3.4",
     packages=find_packages(exclude=["compiler*"]),
     package_data={
-        "pyrogram.client.ext": ["mime.types"]
+        "pyrogram.client.ext": ["mime.types"],
+        "pyrogram.client.storage": ["schema.sql"]
     },
     zip_safe=False,
     install_requires=requires,