From 9c4e9e166e528d2ef990bcb3f2093a877d65b642 Mon Sep 17 00:00:00 2001
From: bakatrouble <bakatrouble@gmail.com>
Date: Fri, 22 Feb 2019 02:13:51 +0300
Subject: [PATCH] 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 = '<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))