├── .gitignore ├── LICENSE ├── README.rst ├── requirements.txt ├── setup.py └── teleredis ├── __init__.py └── teleredis.py /.gitignore: -------------------------------------------------------------------------------- 1 | build/ 2 | dist/ 3 | releases/ 4 | *.egg-info/ 5 | *.egg 6 | __pycache__/ 7 | *.py[cod] 8 | .installed.cfg 9 | .idea/ 10 | *.pyc 11 | *.pyo 12 | *.pyd 13 | *.log 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Konstantin M. (ezdev128) 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | Telethon Redis session 2 | =========================== 3 | 4 | A `Telethon`_ session storage implementation backed by Redis. 5 | 6 | .. _Telethon: https://github.com/LonamiWebs/Telethon 7 | 8 | Usage 9 | ----- 10 | This session implementation can store multiple Sessions in the same key hive. 11 | 12 | 13 | 14 | Installing 15 | ---------- 16 | 17 | .. code-block:: sh 18 | 19 | pip3 install teleredis 20 | 21 | 22 | Upgrading 23 | ---------- 24 | 25 | .. code-block:: sh 26 | 27 | pip3 install -U teleredis 28 | 29 | 30 | Quick start 31 | ----------- 32 | .. code-block:: python 33 | 34 | from telethon import TelegramClient 35 | from teleredis import RedisSession 36 | import redis 37 | 38 | # These example values won't work. You must get your own api_id and 39 | # api_hash from https://my.telegram.org, under API Development. 40 | api_id = 12345 41 | api_hash = '0123456789abcdef0123456789abcdef' 42 | 43 | redis_connector = redis.Redis(host='localhost', port=6379, db=0, decode_responses=False) 44 | session = RedisSession('session_name', redis_connector) 45 | client = TelegramClient(session, api_id, api_hash) 46 | client.start() 47 | 48 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | redis 2 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import setuptools 4 | import re 5 | 6 | with open("teleredis/__init__.py", encoding="utf-8") as f: 7 | version = re.search(r"^__version__[\s\t=]*[\"']*([\w\d.\-_+]+)[\"']*$", 8 | f.read(), re.M).group(1) 9 | 10 | package_name = "teleredis" 11 | 12 | setuptools.setup( 13 | name=package_name, 14 | packages=[package_name], 15 | version=version, 16 | 17 | url="https://github.com/ezdev128/telethon-session-redis", 18 | download_url="https://github.com/ezdev128/telethon-session-redis/releases", 19 | 20 | author="Konstantin M.", 21 | author_email="ezdev128@yandex.com", 22 | 23 | description="Redis backend for Telethon session storage", 24 | long_description=open("README.rst", encoding="utf-8").read(), 25 | 26 | 27 | license="MIT", 28 | 29 | classifiers=[ 30 | # 3 - Alpha 31 | # 4 - Beta 32 | # 5 - Production/Stable 33 | "Development Status :: 4 - Beta", 34 | "Intended Audience :: Developers", 35 | "License :: OSI Approved :: MIT License", 36 | "Programming Language :: Python", 37 | "Programming Language :: Python :: 3", 38 | "Programming Language :: Python :: 3.4", 39 | "Programming Language :: Python :: 3.5", 40 | "Programming Language :: Python :: 3.6", 41 | ], 42 | keywords="telegram session sessions redis", 43 | python_requires="~=3.4", 44 | 45 | install_requires=[ 46 | "redis>=2.0", 47 | "Telethon>=0.17" 48 | ], 49 | ) 50 | -------------------------------------------------------------------------------- /teleredis/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from .teleredis import RedisSession, PackFunction, UnpackFunction 3 | 4 | __version__ = "0.1.3" 5 | __author__ = "Konstantin M. " 6 | __all__ = [ 7 | "RedisSession", 8 | ] 9 | 10 | -------------------------------------------------------------------------------- /teleredis/teleredis.py: -------------------------------------------------------------------------------- 1 | 2 | from telethon.sessions.memory import MemorySession, _SentFileType 3 | from telethon.crypto import AuthKey 4 | from telethon import utils 5 | from telethon.tl.types import ( 6 | InputPhoto, InputDocument, PeerUser, PeerChat, PeerChannel 7 | ) 8 | import logging 9 | import json 10 | import base64 11 | import time 12 | import redis 13 | import pickle 14 | from enum import Enum 15 | 16 | 17 | class PackFunction(Enum): 18 | JSON = 0 19 | PICKLE = 1 20 | 21 | 22 | class UnpackFunction(Enum): 23 | JSON = 0 24 | PICKLE = 1 25 | 26 | 27 | DEFAULT_TS_STR_FORMAT = "%F %T" 28 | DEFAULT_HIVE_PREFIX = "telethon:client" 29 | DEFAULT_PACK_FUNC = PackFunction.JSON 30 | DEFAULT_UNPACK_FUNC = UnpackFunction.JSON 31 | 32 | 33 | __log__ = logging.getLogger(__name__) 34 | 35 | 36 | class RedisSession(MemorySession): 37 | 38 | session_name = None 39 | redis_connection = None 40 | hive_prefix = None 41 | sess_prefix = None 42 | use_indents = True 43 | add_timestamps = False 44 | ts_format = None 45 | pack_func = None 46 | unpack_func = None 47 | 48 | def __init__(self, session_name=None, redis_connection=None, hive_prefix=None, 49 | use_indents=False, add_timestamps=False, ts_format=None, 50 | pack_func=None, unpack_func=None): 51 | if not isinstance(session_name, (str, bytes)): 52 | raise TypeError("Session name must be a string or bytes.") 53 | 54 | if not redis_connection or not isinstance(redis_connection, redis.StrictRedis): 55 | raise TypeError('The given redis_connection must be a Redis instance.') 56 | 57 | super().__init__() 58 | 59 | self.session_name = session_name if isinstance(session_name, str) else session_name.decode() 60 | self.redis_connection = redis_connection 61 | self.hive_prefix = hive_prefix or DEFAULT_HIVE_PREFIX 62 | self.use_indents = use_indents 63 | self.add_timestamps = add_timestamps 64 | self.ts_format = ts_format or DEFAULT_TS_STR_FORMAT 65 | self.pack_func = pack_func or DEFAULT_PACK_FUNC 66 | self.unpack_func = unpack_func or DEFAULT_UNPACK_FUNC 67 | self.sess_prefix = "{}:{}".format(self.hive_prefix, self.session_name) 68 | self.save_entities = True 69 | self.feed_session() 70 | 71 | def _pack(self, o, **kwargs): 72 | if self.pack_func == PackFunction.JSON: 73 | if self.use_indents: 74 | kwargs["indent"] = 2 75 | return json.dumps(o, **kwargs) if self.pack_func == PackFunction.JSON else pickle.dumps(o, **kwargs) 76 | 77 | def _unpack(self, o, **kwargs): 78 | if self.unpack_func == UnpackFunction.JSON and isinstance(o, bytes): 79 | o = o.decode() 80 | return json.loads(o, **kwargs) if self.unpack_func == UnpackFunction.JSON else pickle.loads(o, **kwargs) 81 | 82 | def feed_session(self): 83 | try: 84 | s = self._get_sessions() 85 | if len(s) == 0: 86 | self._auth_key = AuthKey(data=bytes()) 87 | return 88 | 89 | s = self.redis_connection.get(s[-1]) 90 | if not s: 91 | # No sessions 92 | self._auth_key = AuthKey(data=bytes()) 93 | return 94 | 95 | s = self._unpack(s) 96 | self._dc_id = s["dc_id"] 97 | self._server_address = s["server_address"] 98 | self._port = s["port"] 99 | auth_key = base64.standard_b64decode(s["auth_key"]) 100 | self._auth_key = AuthKey(data=auth_key) 101 | except Exception as ex: 102 | __log__.exception(ex.args) 103 | 104 | def _update_sessions(self): 105 | """ 106 | Stores session into redis. 107 | """ 108 | auth_key = self._auth_key.key if self._auth_key else bytes() 109 | if not self._dc_id: 110 | return 111 | 112 | s = { 113 | "dc_id": self._dc_id, 114 | "server_address": self._server_address, 115 | "port": self._port, 116 | "auth_key": base64.standard_b64encode(auth_key).decode(), 117 | } 118 | 119 | if self.add_timestamps: 120 | s.update({ 121 | "ts_ts": time.time(), 122 | "ts_str": time.strftime(DEFAULT_TS_STR_FORMAT, time.localtime()), 123 | }) 124 | 125 | key = "{}:sessions:{}".format(self.sess_prefix, self._dc_id) 126 | try: 127 | self.redis_connection.set(key, self._pack(s)) 128 | except Exception as ex: 129 | __log__.exception(ex.args) 130 | 131 | def set_dc(self, dc_id, server_address, port): 132 | """ 133 | Sets the information of the data center address and port that 134 | the library should connect to, as well as the data center ID, 135 | which is currently unused. 136 | """ 137 | super().set_dc(dc_id, server_address, port) 138 | self._update_sessions() 139 | 140 | auth_key = bytes() 141 | 142 | if not self._dc_id: 143 | self._auth_key = AuthKey(data=auth_key) 144 | return 145 | 146 | key = "{}:sessions:{}".format(self.sess_prefix, self._dc_id) 147 | s = self.redis_connection.get(key) 148 | if s: 149 | s = self._unpack(s) 150 | auth_key = base64.standard_b64decode(s["auth_key"]) 151 | self._auth_key = AuthKey(data=auth_key) 152 | 153 | @MemorySession.auth_key.setter 154 | def auth_key(self, value): 155 | """ 156 | Sets the ``AuthKey`` to be used for the saved data center. 157 | """ 158 | self._auth_key = value 159 | self._update_sessions() 160 | 161 | def list_sessions(self): 162 | """ 163 | Lists available sessions. Not used by the library itself. 164 | """ 165 | return self._get_sessions(strip_prefix=True) 166 | 167 | def process_entities(self, tlo): 168 | """ 169 | Processes the input ``TLObject`` or ``list`` and saves 170 | whatever information is relevant (e.g., ID or access hash). 171 | """ 172 | 173 | if not self.save_entities: 174 | return 175 | 176 | rows = self._entities_to_rows(tlo) 177 | if not rows or len(rows) == 0 or len(rows[0]) == 0: 178 | return 179 | 180 | try: 181 | rows = rows[0] 182 | key = "{}:entities:{}".format(self.sess_prefix, rows[0]) 183 | s = { 184 | "id": rows[0], 185 | "hash": rows[1], 186 | "username": rows[2], 187 | "phone": rows[3], 188 | "name": rows[4], 189 | } 190 | 191 | if self.add_timestamps: 192 | s.update({ 193 | "ts_ts": time.time(), 194 | "ts_str": time.strftime(DEFAULT_TS_STR_FORMAT, time.localtime()), 195 | }) 196 | 197 | self.redis_connection.set(key, self._pack(s)) 198 | except Exception as ex: 199 | __log__.exception(ex.args) 200 | 201 | def _get_entities(self, strip_prefix=False): 202 | """ 203 | Returns list of entities. if strip_prefix is False - returns redis keys, 204 | else returns list of id's 205 | """ 206 | key_pattern = "{}:{}:entities:".format(self.hive_prefix, self.session_name) 207 | try: 208 | entities = self.redis_connection.keys(key_pattern+"*") 209 | if not strip_prefix: 210 | return entities 211 | return [s.decode().replace(key_pattern, "") for s in entities] 212 | except Exception as ex: 213 | __log__.exception(ex.args) 214 | return [] 215 | 216 | def _get_sessions(self, strip_prefix=False): 217 | """ 218 | Returns list of sessions. if strip_prefix is False - returns redis keys, 219 | else returns list of id's 220 | """ 221 | key_pattern = "{}:{}:sessions:".format(self.hive_prefix, self.session_name) 222 | try: 223 | sessions = self.redis_connection.keys(key_pattern+"*") 224 | return [s.decode().replace(key_pattern, "") if strip_prefix else s.decode() for s in sessions] 225 | except Exception as ex: 226 | __log__.exception(ex.args) 227 | return [] 228 | 229 | def get_entity_rows_by_phone(self, phone): 230 | try: 231 | for key in self._get_entities(): 232 | entity = self._unpack(self.redis_connection.get(key)) 233 | if "phone" in entity and entity["phone"] == phone: 234 | return entity["id"], entity["hash"] 235 | except Exception as ex: 236 | __log__.exception(ex.args) 237 | return None 238 | 239 | def get_entity_rows_by_username(self, username): 240 | try: 241 | for key in self._get_entities(): 242 | entity = self._unpack(self.redis_connection.get(key)) 243 | if "username" in entity and entity["username"] == username: 244 | return entity["id"], entity["hash"] 245 | except Exception as ex: 246 | __log__.exception(ex.args) 247 | return None 248 | 249 | def get_entity_rows_by_name(self, name): 250 | try: 251 | for key in self._get_entities(): 252 | entity = self._unpack(self.redis_connection.get(key)) 253 | if "name" in entity and entity["name"] == name: 254 | return entity["id"], entity["hash"] 255 | except Exception as ex: 256 | __log__.exception(ex.args) 257 | 258 | return None 259 | 260 | def get_entity_rows_by_id(self, entity_id, exact=True): 261 | if exact: 262 | key = "{}:entities:{}".format(self.sess_prefix, entity_id) 263 | s = self.redis_connection.get(key) 264 | if not s: 265 | return None 266 | try: 267 | s = self._unpack(s) 268 | return entity_id, s["hash"] 269 | except Exception as ex: 270 | __log__.exception(ex.args) 271 | return None 272 | else: 273 | ids = ( 274 | utils.get_peer_id(PeerUser(entity_id)), 275 | utils.get_peer_id(PeerChat(entity_id)), 276 | utils.get_peer_id(PeerChannel(entity_id)) 277 | ) 278 | 279 | try: 280 | for key in self._get_entities(): 281 | entity = self._unpack(self.redis_connection.get(key)) 282 | if "id" in entity and entity["id"] in ids: 283 | return entity["id"], entity["hash"] 284 | except Exception as ex: 285 | __log__.exception(ex.args) 286 | 287 | def get_file(self, md5_digest, file_size, cls): 288 | key = "{}:sent_files:{}".format(self.sess_prefix, md5_digest) 289 | s = self.redis_connection.get(key) 290 | if s: 291 | try: 292 | s = self._unpack(s) 293 | return md5_digest, file_size \ 294 | if s["file_size"] == file_size and s["type"] == _SentFileType.from_type(cls).value \ 295 | else None 296 | except Exception as ex: 297 | __log__.exception(ex.args) 298 | return None 299 | 300 | def cache_file(self, md5_digest, file_size, instance): 301 | if not isinstance(instance, (InputDocument, InputPhoto)): 302 | raise TypeError('Cannot cache {} instance'.format(type(instance))) 303 | 304 | key = "{}:sent_files:{}".format(self.sess_prefix, md5_digest) 305 | s = { 306 | "md5_digest": md5_digest, 307 | "file_size": file_size, 308 | "type": _SentFileType.from_type(type(instance)).value, 309 | "id": instance.id, 310 | "hash": instance.access_hash, 311 | } 312 | 313 | if self.add_timestamps: 314 | s.update({ 315 | "ts_ts": time.time(), 316 | "ts_str": time.strftime(DEFAULT_TS_STR_FORMAT, time.localtime()), 317 | }) 318 | 319 | try: 320 | self.redis_connection.set(key, self._pack(s)) 321 | except Exception as ex: 322 | __log__.exception(ex.args) 323 | --------------------------------------------------------------------------------