mir/pyrogram
2
0
Fork 0
mirror of https://github.com/pyrogram/pyrogram synced 2025-08-28 12:57:52 +00:00
Code Issues Releases Wiki Activity

Merge pull request #220 from bakatrouble/session_storage

Browse Source 
Implement extendable session storage, move to SQLite as main storage (WIP)
This commit is contained in:
Dan 2019-06-19 16:12:36 +02:00 committed by GitHub
commit 2ff74270ad
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
18 changed files with 745 additions and 247 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
MANIFEST.in
View File

@ -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

2
docs/source/api/methods.rst
View File

@ -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()

2
docs/source/glossary.rst
View File

@ -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.

1
docs/source/index.rst
View File

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

95
docs/source/topics/storage-engines.rst Normal file
View File

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

289
pyrogram/client/client.py
View File

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

10
pyrogram/client/ext/base_client.py
View File

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

53
pyrogram/client/ext/syncer.py
View File

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

6
pyrogram/client/ext/utils.py
View File

@ -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:

1
pyrogram/client/methods/contacts/get_contacts.py
View File

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

21
pyrogram/client/storage/__init__.py Normal file
View File

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

102
pyrogram/client/storage/file_storage.py Normal file
View File

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

241
pyrogram/client/storage/memory_storage.py Normal file
View File

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

34
pyrogram/client/storage/schema.sql Normal file
View File

@ -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;

98
pyrogram/client/storage/storage.py Normal file
View File

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

10
pyrogram/session/auth.py
View File

@ -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

22
pyrogram/session/session.py
View File

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

3
setup.py
View File

@ -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,