From cc3d4145d85ad3b0b31faba38206d0e1f4dbaab5 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Sun, 19 Sep 2021 18:15:09 +0200
Subject: [PATCH] Update and persist session state on successful login

---
 .gitignore                             |  1 +
 telethon/_client/auth.py               | 21 ++++++++++++++++-----
 telethon/_client/telegrambaseclient.py | 11 +----------
 telethon/_client/telegramclient.py     |  4 ++--
 4 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/.gitignore b/.gitignore
index 1e497d62..3856ed42 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,6 +6,7 @@
 
 # User session
 *.session
+sessions/
 /usermedia/
 
 # Builds and testing
diff --git a/telethon/_client/auth.py b/telethon/_client/auth.py
index 296656f7..365e7257 100644
--- a/telethon/_client/auth.py
+++ b/telethon/_client/auth.py
@@ -253,7 +253,7 @@ async def sign_in(
         self._tos = result.terms_of_service
         raise errors.PhoneNumberUnoccupiedError(request=request)
 
-    return _on_login(self, result.user)
+    return await _update_session_state(self, result.user)
 
 async def sign_up(
         self: 'TelegramClient',
@@ -306,17 +306,28 @@ async def sign_up(
         await self(
             _tl.fn.help.AcceptTermsOfService(self._tos.id))
 
-    return _on_login(self, result.user)
+    return await _update_session_state(self, result.user)
 
-def _on_login(self, user):
+async def _update_session_state(self, user, save=True):
     """
     Callback called whenever the login or sign up process completes.
     Returns the input user parameter.
     """
-    self._bot = bool(user.bot)
-    self._self_input_peer = utils.get_input_peer(user, allow_self=False)
     self._authorized = True
 
+    self._session_state.user_id = user.id
+    self._session_state.bot = user.bot
+
+    state = await self(_tl.fn.updates.GetState())
+    self._session_state.pts = state.pts
+    self._session_state.qts = state.qts
+    self._session_state.date = int(state.date.timestamp())
+    self._session_state.seq = state.seq
+
+    await self.session.set_state(self._session_state)
+    if save:
+        await self.session.save()
+
     return user
 
 async def send_code_request(
diff --git a/telethon/_client/telegrambaseclient.py b/telethon/_client/telegrambaseclient.py
index a447ae61..38ad77ed 100644
--- a/telethon/_client/telegrambaseclient.py
+++ b/telethon/_client/telegrambaseclient.py
@@ -385,16 +385,7 @@ async def connect(self: 'TelegramClient') -> None:
         # logged in). Attempt to fetch the user now. If it works, also get the update state.
         me = await self.get_me()
         if me:
-            self._session_state.user_id = me.id
-            self._session_state.bot = me.bot
-
-            state = await self(_tl.fn.updates.GetState())
-            self._session_state.pts = state.pts
-            self._session_state.qts = state.qts
-            self._session_state.date = int(state.date.timestamp())
-            self._session_state.seq = state.seq
-
-            await self.session.set_state(self._session_state)
+            await self._update_session_state(me, save=False)
 
     await self.session.save()
 
diff --git a/telethon/_client/telegramclient.py b/telethon/_client/telegramclient.py
index 10db3557..15a7db75 100644
--- a/telethon/_client/telegramclient.py
+++ b/telethon/_client/telegramclient.py
@@ -3598,8 +3598,8 @@ class TelegramClient:
     async def _handle_auto_reconnect(self: 'TelegramClient'):
         return await updates._handle_auto_reconnect(**locals())
 
-    def _self_id(self: 'TelegramClient') -> typing.Optional[int]:
-        return users._self_id(**locals())
+    async def _update_session_state(self, user, save=True):
+        return await auth._update_session_state(**locals())
 
     # endregion Private