Merge pyrogram#221, string sessions now work for bots too

bakatrouble · bakatrouble · commit 9c4e9e166e52 · 2019-02-22T02:13:51.000+03:00
diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
@@ -281,14 +281,15 @@ def start(self):
         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
@@ -70,7 +70,6 @@ class StopTransmission(StopIteration):
 
     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
@@ -35,6 +35,7 @@ def __init__(self, client: 'pyrogram.client.BaseClient', session_data):
         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
@@ -50,6 +50,7 @@ def load_session(self):
         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 @@ def save_session(self, sync=False):
             '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
@@ -60,6 +60,14 @@ def date(self) -> int:
     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
@@ -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 = '<B?q256s'
+    PACK_FORMAT = '<B?q?256s'
+
+    def _unpack(self, data):
+        return struct.unpack(self.PACK_FORMAT, data)
+
+    def _pack(self):
+        return struct.pack(self.PACK_FORMAT, self.dc_id, self.test_mode, self.user_id, self.is_bot, self.auth_key)
 
     def load_session(self):
         try:
             session_string = self.session_data[1:]
             session_string += '=' * (4 - len(session_string) % 4)  # restore padding
             decoded = base64.b64decode(session_string, b'-_')
-            self.dc_id, self.test_mode, self.user_id, self.auth_key = struct.unpack(self.PACK_FORMAT, decoded)
+            self.dc_id, self.test_mode, self.user_id, self.is_bot, self.auth_key = self._unpack(decoded)
         except (struct.error, binascii.Error):
             raise SessionDoesNotExist()
 
     def save_session(self, sync=False):
         if not sync:
-            packed = struct.pack(self.PACK_FORMAT, self.dc_id, self.test_mode, self.user_id, self.auth_key)
+            packed = self._pack()
             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))
