diff --git a/MANIFEST.in b/MANIFEST.in index 97d04588..79c547f6 100644 --- a/MANIFEST.in +++ b/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 diff --git a/docs/source/api/methods.rst b/docs/source/api/methods.rst index 4a3eefd8..2a08b37f 100644 --- a/docs/source/api/methods.rst +++ b/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() diff --git a/docs/source/glossary.rst b/docs/source/glossary.rst index bcb1193c..d5a1bffd 100644 --- a/docs/source/glossary.rst +++ b/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. diff --git a/docs/source/index.rst b/docs/source/index.rst index 0bc175ee..b9682827 100644 --- a/docs/source/index.rst +++ b/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 diff --git a/docs/source/topics/storage-engines.rst b/docs/source/topics/storage-engines.rst new file mode 100644 index 00000000..933a21b3 --- /dev/null +++ b/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 `_, 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()) + diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 9001a37e..2d18d178 100644 --- a/pyrogram/client/client.py +++ b/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 . -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 - to save a file to disk that stores details needed for reconnecting without asking again for credentials. + 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. 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.config_file = config_file + self.workdir = Path(workdir) + 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, - self.dc_id, - self.auth_key - ) + self.session = Session(self, self.storage.dc_id, self.storage.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.is_bot: + if not self.storage.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: - self.peers_by_username = {} - self.peers_by_phone = {} - + if abs(now - self.storage.date) > Client.OFFLINE_SLEEP: 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.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create() - - self.session = Session( - self, - self.dc_id, - self.auth_key - ) + self.storage.dc_id = e.x + 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.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.dc_id, - self.test_mode, - self.ipv6, - self._proxy - ).create() - - self.session = Session( - self, - self.dc_id, - self.auth_key - ) + self.session = Session(self, self.storage.dc_id, self.storage.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: - if isinstance(entity, types.User): - user_id = entity.id + for peer in peers: + username = None + 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: - self.peers_by_phone[phone] = input_peer + username = getattr(peer, "username", None) - if isinstance(entity, (types.Chat, types.ChatForbidden)): - chat_id = entity.id - peer_id = -chat_id - - 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 peer.broadcast: + peer_type = "channel" + else: + peer_type = "supergroup" 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: - 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) + self.storage.open() - for k, v in s.get("peers_by_id", {}).items(): - self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v) + session_empty = any([ + 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(): - peer = self.peers_by_id.get(v, None) + if session_empty: + self.storage.dc_id = 1 + self.storage.date = 0 - if peer: - self.peers_by_username[k] = peer - - for k, v in s.get("peers_by_phone", {}).items(): - peer = self.peers_by_id.get(v, None) - - if peer: - self.peers_by_phone[k] = peer + self.storage.test_mode = self.test_mode + self.storage.auth_key = Auth(self, self.storage.dc_id).create() + self.storage.user_id = None + self.storage.is_bot = None 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( - self, - dc_id, - Auth(dc_id, self.test_mode, self.ipv6, self._proxy).create(), - is_media=True - ) + session = Session(self, dc_id, Auth(self, dc_id).create(), is_media=True) session.start() @@ -1546,12 +1492,7 @@ class Client(Methods, BaseClient): ) ) else: - session = Session( - self, - dc_id, - self.auth_key, - is_media=True - ) + session = Session(self, dc_id, self.storage.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( - self, - r.dc_id, - Auth(r.dc_id, self.test_mode, self.ipv6, self._proxy).create(), - is_media=True, - is_cdn=True - ) + cdn_session = Session(self, r.dc_id, Auth(self, r.dc_id).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() diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py index def290e6..88623f4a 100644 --- a/pyrogram/client/ext/base_client.py +++ b/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.dc_id = None - self.auth_key = None - self.user_id = None - self.date = None + self.storage = 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) diff --git a/pyrogram/client/ext/syncer.py b/pyrogram/client/ext/syncer.py index c3921205..42e1f95a 100644 --- a/pyrogram/client/ext/syncer.py +++ b/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 . -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() - auth_key = [auth_key[i: i + 43] for i in range(0, len(auth_key), 43)] - - 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()) + start = time.time() + client.storage.save() except Exception as e: log.critical(e, exc_info=True) else: - shutil.move(temporary, persistent) - log.info("Synced {}".format(client.session_name)) - finally: - try: - os.remove(temporary) - except OSError: - pass + log.info('Synced "{}" in {:.6} ms'.format( + client.storage.name, + (time.time() - start) * 1000 + )) diff --git a/pyrogram/client/ext/utils.py b/pyrogram/client/ext/utils.py index fa107fab..e0a797e2 100644 --- a/pyrogram/client/ext/utils.py +++ b/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: diff --git a/pyrogram/client/methods/contacts/get_contacts.py b/pyrogram/client/methods/contacts/get_contacts.py index 8ca321dc..40cb344e 100644 --- a/pyrogram/client/methods/contacts/get_contacts.py +++ b/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) diff --git a/pyrogram/client/storage/__init__.py b/pyrogram/client/storage/__init__.py new file mode 100644 index 00000000..00d2f144 --- /dev/null +++ b/pyrogram/client/storage/__init__.py @@ -0,0 +1,21 @@ +# Pyrogram - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-2019 Dan Tès +# +# 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 . + +from .memory_storage import MemoryStorage +from .file_storage import FileStorage +from .storage import Storage diff --git a/pyrogram/client/storage/file_storage.py b/pyrogram/client/storage/file_storage.py new file mode 100644 index 00000000..ee5000c5 --- /dev/null +++ b/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 +# +# 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 . + +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) diff --git a/pyrogram/client/storage/memory_storage.py b/pyrogram/client/storage/memory_storage.py new file mode 100644 index 00000000..7eb3a7d0 --- /dev/null +++ b/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 +# +# 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 . + +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 "") + + # 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) diff --git a/pyrogram/client/storage/schema.sql b/pyrogram/client/storage/schema.sql new file mode 100644 index 00000000..1f5af6d2 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/pyrogram/client/storage/storage.py b/pyrogram/client/storage/storage.py new file mode 100644 index 00000000..e0810645 --- /dev/null +++ b/pyrogram/client/storage/storage.py @@ -0,0 +1,98 @@ +# Pyrogram - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-2019 Dan Tès +# +# 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 . + + +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 diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py index fb6e7ca3..b05b2855 100644 --- a/pyrogram/session/auth.py +++ b/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.ipv6 = ipv6 - self.proxy = proxy + self.test_mode = client.storage.test_mode + self.ipv6 = client.ipv6 + self.proxy = client.proxy self.connection = None diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index 8ef5570c..5947fc0f 100644 --- a/pyrogram/session/session.py +++ b/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, - client: pyrogram, - dc_id: int, - auth_key: bytes, - is_media: bool = False, - is_cdn: bool = False): + def __init__( + self, + client: pyrogram, + dc_id: int, + auth_key: bytes, + 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() diff --git a/setup.py b/setup.py index 146dae9e..d4255e03 100644 --- a/setup.py +++ b/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,