Fix-up 2b090f8 not actually resaving new auth_keys

This commit is contained in:
Lonami Exo 2018-06-27 10:15:59 +02:00
parent b834b6c16c
commit 58aadce5ba
2 changed files with 18 additions and 7 deletions

View File

@ -208,7 +208,8 @@ class TelegramBaseClient(abc.ABC):
state, connection, self._loop,
retries=self._connection_retries,
auto_reconnect=self._auto_reconnect,
update_callback=self._handle_update
update_callback=self._handle_update,
auth_key_callback=self._auth_key_callback
)
# Cache :tl:`ExportedAuthorization` as ``dc_id: MTProtoState``
@ -271,7 +272,6 @@ class TelegramBaseClient(abc.ABC):
"""
Connects to Telegram.
"""
old_auth = self.session.auth_key
await self._sender.connect(
self.session.server_address, self.session.port)
@ -280,10 +280,6 @@ class TelegramBaseClient(abc.ABC):
self._updates_handle = self._loop.create_task(self._update_loop())
if old_auth != self._sender.state.auth_key:
self.session.auth_key = self._sender.state.auth_key
self.session.save()
def is_connected(self):
"""
Returns ``True`` if the user has connected.
@ -339,6 +335,15 @@ class TelegramBaseClient(abc.ABC):
await self._disconnect()
return await self.connect()
def _auth_key_callback(self, auth_key):
"""
Callback from the sender whenever it needed to generate a
new authorization key. This means we are not authorized.
"""
self._authorized = None
self.session.auth_key = auth_key
self.session.save()
# endregion
# region Working with different connections/Data Centers

View File

@ -41,7 +41,8 @@ class MTProtoSender:
key exists yet.
"""
def __init__(self, state, connection, loop, *,
retries=5, auto_reconnect=True, update_callback=None):
retries=5, auto_reconnect=True, update_callback=None,
auth_key_callback=None):
self.state = state
self._connection = connection
self._loop = loop
@ -50,6 +51,7 @@ class MTProtoSender:
self._retries = retries
self._auto_reconnect = auto_reconnect
self._update_callback = update_callback
self._auth_key_callback = auth_key_callback
# Whether the user has explicitly connected or disconnected.
#
@ -255,6 +257,10 @@ class MTProtoSender:
__log__.debug('New auth_key attempt {}...'.format(retry))
self.state.auth_key, self.state.time_offset =\
await authenticator.do_authentication(plain)
if self._auth_key_callback:
self._auth_key_callback(self.state.auth_key)
break
except (SecurityError, AssertionError) as e:
__log__.warning('Attempt {} at new auth_key failed: {}'