Save all found entities to the session file

telethon/tl/session.py

@@ -2,11 +2,15 @@ import json
 import os
 import platform
 import time
-from threading import Lock
 from base64 import b64encode, b64decode
 from os.path import isfile as file_exists
+from threading import Lock
 
-from .. import helpers as utils
+from .. import helpers, utils
+from ..tl.types import (
+    InputPeerUser, InputPeerChat, InputPeerChannel,
+    PeerUser, PeerChat, PeerChannel
+)
 
 
 class Session:
@@ -51,7 +55,7 @@ class Session:
         # Cross-thread safety
         self._lock = Lock()
 
-        self.id = utils.generate_random_long(signed=False)
+        self.id = helpers.generate_random_long(signed=False)
         self._sequence = 0
         self.time_offset = 0
         self._last_msg_id = 0  # Long
@@ -62,6 +66,8 @@ class Session:
         self.auth_key = None
         self.layer = 0
         self.salt = 0  # Unsigned long
+        self._input_entities = {}  # {marked_id: hash}
+        self._entities_lock = Lock()
 
     def save(self):
         """Saves the current session object as session_user_id.session"""
@@ -74,7 +80,8 @@ class Session:
                     'server_address': self.server_address,
                     'auth_key_data':
                         b64encode(self.auth_key.key).decode('ascii')
-                        if self.auth_key else None
+                        if self.auth_key else None,
+                    'entities': list(self._input_entities.items())
                 }, file)
 
     def delete(self):
@@ -122,6 +129,9 @@ class Session:
                         key = b64decode(data['auth_key_data'])
                         result.auth_key = AuthKey(data=key)
 
+                    for e_mid, e_hash in data.get('entities', []):
+                        result._input_entities[e_mid] = e_hash
+
             except (json.decoder.JSONDecodeError, UnicodeDecodeError):
                 pass
 
@@ -164,3 +174,43 @@ class Session:
         now = int(time.time())
         correct = correct_msg_id >> 32
         self.time_offset = correct - now
+
+    def add_entities(self, entities):
+        """Adds new input entities to the local database of them.
+           Unknown types will be ignored.
+        """
+        if not entities:
+            return False
+
+        new = {}
+        for e in entities:
+            try:
+                p = utils.get_input_peer(e)
+                new[utils.get_peer_id(p, add_mark=True)] = \
+                    getattr(p, 'access_hash', 0)  # chats won't have hash
+            except ValueError:
+                pass
+
+        with self._entities_lock:
+            before = len(self._input_entities)
+            self._input_entities.update(new)
+            return len(self._input_entities) != before
+
+    def get_input_entity(self, peer):
+        """Gets an input entity known its Peer or a marked ID,
+           or raises KeyError if not found/invalid.
+        """
+        if not isinstance(peer, int):
+            peer = utils.get_peer_id(peer, add_mark=True)
+
+        entity_hash = self._input_entities[peer]
+        entity_id, peer_class = utils.resolve_id(peer)
+
+        if peer_class == PeerUser:
+            return InputPeerUser(entity_id, entity_hash)
+        if peer_class == PeerChat:
+            return InputPeerChat(entity_id)
+        if peer_class == PeerChannel:
+            return InputPeerChannel(entity_id, entity_hash)
+
+        raise KeyError()