From 9d32b28f94e4c0f263e80e758fa530a3471cf91b Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Thu, 21 Feb 2019 20:12:11 +0300 Subject: [PATCH 01/22] Implement extendable session storage and JSON session storage --- pyrogram/client/client.py | 59 ++------- pyrogram/client/ext/base_client.py | 14 +-- pyrogram/client/ext/syncer.py | 41 +------ pyrogram/client/session_storage/__init__.py | 21 ++++ .../session_storage/base_session_storage.py | 50 ++++++++ .../session_storage/json_session_storage.py | 116 ++++++++++++++++++ .../session_storage/session_storage_mixin.py | 73 +++++++++++ 7 files changed, 278 insertions(+), 96 deletions(-) create mode 100644 pyrogram/client/session_storage/__init__.py create mode 100644 pyrogram/client/session_storage/base_session_storage.py create mode 100644 pyrogram/client/session_storage/json_session_storage.py create mode 100644 pyrogram/client/session_storage/session_storage_mixin.py diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index f62c046c..9a9f8482 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -36,7 +36,7 @@ from importlib import import_module from pathlib import Path from signal import signal, SIGINT, SIGTERM, SIGABRT from threading import Thread -from typing import Union, List +from typing import Union, List, Type from pyrogram.api import functions, types from pyrogram.api.core import Object @@ -56,6 +56,7 @@ from pyrogram.session import Auth, Session from .dispatcher import Dispatcher from .ext import utils, Syncer, BaseClient from .methods import Methods +from .session_storage import BaseSessionStorage, JsonSessionStorage, SessionDoesNotExist log = logging.getLogger(__name__) @@ -199,8 +200,9 @@ class Client(Methods, BaseClient): config_file: str = BaseClient.CONFIG_FILE, plugins: dict = None, no_updates: bool = None, - takeout: bool = None): - super().__init__() + takeout: bool = None, + session_storage_cls: Type[BaseSessionStorage] = JsonSessionStorage): + super().__init__(session_storage_cls(self)) self.session_name = session_name self.api_id = int(api_id) if api_id else None @@ -296,8 +298,8 @@ class Client(Methods, BaseClient): now = time.time() if abs(now - self.date) > Client.OFFLINE_SLEEP: - self.peers_by_username = {} - self.peers_by_phone = {} + self.peers_by_username.clear() + self.peers_by_phone.clear() self.get_initial_dialogs() self.get_contacts() @@ -1101,33 +1103,10 @@ class Client(Methods, BaseClient): 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.session_storage.load_session(self.session_name) + except SessionDoesNotExist: + log.info('Session {} was not found, initializing new one') 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) - - for k, v in s.get("peers_by_id", {}).items(): - self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v) - - for k, v in s.get("peers_by_username", {}).items(): - peer = self.peers_by_id.get(v, None) - - 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 def load_plugins(self): if self.plugins.get("enabled", False): @@ -1234,23 +1213,7 @@ class Client(Methods, BaseClient): log.warning('No plugin loaded from "{}"'.format(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 - ), - f, - indent=4 - ) + self.session_storage.save_session(self.session_name) def get_initial_dialogs_chunk(self, offset_date: int = 0): diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py index d2c348a8..87f11e23 100644 --- a/pyrogram/client/ext/base_client.py +++ b/pyrogram/client/ext/base_client.py @@ -24,9 +24,10 @@ from threading import Lock from pyrogram import __version__ from ..style import Markdown, HTML from ...session.internals import MsgId +from ..session_storage import SessionStorageMixin, BaseSessionStorage -class BaseClient: +class BaseClient(SessionStorageMixin): class StopTransmission(StopIteration): pass @@ -67,20 +68,13 @@ class BaseClient: 13: "video_note" } - def __init__(self): + def __init__(self, session_storage: BaseSessionStorage): + self.session_storage = session_storage self.bot_token = None - self.dc_id = None - self.auth_key = None - self.user_id = None - self.date = None self.rnd_id = MsgId self.channels_pts = {} - self.peers_by_id = {} - self.peers_by_username = {} - self.peers_by_phone = {} - self.markdown = Markdown(self.peers_by_id) self.html = HTML(self.peers_by_id) diff --git a/pyrogram/client/ext/syncer.py b/pyrogram/client/ext/syncer.py index e169d2a3..8930b13e 100644 --- a/pyrogram/client/ext/syncer.py +++ b/pyrogram/client/ext/syncer.py @@ -81,47 +81,12 @@ 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)) - + client.date = int(time.time()) 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()), - 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()) + client.session_storage.save_session(client.session_name, sync=True) 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 + client.session_storage.sync_cleanup(client.session_name) diff --git a/pyrogram/client/session_storage/__init__.py b/pyrogram/client/session_storage/__init__.py new file mode 100644 index 00000000..6ee92ebc --- /dev/null +++ b/pyrogram/client/session_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 .session_storage_mixin import SessionStorageMixin +from .base_session_storage import BaseSessionStorage, SessionDoesNotExist +from .json_session_storage import JsonSessionStorage diff --git a/pyrogram/client/session_storage/base_session_storage.py b/pyrogram/client/session_storage/base_session_storage.py new file mode 100644 index 00000000..75e416b4 --- /dev/null +++ b/pyrogram/client/session_storage/base_session_storage.py @@ -0,0 +1,50 @@ +# 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 abc + +import pyrogram + + +class SessionDoesNotExist(Exception): + pass + + +class BaseSessionStorage(abc.ABC): + def __init__(self, client: 'pyrogram.client.BaseClient'): + self.client = client + self.dc_id = 1 + self.test_mode = None + self.auth_key = None + self.user_id = None + self.date = 0 + self.peers_by_id = {} + self.peers_by_username = {} + self.peers_by_phone = {} + + @abc.abstractmethod + def load_session(self, name: str): + ... + + @abc.abstractmethod + def save_session(self, name: str, sync=False): + ... + + @abc.abstractmethod + def sync_cleanup(self, name: str): + ... diff --git a/pyrogram/client/session_storage/json_session_storage.py b/pyrogram/client/session_storage/json_session_storage.py new file mode 100644 index 00000000..679a21f3 --- /dev/null +++ b/pyrogram/client/session_storage/json_session_storage.py @@ -0,0 +1,116 @@ +# 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 shutil + +from ..ext import utils +from . import BaseSessionStorage, SessionDoesNotExist + + +log = logging.getLogger(__name__) + + +class JsonSessionStorage(BaseSessionStorage): + def _get_file_name(self, name: str): + if not name.endswith('.session'): + name += '.session' + return os.path.join(self.client.workdir, name) + + def load_session(self, name: str): + file_path = self._get_file_name(name) + log.info('Loading JSON session from {}'.format(file_path)) + + try: + with open(file_path, encoding='utf-8') as f: + s = json.load(f) + except FileNotFoundError: + raise SessionDoesNotExist() + + self.dc_id = s["dc_id"] + self.test_mode = s["test_mode"] + self.auth_key = base64.b64decode("".join(s["auth_key"])) # join split key + self.user_id = s["user_id"] + self.date = s.get("date", 0) + + for k, v in s.get("peers_by_id", {}).items(): + self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v) + + for k, v in s.get("peers_by_username", {}).items(): + peer = self.peers_by_id.get(v, None) + + 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 + + def save_session(self, name: str, sync=False): + file_path = self._get_file_name(name) + + if sync: + file_path += '.tmp' + + log.info('Saving JSON session to {}, sync={}'.format(file_path, sync)) + + auth_key = base64.b64encode(self.auth_key).decode() + auth_key = [auth_key[i: i + 43] for i in range(0, len(auth_key), 43)] # split key in lines of 43 chars + + os.makedirs(self.client.workdir, exist_ok=True) + + data = { + 'dc_id': self.dc_id, + 'test_mode': self.test_mode, + 'auth_key': auth_key, + 'user_id': self.user_id, + 'date': self.date, + 'peers_by_id': { + k: getattr(v, "access_hash", None) + for k, v in self.peers_by_id.copy().items() + }, + 'peers_by_username': { + k: utils.get_peer_id(v) + for k, v in self.peers_by_username.copy().items() + }, + 'peers_by_phone': { + k: utils.get_peer_id(v) + for k, v in self.peers_by_phone.copy().items() + } + } + + with open(file_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=4) + + f.flush() + os.fsync(f.fileno()) + + # execution won't be here if an error has occurred earlier + if sync: + shutil.move(file_path, self._get_file_name(name)) + + def sync_cleanup(self, name: str): + try: + os.remove(self._get_file_name(name) + '.tmp') + except OSError: + pass diff --git a/pyrogram/client/session_storage/session_storage_mixin.py b/pyrogram/client/session_storage/session_storage_mixin.py new file mode 100644 index 00000000..bfe9a590 --- /dev/null +++ b/pyrogram/client/session_storage/session_storage_mixin.py @@ -0,0 +1,73 @@ +# 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 typing import Dict + + +class SessionStorageMixin: + @property + def dc_id(self) -> int: + return self.session_storage.dc_id + + @dc_id.setter + def dc_id(self, val): + self.session_storage.dc_id = val + + @property + def test_mode(self) -> bool: + return self.session_storage.test_mode + + @test_mode.setter + def test_mode(self, val): + self.session_storage.test_mode = val + + @property + def auth_key(self) -> bytes: + return self.session_storage.auth_key + + @auth_key.setter + def auth_key(self, val): + self.session_storage.auth_key = val + + @property + def user_id(self): + return self.session_storage.user_id + + @user_id.setter + def user_id(self, val) -> int: + self.session_storage.user_id = val + + @property + def date(self) -> int: + return self.session_storage.date + + @date.setter + def date(self, val): + self.session_storage.date = val + + @property + def peers_by_id(self) -> Dict[str, int]: + return self.session_storage.peers_by_id + + @property + def peers_by_username(self) -> Dict[str, int]: + return self.session_storage.peers_by_username + + @property + def peers_by_phone(self) -> Dict[str, int]: + return self.session_storage.peers_by_phone From 431a983d5b66522604f0685ef078d40735cea64c Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Thu, 21 Feb 2019 21:18:53 +0300 Subject: [PATCH 02/22] Fix logging and cleanup imports in client.py --- pyrogram/client/client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 9a9f8482..0e8d5554 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -16,9 +16,7 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . -import base64 import binascii -import json import logging import math import mimetypes @@ -1105,7 +1103,10 @@ class Client(Methods, BaseClient): try: self.session_storage.load_session(self.session_name) except SessionDoesNotExist: - log.info('Session {} was not found, initializing new one') + session_name = self.session_name[:32] + if session_name != self.session_name: + session_name += '...' + log.info('Could not load session "{}", initializing new one'.format(self.session_name)) self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create() def load_plugins(self): From b04cf9ec9297ce0884943879b1fac07fa7e2933f Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Thu, 21 Feb 2019 21:43:57 +0300 Subject: [PATCH 03/22] Add string session storage --- pyrogram/client/session_storage/__init__.py | 1 + .../session_storage/string_session_storage.py | 38 +++++++++++++++++++ 2 files changed, 39 insertions(+) create mode 100644 pyrogram/client/session_storage/string_session_storage.py diff --git a/pyrogram/client/session_storage/__init__.py b/pyrogram/client/session_storage/__init__.py index 6ee92ebc..ced103ce 100644 --- a/pyrogram/client/session_storage/__init__.py +++ b/pyrogram/client/session_storage/__init__.py @@ -19,3 +19,4 @@ from .session_storage_mixin import SessionStorageMixin from .base_session_storage import BaseSessionStorage, SessionDoesNotExist from .json_session_storage import JsonSessionStorage +from .string_session_storage import StringSessionStorage diff --git a/pyrogram/client/session_storage/string_session_storage.py b/pyrogram/client/session_storage/string_session_storage.py new file mode 100644 index 00000000..9b6ebf0e --- /dev/null +++ b/pyrogram/client/session_storage/string_session_storage.py @@ -0,0 +1,38 @@ +import base64 +import binascii +import struct + +from . import BaseSessionStorage, SessionDoesNotExist + + +def StringSessionStorage(print_session: bool = False): + class StringSessionStorageClass(BaseSessionStorage): + """ + Packs session data as following (forcing little-endian byte order): + Char dc_id (1 byte, unsigned) + Boolean test_mode (1 byte) + Long long user_id (8 bytes, signed) + Bytes auth_key (256 bytes) + + Uses Base64 encoding for printable representation + """ + PACK_FORMAT = ' Date: Fri, 22 Feb 2019 00:03:58 +0300 Subject: [PATCH 04/22] Refactor session storages: use session_name arg to detect storage type --- pyrogram/client/client.py | 34 +++++++----- pyrogram/client/ext/syncer.py | 4 +- pyrogram/client/session_storage/__init__.py | 4 +- .../session_storage/base_session_storage.py | 17 ++++-- .../session_storage/json_session_storage.py | 14 ++--- .../session_storage/string_session_storage.py | 53 +++++++++---------- 6 files changed, 71 insertions(+), 55 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 0e8d5554..f17a054b 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -49,12 +49,15 @@ from pyrogram.api.errors import ( from pyrogram.client.handlers import DisconnectHandler from pyrogram.client.handlers.handler import Handler from pyrogram.client.methods.password.utils import compute_check +from pyrogram.client.session_storage import BaseSessionConfig from pyrogram.crypto import AES from pyrogram.session import Auth, Session from .dispatcher import Dispatcher from .ext import utils, Syncer, BaseClient from .methods import Methods -from .session_storage import BaseSessionStorage, JsonSessionStorage, SessionDoesNotExist +from .session_storage import SessionDoesNotExist +from .session_storage.json_session_storage import JsonSessionStorage +from .session_storage.string_session_storage import StringSessionStorage log = logging.getLogger(__name__) @@ -176,7 +179,7 @@ class Client(Methods, BaseClient): """ def __init__(self, - session_name: str, + session_name: Union[str, BaseSessionConfig], api_id: Union[int, str] = None, api_hash: str = None, app_version: str = None, @@ -198,11 +201,21 @@ class Client(Methods, BaseClient): config_file: str = BaseClient.CONFIG_FILE, plugins: dict = None, no_updates: bool = None, - takeout: bool = None, - session_storage_cls: Type[BaseSessionStorage] = JsonSessionStorage): - super().__init__(session_storage_cls(self)) + takeout: bool = None): - self.session_name = session_name + if isinstance(session_name, str): + if session_name.startswith(':'): + session_storage = StringSessionStorage(self, session_name) + else: + session_storage = JsonSessionStorage(self, session_name) + elif isinstance(session_name, BaseSessionConfig): + session_storage = session_name.session_storage_cls(self, session_name) + else: + raise RuntimeError('Wrong session_name passed, expected str or BaseSessionConfig subclass') + + super().__init__(session_storage) + + self.session_name = str(session_name) # TODO: build correct session name self.api_id = int(api_id) if api_id else None self.api_hash = api_hash self.app_version = app_version @@ -1101,12 +1114,9 @@ class Client(Methods, BaseClient): def load_session(self): try: - self.session_storage.load_session(self.session_name) + self.session_storage.load_session() except SessionDoesNotExist: - session_name = self.session_name[:32] - if session_name != self.session_name: - session_name += '...' - log.info('Could not load session "{}", initializing new one'.format(self.session_name)) + log.info('Could not load session "{}", initiate new one'.format(self.session_name)) self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create() def load_plugins(self): @@ -1214,7 +1224,7 @@ class Client(Methods, BaseClient): log.warning('No plugin loaded from "{}"'.format(root)) def save_session(self): - self.session_storage.save_session(self.session_name) + self.session_storage.save_session() def get_initial_dialogs_chunk(self, offset_date: int = 0): diff --git a/pyrogram/client/ext/syncer.py b/pyrogram/client/ext/syncer.py index 8930b13e..70955624 100644 --- a/pyrogram/client/ext/syncer.py +++ b/pyrogram/client/ext/syncer.py @@ -83,10 +83,10 @@ class Syncer: def sync(cls, client): client.date = int(time.time()) try: - client.session_storage.save_session(client.session_name, sync=True) + client.session_storage.save_session(sync=True) except Exception as e: log.critical(e, exc_info=True) else: log.info("Synced {}".format(client.session_name)) finally: - client.session_storage.sync_cleanup(client.session_name) + client.session_storage.sync_cleanup() diff --git a/pyrogram/client/session_storage/__init__.py b/pyrogram/client/session_storage/__init__.py index ced103ce..611ec9b7 100644 --- a/pyrogram/client/session_storage/__init__.py +++ b/pyrogram/client/session_storage/__init__.py @@ -17,6 +17,4 @@ # along with Pyrogram. If not, see . from .session_storage_mixin import SessionStorageMixin -from .base_session_storage import BaseSessionStorage, SessionDoesNotExist -from .json_session_storage import JsonSessionStorage -from .string_session_storage import StringSessionStorage +from .base_session_storage import BaseSessionStorage, BaseSessionConfig, SessionDoesNotExist diff --git a/pyrogram/client/session_storage/base_session_storage.py b/pyrogram/client/session_storage/base_session_storage.py index 75e416b4..a5c879f1 100644 --- a/pyrogram/client/session_storage/base_session_storage.py +++ b/pyrogram/client/session_storage/base_session_storage.py @@ -17,6 +17,7 @@ # along with Pyrogram. If not, see . import abc +from typing import Type import pyrogram @@ -26,8 +27,9 @@ class SessionDoesNotExist(Exception): class BaseSessionStorage(abc.ABC): - def __init__(self, client: 'pyrogram.client.BaseClient'): + def __init__(self, client: 'pyrogram.client.BaseClient', session_data): self.client = client + self.session_data = session_data self.dc_id = 1 self.test_mode = None self.auth_key = None @@ -38,13 +40,20 @@ class BaseSessionStorage(abc.ABC): self.peers_by_phone = {} @abc.abstractmethod - def load_session(self, name: str): + def load_session(self): ... @abc.abstractmethod - def save_session(self, name: str, sync=False): + def save_session(self, sync=False): ... @abc.abstractmethod - def sync_cleanup(self, name: str): + def sync_cleanup(self): + ... + + +class BaseSessionConfig(abc.ABC): + @property + @abc.abstractmethod + def session_storage_cls(self) -> Type[BaseSessionStorage]: ... diff --git a/pyrogram/client/session_storage/json_session_storage.py b/pyrogram/client/session_storage/json_session_storage.py index 679a21f3..f41091af 100644 --- a/pyrogram/client/session_storage/json_session_storage.py +++ b/pyrogram/client/session_storage/json_session_storage.py @@ -35,8 +35,8 @@ class JsonSessionStorage(BaseSessionStorage): name += '.session' return os.path.join(self.client.workdir, name) - def load_session(self, name: str): - file_path = self._get_file_name(name) + def load_session(self): + file_path = self._get_file_name(self.session_data) log.info('Loading JSON session from {}'.format(file_path)) try: @@ -66,8 +66,8 @@ class JsonSessionStorage(BaseSessionStorage): if peer: self.peers_by_phone[k] = peer - def save_session(self, name: str, sync=False): - file_path = self._get_file_name(name) + def save_session(self, sync=False): + file_path = self._get_file_name(self.session_data) if sync: file_path += '.tmp' @@ -107,10 +107,10 @@ class JsonSessionStorage(BaseSessionStorage): # execution won't be here if an error has occurred earlier if sync: - shutil.move(file_path, self._get_file_name(name)) + shutil.move(file_path, self._get_file_name(self.session_data)) - def sync_cleanup(self, name: str): + def sync_cleanup(self): try: - os.remove(self._get_file_name(name) + '.tmp') + os.remove(self._get_file_name(self.session_data) + '.tmp') except OSError: pass diff --git a/pyrogram/client/session_storage/string_session_storage.py b/pyrogram/client/session_storage/string_session_storage.py index 9b6ebf0e..c01a2b35 100644 --- a/pyrogram/client/session_storage/string_session_storage.py +++ b/pyrogram/client/session_storage/string_session_storage.py @@ -5,34 +5,33 @@ import struct from . import BaseSessionStorage, SessionDoesNotExist -def StringSessionStorage(print_session: bool = False): - class StringSessionStorageClass(BaseSessionStorage): - """ - Packs session data as following (forcing little-endian byte order): - Char dc_id (1 byte, unsigned) - Boolean test_mode (1 byte) - Long long user_id (8 bytes, signed) - Bytes auth_key (256 bytes) +class StringSessionStorage(BaseSessionStorage): + """ + Packs session data as following (forcing little-endian byte order): + Char dc_id (1 byte, unsigned) + Boolean test_mode (1 byte) + Long long user_id (8 bytes, signed) + Bytes auth_key (256 bytes) - Uses Base64 encoding for printable representation - """ - PACK_FORMAT = ' Date: Fri, 22 Feb 2019 01:34:08 +0300 Subject: [PATCH 05/22] Add bot_token argument (closes #123) --- pyrogram/client/client.py | 25 ++++++++++++++++++++----- pyrogram/client/ext/base_client.py | 2 +- pyrogram/client/ext/syncer.py | 1 + 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index f62c046c..da2ddc5b 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -29,6 +29,7 @@ import struct import tempfile import threading import time +import warnings from configparser import ConfigParser from datetime import datetime from hashlib import sha256, md5 @@ -67,9 +68,8 @@ class Client(Methods, BaseClient): Args: session_name (``str``): - Name to uniquely identify a session of either a User or a Bot. - For Users: pass a string of your choice, e.g.: "my_main_account". - For Bots: pass your Bot API token, e.g.: "123456:ABC-DEF1234ghIkl-zyx57W2v1u123ew11" + Name to uniquely identify a session of either a User or a Bot, e.g.: "my_main_account". + You still can use bot token here, but it will be deprecated in next release. Note: as long as a valid User session file exists, Pyrogram won't ask you again to input your phone number. api_id (``int``, *optional*): @@ -144,6 +144,10 @@ class Client(Methods, BaseClient): a new Telegram account in case the phone number you passed is not registered yet. Only applicable for new sessions. + bot_token (``str``, *optional*): + Pass your Bot API token to create a bot session, e.g.: "123456:ABC-DEF1234ghIkl-zyx57W2v1u123ew11" + Only applicable for new sessions. + last_name (``str``, *optional*): Same purpose as *first_name*; pass a Last Name to avoid entering it manually. It can be an empty string: "". Only applicable for new sessions. @@ -192,6 +196,7 @@ class Client(Methods, BaseClient): password: str = None, recovery_code: callable = None, force_sms: bool = False, + bot_token: str = None, first_name: str = None, last_name: str = None, workers: int = BaseClient.WORKERS, @@ -218,6 +223,7 @@ class Client(Methods, BaseClient): self.password = password self.recovery_code = recovery_code self.force_sms = force_sms + self.bot_token = bot_token self.first_name = first_name self.last_name = last_name self.workers = workers @@ -263,8 +269,13 @@ class Client(Methods, BaseClient): 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] + warnings.warn('\nYou are using a bot token as session name.\n' + 'It will be deprecated in next update, please use session file name to load ' + 'existing sessions and bot_token argument to create new sessions.', + DeprecationWarning, stacklevel=2) self.load_config() self.load_session() @@ -284,11 +295,12 @@ class Client(Methods, BaseClient): if self.bot_token is None: self.authorize_user() else: + self.is_bot = True self.authorize_bot() self.save_session() - if self.bot_token is None: + 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)) @@ -1113,6 +1125,8 @@ class Client(Methods, BaseClient): 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(): self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v) @@ -1246,7 +1260,8 @@ class Client(Methods, BaseClient): test_mode=self.test_mode, auth_key=auth_key, user_id=self.user_id, - date=self.date + date=self.date, + is_bot=self.is_bot, ), f, indent=4 diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py index d2c348a8..8ca784aa 100644 --- a/pyrogram/client/ext/base_client.py +++ b/pyrogram/client/ext/base_client.py @@ -68,7 +68,7 @@ class BaseClient: } def __init__(self): - self.bot_token = None + self.is_bot = False self.dc_id = None self.auth_key = None self.user_id = None diff --git a/pyrogram/client/ext/syncer.py b/pyrogram/client/ext/syncer.py index e169d2a3..71dc3f35 100644 --- a/pyrogram/client/ext/syncer.py +++ b/pyrogram/client/ext/syncer.py @@ -94,6 +94,7 @@ class Syncer: auth_key=auth_key, user_id=client.user_id, date=int(time.time()), + is_bot=client.is_bot, peers_by_id={ k: getattr(v, "access_hash", None) for k, v in client.peers_by_id.copy().items() From 9c4e9e166e528d2ef990bcb3f2093a877d65b642 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Fri, 22 Feb 2019 02:13:51 +0300 Subject: [PATCH 06/22] Merge #221, string sessions now work for bots too --- pyrogram/client/client.py | 17 +++++++++-------- pyrogram/client/ext/base_client.py | 1 - .../session_storage/base_session_storage.py | 1 + .../session_storage/json_session_storage.py | 2 ++ .../session_storage/session_storage_mixin.py | 8 ++++++++ .../session_storage/string_session_storage.py | 13 ++++++++++--- 6 files changed, 30 insertions(+), 12 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 42a2566a..429abab3 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -281,14 +281,15 @@ 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] - warnings.warn('\nYou are using a bot token as session name.\n' - 'It will be deprecated in next update, please use session file name to load ' - 'existing sessions and bot_token argument to create new sessions.', - DeprecationWarning, stacklevel=2) + if isinstance(self.session_storage, JsonSessionStorage): + if self.BOT_TOKEN_RE.match(self.session_storage.session_data): + self.is_bot = True + self.bot_token = self.session_storage.session_data + self.session_storage.session_data = self.session_storage.session_data.split(":")[0] + warnings.warn('\nYou are using a bot token as session name.\n' + 'It will be deprecated in next update, please use session file name to load ' + 'existing sessions and bot_token argument to create new sessions.', + DeprecationWarning, stacklevel=2) self.load_config() self.load_session() diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py index a354ba76..3f40865f 100644 --- a/pyrogram/client/ext/base_client.py +++ b/pyrogram/client/ext/base_client.py @@ -70,7 +70,6 @@ class BaseClient(SessionStorageMixin): def __init__(self, session_storage: BaseSessionStorage): self.session_storage = session_storage - self.is_bot = False self.rnd_id = MsgId self.channels_pts = {} diff --git a/pyrogram/client/session_storage/base_session_storage.py b/pyrogram/client/session_storage/base_session_storage.py index a5c879f1..92473956 100644 --- a/pyrogram/client/session_storage/base_session_storage.py +++ b/pyrogram/client/session_storage/base_session_storage.py @@ -35,6 +35,7 @@ class BaseSessionStorage(abc.ABC): self.auth_key = None self.user_id = None self.date = 0 + self.is_bot = False self.peers_by_id = {} self.peers_by_username = {} self.peers_by_phone = {} diff --git a/pyrogram/client/session_storage/json_session_storage.py b/pyrogram/client/session_storage/json_session_storage.py index f41091af..1e1e0ca4 100644 --- a/pyrogram/client/session_storage/json_session_storage.py +++ b/pyrogram/client/session_storage/json_session_storage.py @@ -50,6 +50,7 @@ class JsonSessionStorage(BaseSessionStorage): self.auth_key = base64.b64decode("".join(s["auth_key"])) # join split key self.user_id = s["user_id"] self.date = s.get("date", 0) + self.is_bot = s.get('is_bot', self.client.is_bot) for k, v in s.get("peers_by_id", {}).items(): self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v) @@ -85,6 +86,7 @@ class JsonSessionStorage(BaseSessionStorage): 'auth_key': auth_key, 'user_id': self.user_id, 'date': self.date, + 'is_bot': self.is_bot, 'peers_by_id': { k: getattr(v, "access_hash", None) for k, v in self.peers_by_id.copy().items() diff --git a/pyrogram/client/session_storage/session_storage_mixin.py b/pyrogram/client/session_storage/session_storage_mixin.py index bfe9a590..7d783ca7 100644 --- a/pyrogram/client/session_storage/session_storage_mixin.py +++ b/pyrogram/client/session_storage/session_storage_mixin.py @@ -60,6 +60,14 @@ class SessionStorageMixin: def date(self, val): self.session_storage.date = val + @property + def is_bot(self): + return self.session_storage.is_bot + + @is_bot.setter + def is_bot(self, val) -> int: + self.session_storage.is_bot = val + @property def peers_by_id(self) -> Dict[str, int]: return self.session_storage.peers_by_id diff --git a/pyrogram/client/session_storage/string_session_storage.py b/pyrogram/client/session_storage/string_session_storage.py index c01a2b35..5b1a8cc1 100644 --- a/pyrogram/client/session_storage/string_session_storage.py +++ b/pyrogram/client/session_storage/string_session_storage.py @@ -11,24 +11,31 @@ class StringSessionStorage(BaseSessionStorage): Char dc_id (1 byte, unsigned) Boolean test_mode (1 byte) Long long user_id (8 bytes, signed) + Boolean is_bot (1 byte) Bytes auth_key (256 bytes) Uses Base64 encoding for printable representation """ - PACK_FORMAT = ' Date: Fri, 22 Feb 2019 03:37:19 +0300 Subject: [PATCH 07/22] add in-memory session storage, refactor session storages, remove mixin --- pyrogram/client/client.py | 112 +++++++++--------- pyrogram/client/ext/base_client.py | 10 +- pyrogram/client/ext/syncer.py | 4 +- .../client/methods/contacts/get_contacts.py | 2 +- pyrogram/client/session_storage/__init__.py | 6 +- .../{session_storage_mixin.py => abstract.py} | 89 ++++++++++---- .../session_storage/base_session_storage.py | 60 ---------- .../{json_session_storage.py => json.py} | 65 +++++----- pyrogram/client/session_storage/memory.py | 85 +++++++++++++ .../{string_session_storage.py => string.py} | 19 +-- pyrogram/session/session.py | 3 +- 11 files changed, 267 insertions(+), 188 deletions(-) rename pyrogram/client/session_storage/{session_storage_mixin.py => abstract.py} (50%) delete mode 100644 pyrogram/client/session_storage/base_session_storage.py rename pyrogram/client/session_storage/{json_session_storage.py => json.py} (58%) create mode 100644 pyrogram/client/session_storage/memory.py rename pyrogram/client/session_storage/{string_session_storage.py => string.py} (62%) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 429abab3..42bd73d6 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -50,15 +50,15 @@ from pyrogram.api.errors import ( from pyrogram.client.handlers import DisconnectHandler from pyrogram.client.handlers.handler import Handler from pyrogram.client.methods.password.utils import compute_check -from pyrogram.client.session_storage import BaseSessionConfig from pyrogram.crypto import AES from pyrogram.session import Auth, Session from .dispatcher import Dispatcher from .ext import utils, Syncer, BaseClient from .methods import Methods -from .session_storage import SessionDoesNotExist -from .session_storage.json_session_storage import JsonSessionStorage -from .session_storage.string_session_storage import StringSessionStorage +from .session_storage import ( + SessionDoesNotExist, SessionStorage, MemorySessionStorage, JsonSessionStorage, + StringSessionStorage +) log = logging.getLogger(__name__) @@ -183,7 +183,7 @@ class Client(Methods, BaseClient): """ def __init__(self, - session_name: Union[str, BaseSessionConfig], + session_name: Union[str, SessionStorage], api_id: Union[int, str] = None, api_hash: str = None, app_version: str = None, @@ -209,14 +209,16 @@ class Client(Methods, BaseClient): takeout: bool = None): if isinstance(session_name, str): - if session_name.startswith(':'): + if session_name == ':memory:': + session_storage = MemorySessionStorage(self) + elif session_name.startswith(':'): session_storage = StringSessionStorage(self, session_name) else: session_storage = JsonSessionStorage(self, session_name) - elif isinstance(session_name, BaseSessionConfig): - session_storage = session_name.session_storage_cls(self, session_name) + elif isinstance(session_name, SessionStorage): + session_storage = session_name else: - raise RuntimeError('Wrong session_name passed, expected str or BaseSessionConfig subclass') + raise RuntimeError('Wrong session_name passed, expected str or SessionConfig subclass') super().__init__(session_storage) @@ -230,7 +232,7 @@ class Client(Methods, BaseClient): self.ipv6 = ipv6 # TODO: Make code consistent, use underscore for private/protected fields self._proxy = proxy - self.test_mode = test_mode + self.session_storage.test_mode = test_mode self.phone_number = phone_number self.phone_code = phone_code self.password = password @@ -282,10 +284,10 @@ class Client(Methods, BaseClient): raise ConnectionError("Client has already been started") if isinstance(self.session_storage, JsonSessionStorage): - if self.BOT_TOKEN_RE.match(self.session_storage.session_data): - self.is_bot = True - self.bot_token = self.session_storage.session_data - self.session_storage.session_data = self.session_storage.session_data.split(":")[0] + if self.BOT_TOKEN_RE.match(self.session_storage._session_name): + self.session_storage.is_bot = True + self.bot_token = self.session_storage._session_name + self.session_storage._session_name = self.session_storage._session_name.split(":")[0] warnings.warn('\nYou are using a bot token as session name.\n' 'It will be deprecated in next update, please use session file name to load ' 'existing sessions and bot_token argument to create new sessions.', @@ -297,33 +299,33 @@ class Client(Methods, BaseClient): self.session = Session( self, - self.dc_id, - self.auth_key + self.session_storage.dc_id, + self.session_storage.auth_key ) self.session.start() self.is_started = True try: - if self.user_id is None: + if self.session_storage.user_id is None: if self.bot_token is None: self.authorize_user() else: - self.is_bot = True + self.session_storage.is_bot = True self.authorize_bot() self.save_session() - if not self.is_bot: + if not self.session_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.clear() - self.peers_by_phone.clear() + if abs(now - self.session_storage.date) > Client.OFFLINE_SLEEP: + self.session_storage.peers_by_username.clear() + self.session_storage.peers_by_phone.clear() self.get_initial_dialogs() self.get_contacts() @@ -512,19 +514,20 @@ 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_storage.dc_id = e.x + self.session_storage.auth_key = Auth(self.session_storage.dc_id, self.session_storage.test_mode, + self.ipv6, self._proxy).create() self.session = Session( self, - self.dc_id, - self.auth_key + self.session_storage.dc_id, + self.session_storage.auth_key ) self.session.start() self.authorize_bot() else: - self.user_id = r.user.id + self.session_storage.user_id = r.user.id print("Logged in successfully as @{}".format(r.user.username)) @@ -564,19 +567,19 @@ class Client(Methods, BaseClient): except (PhoneMigrate, NetworkMigrate) as e: self.session.stop() - self.dc_id = e.x + self.session_storage.dc_id = e.x - self.auth_key = Auth( - self.dc_id, - self.test_mode, + self.session_storage.auth_key = Auth( + self.session_storage.dc_id, + self.session_storage.test_mode, self.ipv6, self._proxy ).create() self.session = Session( self, - self.dc_id, - self.auth_key + self.session_storage.dc_id, + self.session_storage.auth_key ) self.session.start() @@ -752,7 +755,7 @@ class Client(Methods, BaseClient): assert self.send(functions.help.AcceptTermsOfService(terms_of_service.id)) self.password = None - self.user_id = r.user.id + self.session_storage.user_id = r.user.id print("Logged in successfully as {}".format(r.user.first_name)) @@ -776,13 +779,13 @@ class Client(Methods, BaseClient): access_hash=access_hash ) - self.peers_by_id[user_id] = input_peer + self.session_storage.peers_by_id[user_id] = input_peer if username is not None: - self.peers_by_username[username.lower()] = input_peer + self.session_storage.peers_by_username[username.lower()] = input_peer if phone is not None: - self.peers_by_phone[phone] = input_peer + self.session_storage.peers_by_phone[phone] = input_peer if isinstance(entity, (types.Chat, types.ChatForbidden)): chat_id = entity.id @@ -792,7 +795,7 @@ class Client(Methods, BaseClient): chat_id=chat_id ) - self.peers_by_id[peer_id] = input_peer + self.session_storage.peers_by_id[peer_id] = input_peer if isinstance(entity, (types.Channel, types.ChannelForbidden)): channel_id = entity.id @@ -810,10 +813,10 @@ class Client(Methods, BaseClient): access_hash=access_hash ) - self.peers_by_id[peer_id] = input_peer + self.session_storage.peers_by_id[peer_id] = input_peer if username is not None: - self.peers_by_username[username.lower()] = input_peer + self.session_storage.peers_by_username[username.lower()] = input_peer def download_worker(self): name = threading.current_thread().name @@ -1127,10 +1130,11 @@ class Client(Methods, BaseClient): def load_session(self): try: - self.session_storage.load_session() + self.session_storage.load() except SessionDoesNotExist: log.info('Could not load session "{}", initiate new one'.format(self.session_name)) - self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create() + self.session_storage.auth_key = Auth(self.session_storage.dc_id, self.session_storage.test_mode, + self.ipv6, self._proxy).create() def load_plugins(self): if self.plugins.get("enabled", False): @@ -1237,7 +1241,7 @@ class Client(Methods, BaseClient): log.warning('No plugin loaded from "{}"'.format(root)) def save_session(self): - self.session_storage.save_session() + self.session_storage.save() def get_initial_dialogs_chunk(self, offset_date: int = 0): @@ -1257,7 +1261,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(len(self.session_storage.peers_by_id))) return r def get_initial_dialogs(self): @@ -1293,7 +1297,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.session_storage.peers_by_id[peer_id] except KeyError: if type(peer_id) is str: if peer_id in ("self", "me"): @@ -1304,17 +1308,17 @@ class Client(Methods, BaseClient): try: int(peer_id) except ValueError: - if peer_id not in self.peers_by_username: + if peer_id not in self.session_storage.peers_by_username: self.send( functions.contacts.ResolveUsername( username=peer_id ) ) - return self.peers_by_username[peer_id] + return self.session_storage.peers_by_username[peer_id] else: try: - return self.peers_by_phone[peer_id] + return self.session_storage.peers_by_phone[peer_id] except KeyError: raise PeerIdInvalid @@ -1341,7 +1345,7 @@ class Client(Methods, BaseClient): ) try: - return self.peers_by_id[peer_id] + return self.session_storage.peers_by_id[peer_id] except KeyError: raise PeerIdInvalid @@ -1411,7 +1415,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.session_storage.dc_id, self.session_storage.auth_key, is_media=True) session.start() try: @@ -1492,7 +1496,7 @@ 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.session_storage.dc_id: exported_auth = self.send( functions.auth.ExportAuthorization( dc_id=dc_id @@ -1502,7 +1506,7 @@ class Client(Methods, BaseClient): session = Session( self, dc_id, - Auth(dc_id, self.test_mode, self.ipv6, self._proxy).create(), + Auth(dc_id, self.session_storage.test_mode, self.ipv6, self._proxy).create(), is_media=True ) @@ -1520,7 +1524,7 @@ class Client(Methods, BaseClient): session = Session( self, dc_id, - self.auth_key, + self.session_storage.auth_key, is_media=True ) @@ -1588,7 +1592,7 @@ class Client(Methods, BaseClient): cdn_session = Session( self, r.dc_id, - Auth(r.dc_id, self.test_mode, self.ipv6, self._proxy).create(), + Auth(r.dc_id, self.session_storage.test_mode, self.ipv6, self._proxy).create(), is_media=True, is_cdn=True ) diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py index 3f40865f..732a600f 100644 --- a/pyrogram/client/ext/base_client.py +++ b/pyrogram/client/ext/base_client.py @@ -24,10 +24,10 @@ from threading import Lock from pyrogram import __version__ from ..style import Markdown, HTML from ...session.internals import MsgId -from ..session_storage import SessionStorageMixin, BaseSessionStorage +from ..session_storage import SessionStorage -class BaseClient(SessionStorageMixin): +class BaseClient: class StopTransmission(StopIteration): pass @@ -68,14 +68,14 @@ class BaseClient(SessionStorageMixin): 13: "video_note" } - def __init__(self, session_storage: BaseSessionStorage): + def __init__(self, session_storage: SessionStorage): self.session_storage = session_storage self.rnd_id = MsgId self.channels_pts = {} - self.markdown = Markdown(self.peers_by_id) - self.html = HTML(self.peers_by_id) + self.markdown = Markdown(self.session_storage.peers_by_id) + self.html = HTML(self.session_storage.peers_by_id) self.session = None self.media_sessions = {} diff --git a/pyrogram/client/ext/syncer.py b/pyrogram/client/ext/syncer.py index 70955624..e13212be 100644 --- a/pyrogram/client/ext/syncer.py +++ b/pyrogram/client/ext/syncer.py @@ -81,9 +81,9 @@ class Syncer: @classmethod def sync(cls, client): - client.date = int(time.time()) + client.session_storage.date = int(time.time()) try: - client.session_storage.save_session(sync=True) + client.session_storage.save(sync=True) except Exception as e: log.critical(e, exc_info=True) else: diff --git a/pyrogram/client/methods/contacts/get_contacts.py b/pyrogram/client/methods/contacts/get_contacts.py index 29b7e176..35b24592 100644 --- a/pyrogram/client/methods/contacts/get_contacts.py +++ b/pyrogram/client/methods/contacts/get_contacts.py @@ -44,5 +44,5 @@ 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))) + log.info("Total contacts: {}".format(len(self.session_storage.peers_by_phone))) return [pyrogram.User._parse(self, user) for user in contacts.users] diff --git a/pyrogram/client/session_storage/__init__.py b/pyrogram/client/session_storage/__init__.py index 611ec9b7..ad2d8900 100644 --- a/pyrogram/client/session_storage/__init__.py +++ b/pyrogram/client/session_storage/__init__.py @@ -16,5 +16,7 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . -from .session_storage_mixin import SessionStorageMixin -from .base_session_storage import BaseSessionStorage, BaseSessionConfig, SessionDoesNotExist +from .abstract import SessionStorage, SessionDoesNotExist +from .memory import MemorySessionStorage +from .json import JsonSessionStorage +from .string import StringSessionStorage diff --git a/pyrogram/client/session_storage/session_storage_mixin.py b/pyrogram/client/session_storage/abstract.py similarity index 50% rename from pyrogram/client/session_storage/session_storage_mixin.py rename to pyrogram/client/session_storage/abstract.py index 7d783ca7..e8f4441e 100644 --- a/pyrogram/client/session_storage/session_storage_mixin.py +++ b/pyrogram/client/session_storage/abstract.py @@ -16,66 +16,103 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . -from typing import Dict +import abc +from typing import Type + +import pyrogram -class SessionStorageMixin: +class SessionDoesNotExist(Exception): + pass + + +class SessionStorage(abc.ABC): + def __init__(self, client: 'pyrogram.client.BaseClient'): + self._client = client + + @abc.abstractmethod + def load(self): + ... + + @abc.abstractmethod + def save(self, sync=False): + ... + + @abc.abstractmethod + def sync_cleanup(self): + ... + @property - def dc_id(self) -> int: - return self.session_storage.dc_id + @abc.abstractmethod + def dc_id(self): + ... @dc_id.setter + @abc.abstractmethod def dc_id(self, val): - self.session_storage.dc_id = val + ... @property - def test_mode(self) -> bool: - return self.session_storage.test_mode + @abc.abstractmethod + def test_mode(self): + ... @test_mode.setter + @abc.abstractmethod def test_mode(self, val): - self.session_storage.test_mode = val + ... @property - def auth_key(self) -> bytes: - return self.session_storage.auth_key + @abc.abstractmethod + def auth_key(self): + ... @auth_key.setter + @abc.abstractmethod def auth_key(self, val): - self.session_storage.auth_key = val + ... @property + @abc.abstractmethod def user_id(self): - return self.session_storage.user_id + ... @user_id.setter - def user_id(self, val) -> int: - self.session_storage.user_id = val + @abc.abstractmethod + def user_id(self, val): + ... @property - def date(self) -> int: - return self.session_storage.date + @abc.abstractmethod + def date(self): + ... @date.setter + @abc.abstractmethod def date(self, val): - self.session_storage.date = val + ... @property + @abc.abstractmethod def is_bot(self): - return self.session_storage.is_bot + ... @is_bot.setter - def is_bot(self, val) -> int: - self.session_storage.is_bot = val + @abc.abstractmethod + def is_bot(self, val): + ... @property - def peers_by_id(self) -> Dict[str, int]: - return self.session_storage.peers_by_id + @abc.abstractmethod + def peers_by_id(self): + ... @property - def peers_by_username(self) -> Dict[str, int]: - return self.session_storage.peers_by_username + @abc.abstractmethod + def peers_by_username(self): + ... @property - def peers_by_phone(self) -> Dict[str, int]: - return self.session_storage.peers_by_phone + @abc.abstractmethod + def peers_by_phone(self): + ... diff --git a/pyrogram/client/session_storage/base_session_storage.py b/pyrogram/client/session_storage/base_session_storage.py deleted file mode 100644 index 92473956..00000000 --- a/pyrogram/client/session_storage/base_session_storage.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 abc -from typing import Type - -import pyrogram - - -class SessionDoesNotExist(Exception): - pass - - -class BaseSessionStorage(abc.ABC): - def __init__(self, client: 'pyrogram.client.BaseClient', session_data): - self.client = client - self.session_data = session_data - self.dc_id = 1 - self.test_mode = None - self.auth_key = None - self.user_id = None - self.date = 0 - self.is_bot = False - self.peers_by_id = {} - self.peers_by_username = {} - self.peers_by_phone = {} - - @abc.abstractmethod - def load_session(self): - ... - - @abc.abstractmethod - def save_session(self, sync=False): - ... - - @abc.abstractmethod - def sync_cleanup(self): - ... - - -class BaseSessionConfig(abc.ABC): - @property - @abc.abstractmethod - def session_storage_cls(self) -> Type[BaseSessionStorage]: - ... diff --git a/pyrogram/client/session_storage/json_session_storage.py b/pyrogram/client/session_storage/json.py similarity index 58% rename from pyrogram/client/session_storage/json_session_storage.py rename to pyrogram/client/session_storage/json.py index 1e1e0ca4..170089a4 100644 --- a/pyrogram/client/session_storage/json_session_storage.py +++ b/pyrogram/client/session_storage/json.py @@ -22,21 +22,26 @@ import logging import os import shutil +import pyrogram from ..ext import utils -from . import BaseSessionStorage, SessionDoesNotExist +from . import MemorySessionStorage, SessionDoesNotExist log = logging.getLogger(__name__) -class JsonSessionStorage(BaseSessionStorage): +class JsonSessionStorage(MemorySessionStorage): + def __init__(self, client: 'pyrogram.client.ext.BaseClient', session_name: str): + super(JsonSessionStorage, self).__init__(client) + self._session_name = session_name + def _get_file_name(self, name: str): if not name.endswith('.session'): name += '.session' - return os.path.join(self.client.workdir, name) + return os.path.join(self._client.workdir, name) - def load_session(self): - file_path = self._get_file_name(self.session_data) + def load(self): + file_path = self._get_file_name(self._session_name) log.info('Loading JSON session from {}'.format(file_path)) try: @@ -45,59 +50,59 @@ class JsonSessionStorage(BaseSessionStorage): except FileNotFoundError: raise SessionDoesNotExist() - self.dc_id = s["dc_id"] - self.test_mode = s["test_mode"] - self.auth_key = base64.b64decode("".join(s["auth_key"])) # join split key - self.user_id = s["user_id"] - self.date = s.get("date", 0) - self.is_bot = s.get('is_bot', self.client.is_bot) + self._dc_id = s["dc_id"] + self._test_mode = s["test_mode"] + self._auth_key = base64.b64decode("".join(s["auth_key"])) # join split key + self._user_id = s["user_id"] + self._date = s.get("date", 0) + self._is_bot = s.get('is_bot', self._is_bot) for k, v in s.get("peers_by_id", {}).items(): - self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v) + self._peers_by_id[int(k)] = utils.get_input_peer(int(k), v) for k, v in s.get("peers_by_username", {}).items(): - peer = self.peers_by_id.get(v, None) + peer = self._peers_by_id.get(v, None) if peer: - self.peers_by_username[k] = 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) + peer = self._peers_by_id.get(v, None) if peer: - self.peers_by_phone[k] = peer + self._peers_by_phone[k] = peer - def save_session(self, sync=False): - file_path = self._get_file_name(self.session_data) + def save(self, sync=False): + file_path = self._get_file_name(self._session_name) if sync: file_path += '.tmp' log.info('Saving JSON session to {}, sync={}'.format(file_path, sync)) - auth_key = base64.b64encode(self.auth_key).decode() + auth_key = base64.b64encode(self._auth_key).decode() auth_key = [auth_key[i: i + 43] for i in range(0, len(auth_key), 43)] # split key in lines of 43 chars - os.makedirs(self.client.workdir, exist_ok=True) + os.makedirs(self._client.workdir, exist_ok=True) data = { - 'dc_id': self.dc_id, - 'test_mode': self.test_mode, + '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, + 'user_id': self._user_id, + 'date': self._date, + 'is_bot': self._is_bot, 'peers_by_id': { k: getattr(v, "access_hash", None) - for k, v in self.peers_by_id.copy().items() + for k, v in self._peers_by_id.copy().items() }, 'peers_by_username': { k: utils.get_peer_id(v) - for k, v in self.peers_by_username.copy().items() + for k, v in self._peers_by_username.copy().items() }, 'peers_by_phone': { k: utils.get_peer_id(v) - for k, v in self.peers_by_phone.copy().items() + for k, v in self._peers_by_phone.copy().items() } } @@ -109,10 +114,10 @@ class JsonSessionStorage(BaseSessionStorage): # execution won't be here if an error has occurred earlier if sync: - shutil.move(file_path, self._get_file_name(self.session_data)) + shutil.move(file_path, self._get_file_name(self._session_name)) def sync_cleanup(self): try: - os.remove(self._get_file_name(self.session_data) + '.tmp') + os.remove(self._get_file_name(self._session_name) + '.tmp') except OSError: pass diff --git a/pyrogram/client/session_storage/memory.py b/pyrogram/client/session_storage/memory.py new file mode 100644 index 00000000..f456f8eb --- /dev/null +++ b/pyrogram/client/session_storage/memory.py @@ -0,0 +1,85 @@ +import pyrogram +from . import SessionStorage, SessionDoesNotExist + + +class MemorySessionStorage(SessionStorage): + def __init__(self, client: 'pyrogram.client.ext.BaseClient'): + super(MemorySessionStorage, self).__init__(client) + self._dc_id = 1 + self._test_mode = None + self._auth_key = None + self._user_id = None + self._date = 0 + self._is_bot = False + self._peers_by_id = {} + self._peers_by_username = {} + self._peers_by_phone = {} + + def load(self): + raise SessionDoesNotExist() + + def save(self, sync=False): + pass + + def sync_cleanup(self): + pass + + @property + def dc_id(self): + return self._dc_id + + @dc_id.setter + def dc_id(self, val): + self._dc_id = val + + @property + def test_mode(self): + return self._test_mode + + @test_mode.setter + def test_mode(self, val): + self._test_mode = val + + @property + def auth_key(self): + return self._auth_key + + @auth_key.setter + def auth_key(self, val): + self._auth_key = val + + @property + def user_id(self): + return self._user_id + + @user_id.setter + def user_id(self, val): + self._user_id = val + + @property + def date(self): + return self._date + + @date.setter + def date(self, val): + self._date = val + + @property + def is_bot(self): + return self._is_bot + + @is_bot.setter + def is_bot(self, val): + self._is_bot = val + + @property + def peers_by_id(self): + return self._peers_by_id + + @property + def peers_by_username(self): + return self._peers_by_username + + @property + def peers_by_phone(self): + return self._peers_by_phone diff --git a/pyrogram/client/session_storage/string_session_storage.py b/pyrogram/client/session_storage/string.py similarity index 62% rename from pyrogram/client/session_storage/string_session_storage.py rename to pyrogram/client/session_storage/string.py index 5b1a8cc1..f8ec740a 100644 --- a/pyrogram/client/session_storage/string_session_storage.py +++ b/pyrogram/client/session_storage/string.py @@ -2,10 +2,11 @@ import base64 import binascii import struct -from . import BaseSessionStorage, SessionDoesNotExist +import pyrogram +from . import MemorySessionStorage, SessionDoesNotExist -class StringSessionStorage(BaseSessionStorage): +class StringSessionStorage(MemorySessionStorage): """ Packs session data as following (forcing little-endian byte order): Char dc_id (1 byte, unsigned) @@ -18,22 +19,26 @@ class StringSessionStorage(BaseSessionStorage): """ PACK_FORMAT = '. import abc -from typing import Type +from typing import Type, Union import pyrogram +from pyrogram.api import types class SessionDoesNotExist(Exception): @@ -102,17 +103,41 @@ class SessionStorage(abc.ABC): def is_bot(self, val): ... - @property @abc.abstractmethod - def peers_by_id(self): + def clear_cache(self): ... - @property @abc.abstractmethod - def peers_by_username(self): + def cache_peer(self, entity: Union[types.User, + types.Chat, types.ChatForbidden, + types.Channel, types.ChannelForbidden]): ... - @property @abc.abstractmethod - def peers_by_phone(self): + def get_peer_by_id(self, val: int): + ... + + @abc.abstractmethod + def get_peer_by_username(self, val: str): + ... + + @abc.abstractmethod + def get_peer_by_phone(self, val: str): + ... + + def get_peer(self, peer_id: Union[int, str]): + if isinstance(peer_id, int): + return self.get_peer_by_id(peer_id) + else: + peer_id = peer_id.lstrip('+@') + if peer_id.isdigit(): + return self.get_peer_by_phone(peer_id) + return self.get_peer_by_username(peer_id) + + @abc.abstractmethod + def peers_count(self): + ... + + @abc.abstractmethod + def contacts_count(self): ... diff --git a/pyrogram/client/session_storage/json.py b/pyrogram/client/session_storage/json.py index 170089a4..aaa6b96f 100644 --- a/pyrogram/client/session_storage/json.py +++ b/pyrogram/client/session_storage/json.py @@ -58,19 +58,19 @@ class JsonSessionStorage(MemorySessionStorage): self._is_bot = s.get('is_bot', self._is_bot) for k, v in s.get("peers_by_id", {}).items(): - self._peers_by_id[int(k)] = utils.get_input_peer(int(k), v) + self._peers_cache['i' + k] = utils.get_input_peer(int(k), v) for k, v in s.get("peers_by_username", {}).items(): - peer = self._peers_by_id.get(v, None) - - if peer: - self._peers_by_username[k] = peer + try: + self._peers_cache['u' + k] = self.get_peer_by_id(v) + except KeyError: + pass 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 + try: + self._peers_cache['p' + k] = self.get_peer_by_id(v) + except KeyError: + pass def save(self, sync=False): file_path = self._get_file_name(self._session_name) @@ -93,16 +93,19 @@ class JsonSessionStorage(MemorySessionStorage): 'date': self._date, 'is_bot': self._is_bot, 'peers_by_id': { - k: getattr(v, "access_hash", None) - for k, v in self._peers_by_id.copy().items() + k[1:]: getattr(v, "access_hash", None) + for k, v in self._peers_cache.copy().items() + if k[0] == 'i' }, 'peers_by_username': { - k: utils.get_peer_id(v) - for k, v in self._peers_by_username.copy().items() + k[1:]: utils.get_peer_id(v) + for k, v in self._peers_cache.copy().items() + if k[0] == 'u' }, 'peers_by_phone': { - k: utils.get_peer_id(v) - for k, v in self._peers_by_phone.copy().items() + k[1:]: utils.get_peer_id(v) + for k, v in self._peers_cache.copy().items() + if k[0] == 'p' } } diff --git a/pyrogram/client/session_storage/memory.py b/pyrogram/client/session_storage/memory.py index f456f8eb..d5f92f0d 100644 --- a/pyrogram/client/session_storage/memory.py +++ b/pyrogram/client/session_storage/memory.py @@ -1,4 +1,5 @@ import pyrogram +from pyrogram.api import types from . import SessionStorage, SessionDoesNotExist @@ -11,9 +12,7 @@ class MemorySessionStorage(SessionStorage): self._user_id = None self._date = 0 self._is_bot = False - self._peers_by_id = {} - self._peers_by_username = {} - self._peers_by_phone = {} + self._peers_cache = {} def load(self): raise SessionDoesNotExist() @@ -72,14 +71,48 @@ class MemorySessionStorage(SessionStorage): def is_bot(self, val): self._is_bot = val - @property - def peers_by_id(self): - return self._peers_by_id + def clear_cache(self): + keys = list(filter(lambda k: k[0] in 'up', self._peers_cache.keys())) + for key in keys: + try: + del self._peers_cache[key] + except KeyError: + pass - @property - def peers_by_username(self): - return self._peers_by_username + def cache_peer(self, entity): + if isinstance(entity, types.User): + input_peer = types.InputPeerUser( + user_id=entity.id, + access_hash=entity.access_hash + ) + self._peers_cache['i' + str(entity.id)] = input_peer + if entity.username: + self._peers_cache['u' + entity.username.lower()] = input_peer + if entity.phone: + self._peers_cache['p' + entity.phone] = input_peer + elif isinstance(entity, (types.Chat, types.ChatForbidden)): + self._peers_cache['i-' + str(entity.id)] = types.InputPeerChat(chat_id=entity.id) + elif isinstance(entity, (types.Channel, types.ChannelForbidden)): + input_peer = types.InputPeerChannel( + channel_id=entity.id, + access_hash=entity.access_hash + ) + self._peers_cache['i-100' + str(entity.id)] = input_peer + username = getattr(entity, "username", None) + if username: + self._peers_cache['u' + username.lower()] = input_peer - @property - def peers_by_phone(self): - return self._peers_by_phone + def get_peer_by_id(self, val): + return self._peers_cache['i' + str(val)] + + def get_peer_by_username(self, val): + return self._peers_cache['u' + val.lower()] + + def get_peer_by_phone(self, val): + return self._peers_cache['p' + val] + + def peers_count(self): + return len(list(filter(lambda k: k[0] == 'i', self._peers_cache.keys()))) + + def contacts_count(self): + return len(list(filter(lambda k: k[0] == 'p', self._peers_cache.keys()))) diff --git a/pyrogram/client/style/html.py b/pyrogram/client/style/html.py index 9a72a565..88e317cd 100644 --- a/pyrogram/client/style/html.py +++ b/pyrogram/client/style/html.py @@ -29,14 +29,15 @@ from pyrogram.api.types import ( InputMessageEntityMentionName as Mention, ) from . import utils +from ..session_storage import SessionStorage class HTML: HTML_RE = re.compile(r"<(\w+)(?: href=([\"'])([^<]+)\2)?>([^>]+)") MENTION_RE = re.compile(r"tg://user\?id=(\d+)") - def __init__(self, peers_by_id): - self.peers_by_id = peers_by_id + def __init__(self, session_storage: SessionStorage): + self.session_storage = session_storage def parse(self, message: str): entities = [] @@ -52,7 +53,10 @@ class HTML: if mention: user_id = int(mention.group(1)) - input_user = self.peers_by_id.get(user_id, None) + try: + input_user = self.session_storage.get_peer_by_id(user_id) + except KeyError: + input_user = None entity = ( Mention(start, len(body), input_user) diff --git a/pyrogram/client/style/markdown.py b/pyrogram/client/style/markdown.py index 05a11a25..6793b643 100644 --- a/pyrogram/client/style/markdown.py +++ b/pyrogram/client/style/markdown.py @@ -29,6 +29,7 @@ from pyrogram.api.types import ( InputMessageEntityMentionName as Mention ) from . import utils +from ..session_storage import SessionStorage class Markdown: @@ -52,8 +53,8 @@ class Markdown: )) MENTION_RE = re.compile(r"tg://user\?id=(\d+)") - def __init__(self, peers_by_id: dict): - self.peers_by_id = peers_by_id + def __init__(self, session_storage: SessionStorage): + self.session_storage = session_storage def parse(self, message: str): message = utils.add_surrogates(str(message)).strip() @@ -69,7 +70,10 @@ class Markdown: if mention: user_id = int(mention.group(1)) - input_user = self.peers_by_id.get(user_id, None) + try: + input_user = self.session_storage.get_peer_by_id(user_id) + except KeyError: + input_user = None entity = ( Mention(start, len(text), input_user) From 03b92b3302a9d316d4e693efa8b9d87b0b991fd0 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Tue, 26 Feb 2019 21:06:30 +0300 Subject: [PATCH 09/22] Implement SQLite session storage --- pyrogram/client/client.py | 2 +- pyrogram/client/session_storage/__init__.py | 1 + pyrogram/client/session_storage/json.py | 6 +- .../client/session_storage/sqlite/0001.sql | 21 +++ .../client/session_storage/sqlite/__init__.py | 132 ++++++++++++++++++ 5 files changed, 159 insertions(+), 3 deletions(-) create mode 100644 pyrogram/client/session_storage/sqlite/0001.sql create mode 100644 pyrogram/client/session_storage/sqlite/__init__.py diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index ad755977..5fc805c0 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -57,7 +57,7 @@ from .ext import utils, Syncer, BaseClient from .methods import Methods from .session_storage import ( SessionDoesNotExist, SessionStorage, MemorySessionStorage, JsonSessionStorage, - StringSessionStorage + StringSessionStorage, SQLiteSessionStorage ) log = logging.getLogger(__name__) diff --git a/pyrogram/client/session_storage/__init__.py b/pyrogram/client/session_storage/__init__.py index ad2d8900..adfcf813 100644 --- a/pyrogram/client/session_storage/__init__.py +++ b/pyrogram/client/session_storage/__init__.py @@ -20,3 +20,4 @@ from .abstract import SessionStorage, SessionDoesNotExist from .memory import MemorySessionStorage from .json import JsonSessionStorage from .string import StringSessionStorage +from .sqlite import SQLiteSessionStorage diff --git a/pyrogram/client/session_storage/json.py b/pyrogram/client/session_storage/json.py index aaa6b96f..570e1525 100644 --- a/pyrogram/client/session_storage/json.py +++ b/pyrogram/client/session_storage/json.py @@ -29,6 +29,8 @@ from . import MemorySessionStorage, SessionDoesNotExist log = logging.getLogger(__name__) +EXTENSION = '.session' + class JsonSessionStorage(MemorySessionStorage): def __init__(self, client: 'pyrogram.client.ext.BaseClient', session_name: str): @@ -36,8 +38,8 @@ class JsonSessionStorage(MemorySessionStorage): self._session_name = session_name def _get_file_name(self, name: str): - if not name.endswith('.session'): - name += '.session' + if not name.endswith(EXTENSION): + name += EXTENSION return os.path.join(self._client.workdir, name) def load(self): diff --git a/pyrogram/client/session_storage/sqlite/0001.sql b/pyrogram/client/session_storage/sqlite/0001.sql new file mode 100644 index 00000000..d81e9554 --- /dev/null +++ b/pyrogram/client/session_storage/sqlite/0001.sql @@ -0,0 +1,21 @@ +create table sessions ( + dc_id integer primary key, + test_mode integer, + auth_key blob, + user_id integer, + date integer, + is_bot integer +); + +create table peers_cache ( + id integer primary key, + hash integer, + username text, + phone integer +); + +create table migrations ( + name text primary key +); + +insert into migrations (name) values ('0001'); diff --git a/pyrogram/client/session_storage/sqlite/__init__.py b/pyrogram/client/session_storage/sqlite/__init__.py new file mode 100644 index 00000000..75931109 --- /dev/null +++ b/pyrogram/client/session_storage/sqlite/__init__.py @@ -0,0 +1,132 @@ +# 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 logging +import os +import sqlite3 + +import pyrogram +from ....api import types +from ...ext import utils +from .. import MemorySessionStorage, SessionDoesNotExist + + +log = logging.getLogger(__name__) + +EXTENSION = '.session.sqlite3' +MIGRATIONS = ['0001'] + + +class SQLiteSessionStorage(MemorySessionStorage): + def __init__(self, client: 'pyrogram.client.ext.BaseClient', session_name: str): + super(SQLiteSessionStorage, self).__init__(client) + self._session_name = session_name + self._conn = None # type: sqlite3.Connection + + def _get_file_name(self, name: str): + if not name.endswith(EXTENSION): + name += EXTENSION + return os.path.join(self._client.workdir, name) + + def _apply_migrations(self, new_db=False): + migrations = MIGRATIONS.copy() + if not new_db: + cursor = self._conn.cursor() + cursor.execute('select name from migrations') + for row in cursor.fetchone(): + migrations.remove(row) + for name in migrations: + with open(os.path.join(os.path.dirname(__file__), '{}.sql'.format(name))) as script: + self._conn.executescript(script.read()) + + def load(self): + file_path = self._get_file_name(self._session_name) + log.info('Loading SQLite session from {}'.format(file_path)) + + if os.path.isfile(file_path): + self._conn = sqlite3.connect(file_path) + self._apply_migrations() + else: + self._conn = sqlite3.connect(file_path) + self._apply_migrations(new_db=True) + + cursor = self._conn.cursor() + cursor.execute('select dc_id, test_mode, auth_key, user_id, "date", is_bot from sessions') + row = cursor.fetchone() + if not row: + raise SessionDoesNotExist() + + self._dc_id = row[0] + self._test_mode = bool(row[1]) + self._auth_key = row[2] + self._user_id = row[3] + self._date = row[4] + self._is_bot = bool(row[5]) + + def cache_peer(self, entity): + peer_id = username = phone = access_hash = None + + if isinstance(entity, types.User): + peer_id = entity.id + username = entity.username.lower() if entity.username else None + phone = entity.phone or None + access_hash = entity.access_hash + elif isinstance(entity, (types.Chat, types.ChatForbidden)): + peer_id = -entity.id + # input_peer = types.InputPeerChat(chat_id=entity.id) + elif isinstance(entity, (types.Channel, types.ChannelForbidden)): + peer_id = int('-100' + str(entity.id)) + username = entity.username.lower() if hasattr(entity, 'username') and entity.username else None + access_hash = entity.access_hash + + self._conn.execute('insert or replace into peers_cache values (?, ?, ?, ?)', + (peer_id, access_hash, username, phone)) + + def get_peer_by_id(self, val): + cursor = self._conn.cursor() + cursor.execute('select id, hash from peers_cache where id = ?', (val,)) + row = cursor.fetchone() + if not row: + raise KeyError(val) + return utils.get_input_peer(row[0], row[1]) + + def get_peer_by_username(self, val): + cursor = self._conn.cursor() + cursor.execute('select id, hash from peers_cache where username = ?', (val,)) + row = cursor.fetchone() + if not row: + raise KeyError(val) + return utils.get_input_peer(row[0], row[1]) + + def get_peer_by_phone(self, val): + cursor = self._conn.cursor() + cursor.execute('select id, hash from peers_cache where phone = ?', (val,)) + row = cursor.fetchone() + if not row: + raise KeyError(val) + return utils.get_input_peer(row[0], row[1]) + + def save(self, sync=False): + log.info('Committing SQLite session') + self._conn.execute('delete from sessions') + self._conn.execute('insert into sessions values (?, ?, ?, ?, ?, ?)', + (self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot)) + self._conn.commit() + + def sync_cleanup(self): + pass From 10fc340efff40dc54e35bc687af5966b3ad077f5 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Tue, 26 Feb 2019 21:43:23 +0300 Subject: [PATCH 10/22] Add session migrating from json; add some indexes to sqlite sessions --- pyrogram/client/client.py | 2 +- .../client/session_storage/sqlite/0001.sql | 3 ++ .../client/session_storage/sqlite/__init__.py | 29 +++++++++++++++---- 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 5fc805c0..d2bc3ee4 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -214,7 +214,7 @@ class Client(Methods, BaseClient): elif session_name.startswith(':'): session_storage = StringSessionStorage(self, session_name) else: - session_storage = JsonSessionStorage(self, session_name) + session_storage = SQLiteSessionStorage(self, session_name) elif isinstance(session_name, SessionStorage): session_storage = session_name else: diff --git a/pyrogram/client/session_storage/sqlite/0001.sql b/pyrogram/client/session_storage/sqlite/0001.sql index d81e9554..c6c51d24 100644 --- a/pyrogram/client/session_storage/sqlite/0001.sql +++ b/pyrogram/client/session_storage/sqlite/0001.sql @@ -18,4 +18,7 @@ create table migrations ( name text primary key ); +create index username_idx on peers_cache(username); +create index phone_idx on peers_cache(phone); + insert into migrations (name) values ('0001'); diff --git a/pyrogram/client/session_storage/sqlite/__init__.py b/pyrogram/client/session_storage/sqlite/__init__.py index 75931109..4fc7ff64 100644 --- a/pyrogram/client/session_storage/sqlite/__init__.py +++ b/pyrogram/client/session_storage/sqlite/__init__.py @@ -18,17 +18,18 @@ import logging import os +import shutil import sqlite3 import pyrogram from ....api import types from ...ext import utils -from .. import MemorySessionStorage, SessionDoesNotExist +from .. import MemorySessionStorage, SessionDoesNotExist, JsonSessionStorage log = logging.getLogger(__name__) -EXTENSION = '.session.sqlite3' +EXTENSION = '.session' MIGRATIONS = ['0001'] @@ -54,13 +55,32 @@ class SQLiteSessionStorage(MemorySessionStorage): with open(os.path.join(os.path.dirname(__file__), '{}.sql'.format(name))) as script: self._conn.executescript(script.read()) + def _migrate_from_json(self): + jss = JsonSessionStorage(self._client, self._session_name) + jss.load() + file_path = self._get_file_name(self._session_name) + self._conn = sqlite3.connect(file_path + '.tmp') + self._apply_migrations(new_db=True) + self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot = \ + jss.dc_id, jss.test_mode, jss.auth_key, jss.user_id, jss.date, jss.is_bot + self.save() + self._conn.close() + shutil.move(file_path + '.tmp', file_path) + log.warning('Session was migrated from JSON, loading...') + self.load() + def load(self): file_path = self._get_file_name(self._session_name) log.info('Loading SQLite session from {}'.format(file_path)) if os.path.isfile(file_path): - self._conn = sqlite3.connect(file_path) - self._apply_migrations() + try: + self._conn = sqlite3.connect(file_path) + self._apply_migrations() + except sqlite3.DatabaseError: + log.warning('Trying to migrate session from JSON...') + self._migrate_from_json() + return else: self._conn = sqlite3.connect(file_path) self._apply_migrations(new_db=True) @@ -88,7 +108,6 @@ class SQLiteSessionStorage(MemorySessionStorage): access_hash = entity.access_hash elif isinstance(entity, (types.Chat, types.ChatForbidden)): peer_id = -entity.id - # input_peer = types.InputPeerChat(chat_id=entity.id) elif isinstance(entity, (types.Channel, types.ChannelForbidden)): peer_id = int('-100' + str(entity.id)) username = entity.username.lower() if hasattr(entity, 'username') and entity.username else None From 033622cfb85efbd8a09abf8b0949f9ccc0495b90 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Wed, 27 Feb 2019 22:49:23 +0300 Subject: [PATCH 11/22] Cleanup json session storage specific code as it is used only for migrations --- pyrogram/client/ext/syncer.py | 2 - pyrogram/client/session_storage/abstract.py | 4 -- pyrogram/client/session_storage/json.py | 67 +------------------ pyrogram/client/session_storage/memory.py | 3 - .../client/session_storage/sqlite/__init__.py | 3 - pyrogram/client/session_storage/string.py | 3 - 6 files changed, 1 insertion(+), 81 deletions(-) diff --git a/pyrogram/client/ext/syncer.py b/pyrogram/client/ext/syncer.py index e13212be..9e7d2303 100644 --- a/pyrogram/client/ext/syncer.py +++ b/pyrogram/client/ext/syncer.py @@ -88,5 +88,3 @@ class Syncer: log.critical(e, exc_info=True) else: log.info("Synced {}".format(client.session_name)) - finally: - client.session_storage.sync_cleanup() diff --git a/pyrogram/client/session_storage/abstract.py b/pyrogram/client/session_storage/abstract.py index 39517a01..134d5c8c 100644 --- a/pyrogram/client/session_storage/abstract.py +++ b/pyrogram/client/session_storage/abstract.py @@ -38,10 +38,6 @@ class SessionStorage(abc.ABC): @abc.abstractmethod def save(self, sync=False): ... - - @abc.abstractmethod - def sync_cleanup(self): - ... @property @abc.abstractmethod diff --git a/pyrogram/client/session_storage/json.py b/pyrogram/client/session_storage/json.py index 570e1525..4a48d3c1 100644 --- a/pyrogram/client/session_storage/json.py +++ b/pyrogram/client/session_storage/json.py @@ -59,70 +59,5 @@ class JsonSessionStorage(MemorySessionStorage): self._date = s.get("date", 0) self._is_bot = s.get('is_bot', self._is_bot) - for k, v in s.get("peers_by_id", {}).items(): - self._peers_cache['i' + k] = utils.get_input_peer(int(k), v) - - for k, v in s.get("peers_by_username", {}).items(): - try: - self._peers_cache['u' + k] = self.get_peer_by_id(v) - except KeyError: - pass - - for k, v in s.get("peers_by_phone", {}).items(): - try: - self._peers_cache['p' + k] = self.get_peer_by_id(v) - except KeyError: - pass - def save(self, sync=False): - file_path = self._get_file_name(self._session_name) - - if sync: - file_path += '.tmp' - - log.info('Saving JSON session to {}, sync={}'.format(file_path, sync)) - - auth_key = base64.b64encode(self._auth_key).decode() - auth_key = [auth_key[i: i + 43] for i in range(0, len(auth_key), 43)] # split key in lines of 43 chars - - os.makedirs(self._client.workdir, exist_ok=True) - - data = { - '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, - 'peers_by_id': { - k[1:]: getattr(v, "access_hash", None) - for k, v in self._peers_cache.copy().items() - if k[0] == 'i' - }, - 'peers_by_username': { - k[1:]: utils.get_peer_id(v) - for k, v in self._peers_cache.copy().items() - if k[0] == 'u' - }, - 'peers_by_phone': { - k[1:]: utils.get_peer_id(v) - for k, v in self._peers_cache.copy().items() - if k[0] == 'p' - } - } - - with open(file_path, "w", encoding="utf-8") as f: - json.dump(data, f, indent=4) - - f.flush() - os.fsync(f.fileno()) - - # execution won't be here if an error has occurred earlier - if sync: - shutil.move(file_path, self._get_file_name(self._session_name)) - - def sync_cleanup(self): - try: - os.remove(self._get_file_name(self._session_name) + '.tmp') - except OSError: - pass + pass diff --git a/pyrogram/client/session_storage/memory.py b/pyrogram/client/session_storage/memory.py index d5f92f0d..c0610e70 100644 --- a/pyrogram/client/session_storage/memory.py +++ b/pyrogram/client/session_storage/memory.py @@ -20,9 +20,6 @@ class MemorySessionStorage(SessionStorage): def save(self, sync=False): pass - def sync_cleanup(self): - pass - @property def dc_id(self): return self._dc_id diff --git a/pyrogram/client/session_storage/sqlite/__init__.py b/pyrogram/client/session_storage/sqlite/__init__.py index 4fc7ff64..0308a4dc 100644 --- a/pyrogram/client/session_storage/sqlite/__init__.py +++ b/pyrogram/client/session_storage/sqlite/__init__.py @@ -146,6 +146,3 @@ class SQLiteSessionStorage(MemorySessionStorage): self._conn.execute('insert into sessions values (?, ?, ?, ?, ?, ?)', (self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot)) self._conn.commit() - - def sync_cleanup(self): - pass diff --git a/pyrogram/client/session_storage/string.py b/pyrogram/client/session_storage/string.py index f8ec740a..11051323 100644 --- a/pyrogram/client/session_storage/string.py +++ b/pyrogram/client/session_storage/string.py @@ -44,6 +44,3 @@ class StringSessionStorage(MemorySessionStorage): encoded = ':' + base64.b64encode(packed, b'-_').decode('latin-1').rstrip('=') split = '\n'.join(['"{}"'.format(encoded[i: i + 50]) for i in range(0, len(encoded), 50)]) print('Created session string:\n{}'.format(split)) - - def sync_cleanup(self): - pass From 8cc61f00ed74fc8290b4d75cc1503275a42d5136 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Fri, 1 Mar 2019 21:23:01 +0300 Subject: [PATCH 12/22] Fix threading with sqlite storage --- .../client/session_storage/sqlite/__init__.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/pyrogram/client/session_storage/sqlite/__init__.py b/pyrogram/client/session_storage/sqlite/__init__.py index 0308a4dc..a16e75e8 100644 --- a/pyrogram/client/session_storage/sqlite/__init__.py +++ b/pyrogram/client/session_storage/sqlite/__init__.py @@ -20,6 +20,7 @@ import logging import os import shutil import sqlite3 +from threading import Lock import pyrogram from ....api import types @@ -38,6 +39,7 @@ class SQLiteSessionStorage(MemorySessionStorage): super(SQLiteSessionStorage, self).__init__(client) self._session_name = session_name self._conn = None # type: sqlite3.Connection + self._lock = Lock() def _get_file_name(self, name: str): if not name.endswith(EXTENSION): @@ -45,6 +47,7 @@ class SQLiteSessionStorage(MemorySessionStorage): return os.path.join(self._client.workdir, name) def _apply_migrations(self, new_db=False): + self._conn.execute('PRAGMA read_uncommitted = true') migrations = MIGRATIONS.copy() if not new_db: cursor = self._conn.cursor() @@ -75,14 +78,14 @@ class SQLiteSessionStorage(MemorySessionStorage): if os.path.isfile(file_path): try: - self._conn = sqlite3.connect(file_path) + self._conn = sqlite3.connect(file_path, isolation_level='EXCLUSIVE', check_same_thread=False) self._apply_migrations() except sqlite3.DatabaseError: log.warning('Trying to migrate session from JSON...') self._migrate_from_json() return else: - self._conn = sqlite3.connect(file_path) + self._conn = sqlite3.connect(file_path, isolation_level='EXCLUSIVE', check_same_thread=False) self._apply_migrations(new_db=True) cursor = self._conn.cursor() @@ -113,8 +116,9 @@ class SQLiteSessionStorage(MemorySessionStorage): username = entity.username.lower() if hasattr(entity, 'username') and entity.username else None access_hash = entity.access_hash - self._conn.execute('insert or replace into peers_cache values (?, ?, ?, ?)', - (peer_id, access_hash, username, phone)) + with self._lock: + self._conn.execute('insert or replace into peers_cache values (?, ?, ?, ?)', + (peer_id, access_hash, username, phone)) def get_peer_by_id(self, val): cursor = self._conn.cursor() @@ -142,7 +146,8 @@ class SQLiteSessionStorage(MemorySessionStorage): def save(self, sync=False): log.info('Committing SQLite session') - self._conn.execute('delete from sessions') - self._conn.execute('insert into sessions values (?, ?, ?, ?, ?, ?)', - (self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot)) - self._conn.commit() + with self._lock: + self._conn.execute('delete from sessions') + self._conn.execute('insert into sessions values (?, ?, ?, ?, ?, ?)', + (self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot)) + self._conn.commit() From 85700b0ffc191458677b6d29452cff121a5d4d13 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Fri, 1 Mar 2019 21:23:53 +0300 Subject: [PATCH 13/22] Do not cache entities without access_hash --- pyrogram/client/client.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 2bcf294f..33b3f137 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -762,6 +762,8 @@ class Client(Methods, BaseClient): types.Chat, types.ChatForbidden, types.Channel, types.ChannelForbidden]]): for entity in entities: + if isinstance(entity, (types.User, types.Channel, types.ChannelForbidden)) and not entity.access_hash: + continue self.session_storage.cache_peer(entity) def download_worker(self): From 682591ea8fdb3f33c2970690b0aafd67c8823c29 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:01:23 +0200 Subject: [PATCH 14/22] Update Auth and Session to accommodate Storage Engines --- pyrogram/client/client.py | 24 +-- pyrogram/client/ext/base_client.py | 1 - pyrogram/client/session_storage/abstract.py | 139 ---------------- pyrogram/client/session_storage/json.py | 63 -------- pyrogram/client/session_storage/memory.py | 115 ------------- .../client/session_storage/sqlite/0001.sql | 24 --- .../client/session_storage/sqlite/__init__.py | 153 ------------------ pyrogram/client/session_storage/string.py | 46 ------ .../{session_storage => storage}/__init__.py | 8 +- pyrogram/client/style/html.py | 9 +- pyrogram/client/style/markdown.py | 9 +- pyrogram/session/auth.py | 10 +- pyrogram/session/session.py | 23 ++- 13 files changed, 36 insertions(+), 588 deletions(-) delete mode 100644 pyrogram/client/session_storage/abstract.py delete mode 100644 pyrogram/client/session_storage/json.py delete mode 100644 pyrogram/client/session_storage/memory.py delete mode 100644 pyrogram/client/session_storage/sqlite/0001.sql delete mode 100644 pyrogram/client/session_storage/sqlite/__init__.py delete mode 100644 pyrogram/client/session_storage/string.py rename pyrogram/client/{session_storage => storage}/__init__.py (78%) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 1aa436b5..885c3334 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -26,14 +26,13 @@ import shutil import tempfile import threading import time -import warnings from configparser import ConfigParser from hashlib import sha256, md5 from importlib import import_module from pathlib import Path from signal import signal, SIGINT, SIGTERM, SIGABRT from threading import Thread -from typing import Union, List, Type +from typing import Union, List from pyrogram.api import functions, types from pyrogram.api.core import TLObject @@ -205,24 +204,9 @@ class Client(Methods, BaseClient): no_updates: bool = None, takeout: bool = None ): + super().__init__() - if isinstance(session_name, str): - if session_name == ':memory:': - session_storage = MemorySessionStorage(self) - elif session_name.startswith(':'): - session_storage = StringSessionStorage(self, session_name) - else: - session_storage = SQLiteSessionStorage(self, session_name) - elif isinstance(session_name, SessionStorage): - session_storage = session_name - else: - raise RuntimeError('Wrong session_name passed, expected str or SessionConfig subclass') - - super().__init__(session_storage) - - super().__init__(session_storage) - - self.session_name = str(session_name) # TODO: build correct session name + self.session_name = session_name self.api_id = int(api_id) if api_id else None self.api_hash = api_hash self.app_version = app_version @@ -232,7 +216,7 @@ class Client(Methods, BaseClient): self.ipv6 = ipv6 # TODO: Make code consistent, use underscore for private/protected fields self._proxy = proxy - self.session_storage.test_mode = test_mode + self.test_mode = test_mode self.bot_token = bot_token self.phone_number = phone_number self.phone_code = phone_code diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py index aaf87823..9276b0eb 100644 --- a/pyrogram/client/ext/base_client.py +++ b/pyrogram/client/ext/base_client.py @@ -27,7 +27,6 @@ from threading import Lock from pyrogram import __version__ from ..style import Markdown, HTML from ...session.internals import MsgId -from ..session_storage import SessionStorage class BaseClient: diff --git a/pyrogram/client/session_storage/abstract.py b/pyrogram/client/session_storage/abstract.py deleted file mode 100644 index 134d5c8c..00000000 --- a/pyrogram/client/session_storage/abstract.py +++ /dev/null @@ -1,139 +0,0 @@ -# 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 abc -from typing import Type, Union - -import pyrogram -from pyrogram.api import types - - -class SessionDoesNotExist(Exception): - pass - - -class SessionStorage(abc.ABC): - def __init__(self, client: 'pyrogram.client.BaseClient'): - self._client = client - - @abc.abstractmethod - def load(self): - ... - - @abc.abstractmethod - def save(self, sync=False): - ... - - @property - @abc.abstractmethod - def dc_id(self): - ... - - @dc_id.setter - @abc.abstractmethod - def dc_id(self, val): - ... - - @property - @abc.abstractmethod - def test_mode(self): - ... - - @test_mode.setter - @abc.abstractmethod - def test_mode(self, val): - ... - - @property - @abc.abstractmethod - def auth_key(self): - ... - - @auth_key.setter - @abc.abstractmethod - def auth_key(self, val): - ... - - @property - @abc.abstractmethod - def user_id(self): - ... - - @user_id.setter - @abc.abstractmethod - def user_id(self, val): - ... - - @property - @abc.abstractmethod - def date(self): - ... - - @date.setter - @abc.abstractmethod - def date(self, val): - ... - - @property - @abc.abstractmethod - def is_bot(self): - ... - - @is_bot.setter - @abc.abstractmethod - def is_bot(self, val): - ... - - @abc.abstractmethod - def clear_cache(self): - ... - - @abc.abstractmethod - def cache_peer(self, entity: Union[types.User, - types.Chat, types.ChatForbidden, - types.Channel, types.ChannelForbidden]): - ... - - @abc.abstractmethod - def get_peer_by_id(self, val: int): - ... - - @abc.abstractmethod - def get_peer_by_username(self, val: str): - ... - - @abc.abstractmethod - def get_peer_by_phone(self, val: str): - ... - - def get_peer(self, peer_id: Union[int, str]): - if isinstance(peer_id, int): - return self.get_peer_by_id(peer_id) - else: - peer_id = peer_id.lstrip('+@') - if peer_id.isdigit(): - return self.get_peer_by_phone(peer_id) - return self.get_peer_by_username(peer_id) - - @abc.abstractmethod - def peers_count(self): - ... - - @abc.abstractmethod - def contacts_count(self): - ... diff --git a/pyrogram/client/session_storage/json.py b/pyrogram/client/session_storage/json.py deleted file mode 100644 index 4a48d3c1..00000000 --- a/pyrogram/client/session_storage/json.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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 shutil - -import pyrogram -from ..ext import utils -from . import MemorySessionStorage, SessionDoesNotExist - - -log = logging.getLogger(__name__) - -EXTENSION = '.session' - - -class JsonSessionStorage(MemorySessionStorage): - def __init__(self, client: 'pyrogram.client.ext.BaseClient', session_name: str): - super(JsonSessionStorage, self).__init__(client) - self._session_name = session_name - - def _get_file_name(self, name: str): - if not name.endswith(EXTENSION): - name += EXTENSION - return os.path.join(self._client.workdir, name) - - def load(self): - file_path = self._get_file_name(self._session_name) - log.info('Loading JSON session from {}'.format(file_path)) - - try: - with open(file_path, encoding='utf-8') as f: - s = json.load(f) - except FileNotFoundError: - raise SessionDoesNotExist() - - self._dc_id = s["dc_id"] - self._test_mode = s["test_mode"] - self._auth_key = base64.b64decode("".join(s["auth_key"])) # join split key - self._user_id = s["user_id"] - self._date = s.get("date", 0) - self._is_bot = s.get('is_bot', self._is_bot) - - def save(self, sync=False): - pass diff --git a/pyrogram/client/session_storage/memory.py b/pyrogram/client/session_storage/memory.py deleted file mode 100644 index c0610e70..00000000 --- a/pyrogram/client/session_storage/memory.py +++ /dev/null @@ -1,115 +0,0 @@ -import pyrogram -from pyrogram.api import types -from . import SessionStorage, SessionDoesNotExist - - -class MemorySessionStorage(SessionStorage): - def __init__(self, client: 'pyrogram.client.ext.BaseClient'): - super(MemorySessionStorage, self).__init__(client) - self._dc_id = 1 - self._test_mode = None - self._auth_key = None - self._user_id = None - self._date = 0 - self._is_bot = False - self._peers_cache = {} - - def load(self): - raise SessionDoesNotExist() - - def save(self, sync=False): - pass - - @property - def dc_id(self): - return self._dc_id - - @dc_id.setter - def dc_id(self, val): - self._dc_id = val - - @property - def test_mode(self): - return self._test_mode - - @test_mode.setter - def test_mode(self, val): - self._test_mode = val - - @property - def auth_key(self): - return self._auth_key - - @auth_key.setter - def auth_key(self, val): - self._auth_key = val - - @property - def user_id(self): - return self._user_id - - @user_id.setter - def user_id(self, val): - self._user_id = val - - @property - def date(self): - return self._date - - @date.setter - def date(self, val): - self._date = val - - @property - def is_bot(self): - return self._is_bot - - @is_bot.setter - def is_bot(self, val): - self._is_bot = val - - def clear_cache(self): - keys = list(filter(lambda k: k[0] in 'up', self._peers_cache.keys())) - for key in keys: - try: - del self._peers_cache[key] - except KeyError: - pass - - def cache_peer(self, entity): - if isinstance(entity, types.User): - input_peer = types.InputPeerUser( - user_id=entity.id, - access_hash=entity.access_hash - ) - self._peers_cache['i' + str(entity.id)] = input_peer - if entity.username: - self._peers_cache['u' + entity.username.lower()] = input_peer - if entity.phone: - self._peers_cache['p' + entity.phone] = input_peer - elif isinstance(entity, (types.Chat, types.ChatForbidden)): - self._peers_cache['i-' + str(entity.id)] = types.InputPeerChat(chat_id=entity.id) - elif isinstance(entity, (types.Channel, types.ChannelForbidden)): - input_peer = types.InputPeerChannel( - channel_id=entity.id, - access_hash=entity.access_hash - ) - self._peers_cache['i-100' + str(entity.id)] = input_peer - username = getattr(entity, "username", None) - if username: - self._peers_cache['u' + username.lower()] = input_peer - - def get_peer_by_id(self, val): - return self._peers_cache['i' + str(val)] - - def get_peer_by_username(self, val): - return self._peers_cache['u' + val.lower()] - - def get_peer_by_phone(self, val): - return self._peers_cache['p' + val] - - def peers_count(self): - return len(list(filter(lambda k: k[0] == 'i', self._peers_cache.keys()))) - - def contacts_count(self): - return len(list(filter(lambda k: k[0] == 'p', self._peers_cache.keys()))) diff --git a/pyrogram/client/session_storage/sqlite/0001.sql b/pyrogram/client/session_storage/sqlite/0001.sql deleted file mode 100644 index c6c51d24..00000000 --- a/pyrogram/client/session_storage/sqlite/0001.sql +++ /dev/null @@ -1,24 +0,0 @@ -create table sessions ( - dc_id integer primary key, - test_mode integer, - auth_key blob, - user_id integer, - date integer, - is_bot integer -); - -create table peers_cache ( - id integer primary key, - hash integer, - username text, - phone integer -); - -create table migrations ( - name text primary key -); - -create index username_idx on peers_cache(username); -create index phone_idx on peers_cache(phone); - -insert into migrations (name) values ('0001'); diff --git a/pyrogram/client/session_storage/sqlite/__init__.py b/pyrogram/client/session_storage/sqlite/__init__.py deleted file mode 100644 index a16e75e8..00000000 --- a/pyrogram/client/session_storage/sqlite/__init__.py +++ /dev/null @@ -1,153 +0,0 @@ -# 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 logging -import os -import shutil -import sqlite3 -from threading import Lock - -import pyrogram -from ....api import types -from ...ext import utils -from .. import MemorySessionStorage, SessionDoesNotExist, JsonSessionStorage - - -log = logging.getLogger(__name__) - -EXTENSION = '.session' -MIGRATIONS = ['0001'] - - -class SQLiteSessionStorage(MemorySessionStorage): - def __init__(self, client: 'pyrogram.client.ext.BaseClient', session_name: str): - super(SQLiteSessionStorage, self).__init__(client) - self._session_name = session_name - self._conn = None # type: sqlite3.Connection - self._lock = Lock() - - def _get_file_name(self, name: str): - if not name.endswith(EXTENSION): - name += EXTENSION - return os.path.join(self._client.workdir, name) - - def _apply_migrations(self, new_db=False): - self._conn.execute('PRAGMA read_uncommitted = true') - migrations = MIGRATIONS.copy() - if not new_db: - cursor = self._conn.cursor() - cursor.execute('select name from migrations') - for row in cursor.fetchone(): - migrations.remove(row) - for name in migrations: - with open(os.path.join(os.path.dirname(__file__), '{}.sql'.format(name))) as script: - self._conn.executescript(script.read()) - - def _migrate_from_json(self): - jss = JsonSessionStorage(self._client, self._session_name) - jss.load() - file_path = self._get_file_name(self._session_name) - self._conn = sqlite3.connect(file_path + '.tmp') - self._apply_migrations(new_db=True) - self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot = \ - jss.dc_id, jss.test_mode, jss.auth_key, jss.user_id, jss.date, jss.is_bot - self.save() - self._conn.close() - shutil.move(file_path + '.tmp', file_path) - log.warning('Session was migrated from JSON, loading...') - self.load() - - def load(self): - file_path = self._get_file_name(self._session_name) - log.info('Loading SQLite session from {}'.format(file_path)) - - if os.path.isfile(file_path): - try: - self._conn = sqlite3.connect(file_path, isolation_level='EXCLUSIVE', check_same_thread=False) - self._apply_migrations() - except sqlite3.DatabaseError: - log.warning('Trying to migrate session from JSON...') - self._migrate_from_json() - return - else: - self._conn = sqlite3.connect(file_path, isolation_level='EXCLUSIVE', check_same_thread=False) - self._apply_migrations(new_db=True) - - cursor = self._conn.cursor() - cursor.execute('select dc_id, test_mode, auth_key, user_id, "date", is_bot from sessions') - row = cursor.fetchone() - if not row: - raise SessionDoesNotExist() - - self._dc_id = row[0] - self._test_mode = bool(row[1]) - self._auth_key = row[2] - self._user_id = row[3] - self._date = row[4] - self._is_bot = bool(row[5]) - - def cache_peer(self, entity): - peer_id = username = phone = access_hash = None - - if isinstance(entity, types.User): - peer_id = entity.id - username = entity.username.lower() if entity.username else None - phone = entity.phone or None - access_hash = entity.access_hash - elif isinstance(entity, (types.Chat, types.ChatForbidden)): - peer_id = -entity.id - elif isinstance(entity, (types.Channel, types.ChannelForbidden)): - peer_id = int('-100' + str(entity.id)) - username = entity.username.lower() if hasattr(entity, 'username') and entity.username else None - access_hash = entity.access_hash - - with self._lock: - self._conn.execute('insert or replace into peers_cache values (?, ?, ?, ?)', - (peer_id, access_hash, username, phone)) - - def get_peer_by_id(self, val): - cursor = self._conn.cursor() - cursor.execute('select id, hash from peers_cache where id = ?', (val,)) - row = cursor.fetchone() - if not row: - raise KeyError(val) - return utils.get_input_peer(row[0], row[1]) - - def get_peer_by_username(self, val): - cursor = self._conn.cursor() - cursor.execute('select id, hash from peers_cache where username = ?', (val,)) - row = cursor.fetchone() - if not row: - raise KeyError(val) - return utils.get_input_peer(row[0], row[1]) - - def get_peer_by_phone(self, val): - cursor = self._conn.cursor() - cursor.execute('select id, hash from peers_cache where phone = ?', (val,)) - row = cursor.fetchone() - if not row: - raise KeyError(val) - return utils.get_input_peer(row[0], row[1]) - - def save(self, sync=False): - log.info('Committing SQLite session') - with self._lock: - self._conn.execute('delete from sessions') - self._conn.execute('insert into sessions values (?, ?, ?, ?, ?, ?)', - (self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot)) - self._conn.commit() diff --git a/pyrogram/client/session_storage/string.py b/pyrogram/client/session_storage/string.py deleted file mode 100644 index 11051323..00000000 --- a/pyrogram/client/session_storage/string.py +++ /dev/null @@ -1,46 +0,0 @@ -import base64 -import binascii -import struct - -import pyrogram -from . import MemorySessionStorage, SessionDoesNotExist - - -class StringSessionStorage(MemorySessionStorage): - """ - Packs session data as following (forcing little-endian byte order): - Char dc_id (1 byte, unsigned) - Boolean test_mode (1 byte) - Long long user_id (8 bytes, signed) - Boolean is_bot (1 byte) - Bytes auth_key (256 bytes) - - Uses Base64 encoding for printable representation - """ - PACK_FORMAT = '. -from .abstract import SessionStorage, SessionDoesNotExist -from .memory import MemorySessionStorage -from .json import JsonSessionStorage -from .string import StringSessionStorage -from .sqlite import SQLiteSessionStorage +from .memory_storage import MemoryStorage +from .file_storage import FileStorage +from .storage import Storage diff --git a/pyrogram/client/style/html.py b/pyrogram/client/style/html.py index 894dbd6c..9c0a372c 100644 --- a/pyrogram/client/style/html.py +++ b/pyrogram/client/style/html.py @@ -31,16 +31,14 @@ from pyrogram.api.types import ( ) from pyrogram.errors import PeerIdInvalid from . import utils -from ..session_storage import SessionStorage class HTML: HTML_RE = re.compile(r"<(\w+)(?: href=([\"'])([^<]+)\2)?>([^>]+)") MENTION_RE = re.compile(r"tg://user\?id=(\d+)") - def __init__(self, session_storage: SessionStorage, client: "pyrogram.BaseClient" = None): + def __init__(self, client: "pyrogram.BaseClient" = None): self.client = client - self.session_storage = session_storage def parse(self, message: str): entities = [] @@ -56,9 +54,10 @@ class HTML: if mention: user_id = int(mention.group(1)) + try: - input_user = self.session_storage.get_peer_by_id(user_id) - except KeyError: + input_user = self.client.resolve_peer(user_id) + except PeerIdInvalid: input_user = None entity = ( diff --git a/pyrogram/client/style/markdown.py b/pyrogram/client/style/markdown.py index 68b54bbb..adb86e94 100644 --- a/pyrogram/client/style/markdown.py +++ b/pyrogram/client/style/markdown.py @@ -31,7 +31,6 @@ from pyrogram.api.types import ( ) from pyrogram.errors import PeerIdInvalid from . import utils -from ..session_storage import SessionStorage class Markdown: @@ -55,9 +54,8 @@ class Markdown: )) MENTION_RE = re.compile(r"tg://user\?id=(\d+)") - def __init__(self, session_storage: SessionStorage, client: "pyrogram.BaseClient" = None): + def __init__(self, client: "pyrogram.BaseClient" = None): self.client = client - self.session_storage = session_storage def parse(self, message: str): message = utils.add_surrogates(str(message or "")).strip() @@ -73,9 +71,10 @@ class Markdown: if mention: user_id = int(mention.group(1)) + try: - input_user = self.session_storage.get_peer_by_id(user_id) - except KeyError: + input_user = self.client.resolve_peer(user_id) + except PeerIdInvalid: input_user = None entity = ( 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 bd7f0f26..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,8 +116,12 @@ class Session: def start(self): while True: - self.connection = Connection(self.dc_id, self.client.session_storage.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() From 6177abbfa4d09b63eee17907a3d2462ac2e14a32 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:04:06 +0200 Subject: [PATCH 15/22] Add Storage abstract class --- pyrogram/client/storage/storage.py | 98 ++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 pyrogram/client/storage/storage.py 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 From 6cc9688e4921a15c82667264eea624844cb48be9 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:04:35 +0200 Subject: [PATCH 16/22] Implement FileStorage and MemoryStorage engines --- pyrogram/client/storage/file_storage.py | 102 +++++++++ pyrogram/client/storage/memory_storage.py | 241 ++++++++++++++++++++++ pyrogram/client/storage/schema.sql | 34 +++ 3 files changed, 377 insertions(+) create mode 100644 pyrogram/client/storage/file_storage.py create mode 100644 pyrogram/client/storage/memory_storage.py create mode 100644 pyrogram/client/storage/schema.sql 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 From 8465c4a97798de7d09bab1048d4154c470d59278 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:06:37 +0200 Subject: [PATCH 17/22] Instruct Python to add schema.sql file to the package --- MANIFEST.in | 2 +- setup.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) 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/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, From edaced35a758a84dff17328cb0095cfb525b5449 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:07:22 +0200 Subject: [PATCH 18/22] Use base64.urlsafe_b64encode/decode instead of manually passing altchars --- pyrogram/client/ext/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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: From 30192de1ad493acd896448e727443d7b03e65761 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:10:37 +0200 Subject: [PATCH 19/22] Update pyrogram/client to accommodate Storage Engines --- pyrogram/client/client.py | 231 ++++++++++-------- pyrogram/client/ext/base_client.py | 8 +- pyrogram/client/ext/syncer.py | 15 +- .../client/methods/contacts/get_contacts.py | 1 - 4 files changed, 141 insertions(+), 114 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 885c3334..2d18d178 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -16,7 +16,6 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . -import binascii import logging import math import mimetypes @@ -51,10 +50,7 @@ from pyrogram.errors import ( from pyrogram.session import Auth, Session from .ext import utils, Syncer, BaseClient, Dispatcher from .methods import Methods -from .session_storage import ( - SessionDoesNotExist, SessionStorage, MemorySessionStorage, JsonSessionStorage, - StringSessionStorage, SQLiteSessionStorage -) +from .storage import Storage, FileStorage, MemoryStorage log = logging.getLogger(__name__) @@ -64,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 @@ -179,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, @@ -226,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): @@ -266,50 +278,32 @@ class Client(Methods, BaseClient): if self.is_started: raise ConnectionError("Client has already been started") - if isinstance(self.session_storage, JsonSessionStorage): - if self.BOT_TOKEN_RE.match(self.session_storage._session_name): - self.session_storage.is_bot = True - self.bot_token = self.session_storage._session_name - self.session_storage._session_name = self.session_storage._session_name.split(":")[0] - warnings.warn('\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.session_storage.dc_id, - self.session_storage.auth_key - ) + self.session = Session(self, self.storage.dc_id, self.storage.auth_key) self.session.start() self.is_started = True try: - if self.session_storage.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.session_storage.is_bot = True + self.storage.is_bot = True self.authorize_bot() - self.save_session() - - if not self.session_storage.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.session_storage.date) > Client.OFFLINE_SLEEP: - self.session_storage.clear_cache() - + if abs(now - self.storage.date) > Client.OFFLINE_SLEEP: self.get_initial_dialogs() self.get_contacts() else: @@ -508,20 +502,15 @@ class Client(Methods, BaseClient): except UserMigrate as e: self.session.stop() - self.session_storage.dc_id = e.x - self.session_storage.auth_key = Auth(self.session_storage.dc_id, self.session_storage.test_mode, - self.ipv6, self._proxy).create() - - self.session = Session( - self, - self.session_storage.dc_id, - self.session_storage.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.session_storage.user_id = r.user.id + self.storage.user_id = r.user.id print("Logged in successfully as @{}".format(r.user.username)) @@ -562,20 +551,10 @@ class Client(Methods, BaseClient): except (PhoneMigrate, NetworkMigrate) as e: self.session.stop() - self.session_storage.dc_id = e.x + self.storage.dc_id = e.x + self.storage.auth_key = Auth(self, self.storage.dc_id).create() - self.session_storage.auth_key = Auth( - self.session_storage.dc_id, - self.session_storage.test_mode, - self.ipv6, - self._proxy - ).create() - - self.session = Session( - self, - self.session_storage.dc_id, - self.session_storage.auth_key - ) + self.session = Session(self, self.storage.dc_id, self.storage.auth_key) self.session.start() except (PhoneNumberInvalid, PhoneNumberBanned) as e: @@ -755,13 +734,13 @@ class Client(Methods, BaseClient): ) self.password = None - self.session_storage.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, @@ -770,11 +749,57 @@ class Client(Methods, BaseClient): ] ) -> bool: is_min = False + parsed_peers = [] - for entity in entities: - if isinstance(entity, (types.User, types.Channel, types.ChannelForbidden)) and not entity.access_hash: + for peer in peers: + username = None + phone_number = None + + 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 + + if username is not None: + 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 + + username = getattr(peer, "username", None) + + if peer.broadcast: + peer_type = "channel" + else: + peer_type = "supergroup" + + if access_hash is None: + is_min = True + continue + + if username is not None: + username = username.lower() + else: continue - self.session_storage.cache_peer(entity) + + parsed_peers.append((peer_id, access_hash, peer_type, username, phone_number)) + + self.storage.update_peers(parsed_peers) return is_min @@ -1035,12 +1060,23 @@ class Client(Methods, BaseClient): self.plugins = None def load_session(self): - try: - self.session_storage.load() - except SessionDoesNotExist: - log.info('Could not load session "{}", initiate new one'.format(self.session_name)) - self.session_storage.auth_key = Auth(self.session_storage.dc_id, self.session_storage.test_mode, - self.ipv6, self._proxy).create() + self.storage.open() + + 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 + ]) + + if session_empty: + self.storage.dc_id = 1 + self.storage.date = 0 + + 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: @@ -1164,9 +1200,6 @@ class Client(Methods, BaseClient): log.warning('[{}] No plugin loaded from "{}"'.format( self.session_name, root)) - def save_session(self): - self.session_storage.save() - def get_initial_dialogs_chunk(self, offset_date: int = 0): while True: try: @@ -1184,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(self.session_storage.peers_count())) + log.info("Total peers: {}".format(self.storage.peers_count)) return r def get_initial_dialogs(self): @@ -1222,7 +1255,7 @@ class Client(Methods, BaseClient): KeyError: In case the peer doesn't exist in the internal database. """ try: - return self.session_storage.get_peer_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"): @@ -1234,7 +1267,7 @@ class Client(Methods, BaseClient): int(peer_id) except ValueError: try: - self.session_storage.get_peer_by_username(peer_id) + return self.storage.get_peer_by_username(peer_id) except KeyError: self.send( functions.contacts.ResolveUsername( @@ -1242,10 +1275,10 @@ class Client(Methods, BaseClient): ) ) - return self.session_storage.get_peer_by_username(peer_id) + return self.storage.get_peer_by_username(peer_id) else: try: - return self.session_storage.get_peer_by_phone(peer_id) + return self.storage.get_peer_by_phone_number(peer_id) except KeyError: raise PeerIdInvalid @@ -1253,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 + )] ) ) ) @@ -1261,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: @@ -1272,7 +1311,7 @@ class Client(Methods, BaseClient): ) try: - return self.session_storage.get_peer_by_id(peer_id) + return self.storage.get_peer_by_id(peer_id) except KeyError: raise PeerIdInvalid @@ -1347,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.session_storage.dc_id, self.session_storage.auth_key, is_media=True) + session = Session(self, self.storage.dc_id, self.storage.auth_key, is_media=True) session.start() try: @@ -1433,19 +1472,14 @@ class Client(Methods, BaseClient): session = self.media_sessions.get(dc_id, None) if session is None: - if dc_id != self.session_storage.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.session_storage.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() @@ -1458,12 +1492,7 @@ class Client(Methods, BaseClient): ) ) else: - session = Session( - self, - dc_id, - self.session_storage.auth_key, - is_media=True - ) + session = Session(self, dc_id, self.storage.auth_key, is_media=True) session.start() @@ -1548,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.session_storage.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() @@ -1650,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 9276b0eb..88623f4a 100644 --- a/pyrogram/client/ext/base_client.py +++ b/pyrogram/client/ext/base_client.py @@ -87,13 +87,13 @@ class BaseClient: mime_types_to_extensions[mime_type] = " ".join(extensions) - def __init__(self, session_storage: SessionStorage): - self.session_storage = session_storage + def __init__(self): + self.storage = None self.rnd_id = MsgId - self.markdown = Markdown(self.session_storage, self) - self.html = HTML(self.session_storage, self) + self.markdown = Markdown(self) + self.html = HTML(self) self.session = None self.media_sessions = {} diff --git a/pyrogram/client/ext/syncer.py b/pyrogram/client/ext/syncer.py index 9e7d2303..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,10 +75,13 @@ class Syncer: @classmethod def sync(cls, client): - client.session_storage.date = int(time.time()) try: - client.session_storage.save(sync=True) + start = time.time() + client.storage.save() except Exception as e: log.critical(e, exc_info=True) else: - log.info("Synced {}".format(client.session_name)) + log.info('Synced "{}" in {:.6} ms'.format( + client.storage.name, + (time.time() - start) * 1000 + )) diff --git a/pyrogram/client/methods/contacts/get_contacts.py b/pyrogram/client/methods/contacts/get_contacts.py index 79677563..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(self.session_storage.contacts_count())) return pyrogram.List(pyrogram.User._parse(self, user) for user in contacts.users) From 0be0e2da5615cd9a403889ce487d3e8878451949 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:11:25 +0200 Subject: [PATCH 20/22] Add export_session_string method to docs --- docs/source/api/methods.rst | 2 ++ 1 file changed, 2 insertions(+) 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() From 1f04ce38fc4f62ec786efb9c8080993bdfe127f2 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:11:53 +0200 Subject: [PATCH 21/22] Fix glossary term --- docs/source/glossary.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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. From d1cd21916a7d20f71afa78767ca660f9f3fa8c8d Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 19 Jun 2019 16:12:06 +0200 Subject: [PATCH 22/22] Add storage-engines.rst page to docs --- docs/source/index.rst | 1 + docs/source/topics/storage-engines.rst | 95 ++++++++++++++++++++++++++ 2 files changed, 96 insertions(+) create mode 100644 docs/source/topics/storage-engines.rst 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()) +