From df0e710fa1c18515dbb2126ffb181a13bab17cce Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Wed, 16 Feb 2022 12:23:19 +0100 Subject: [PATCH] Add a custom SentCode type --- telethon/_client/auth.py | 4 +- telethon/_client/telegramclient.py | 16 +++- telethon/types/__init__.py | 2 + telethon/types/_custom/__init__.py | 1 + telethon/types/_custom/auth.py | 120 +++++++++++++++++++++++++++++ 5 files changed, 138 insertions(+), 5 deletions(-) create mode 100644 telethon/types/_custom/auth.py diff --git a/telethon/_client/auth.py b/telethon/_client/auth.py index 67e6382e..a82aec03 100644 --- a/telethon/_client/auth.py +++ b/telethon/_client/auth.py @@ -335,7 +335,7 @@ async def _replace_session_state(self, *, save=True, **changes): async def send_code_request( self: 'TelegramClient', - phone: str) -> '_tl.auth.SentCode': + phone: str) -> 'SentCode': result = None phone = utils.parse_phone(phone) or self._phone phone_hash = self._phone_code_hash.get(phone) @@ -358,7 +358,7 @@ async def send_code_request( self._phone = phone - return result + return _custom.SentCode._new(result) async def qr_login(self: 'TelegramClient', ignored_ids: typing.List[int] = None) -> _custom.QRLogin: qr_login = _custom.QRLogin(self, ignored_ids or []) diff --git a/telethon/_client/telegramclient.py b/telethon/_client/telegramclient.py index eca8ba84..5ac3fa4d 100644 --- a/telethon/_client/telegramclient.py +++ b/telethon/_client/telegramclient.py @@ -504,7 +504,7 @@ class TelegramClient: @forward_call(auth.send_code_request) async def send_code_request( self: 'TelegramClient', - phone: str) -> '_tl.auth.SentCode': + phone: str) -> 'SentCode': """ Sends the Telegram code needed to login to the given phone number. @@ -513,14 +513,24 @@ class TelegramClient: The phone to which the code will be sent. Returns - An instance of :tl:`SentCode`. + An instance of `SentCode`. Example .. code-block:: python phone = '+34 123 123 123' sent = await client.send_code_request(phone) - print(sent) + print(sent.type) + + # Wait before resending sent.next_type, if any + if sent.next_type: + await asyncio.sleep(sent.timeout or 0) + resent = await client.send_code_request(phone) + print(sent.type) + + # Checking the code locally + code = input('Enter code: ') + print('Code looks OK:', resent.check(code)) """ @forward_call(auth.qr_login) diff --git a/telethon/types/__init__.py b/telethon/types/__init__.py index ac52ff6d..2fa8cd02 100644 --- a/telethon/types/__init__.py +++ b/telethon/types/__init__.py @@ -1,5 +1,7 @@ from .._misc.tlobject import TLObject, TLRequest from ._custom import ( + CodeType, + SentCode, AdminLogEvent, Draft, Dialog, diff --git a/telethon/types/_custom/__init__.py b/telethon/types/_custom/__init__.py index 00a0d00f..5637c1e2 100644 --- a/telethon/types/_custom/__init__.py +++ b/telethon/types/_custom/__init__.py @@ -1,4 +1,5 @@ from .adminlogevent import AdminLogEvent +from .auth import CodeType, SentCode from .draft import Draft from .dialog import Dialog from .inputsizedfile import InputSizedFile diff --git a/telethon/types/_custom/auth.py b/telethon/types/_custom/auth.py new file mode 100644 index 00000000..bfd1ad09 --- /dev/null +++ b/telethon/types/_custom/auth.py @@ -0,0 +1,120 @@ +import asyncio +import re +from enum import Enum, auto +from ... import _tl + + +class CodeType(Enum): + """ + The type of the login code sent. + + When resending the code, it won't be APP a second time. + """ + + APP = auto() + SMS = auto() + CALL = auto() + FLASH_CALL = auto() + MISSED_CALL = auto() + + +class SentCode: + """ + Information about the login code request, returned by `client.send_code_request`. + """ + + @classmethod + def _new(cls, code): + self = cls.__new__(cls) + self._code = code + self._start = asyncio.get_running_loop().time() + return self + + @property + def type(self): + """ + The `CodeType` which was sent. + """ + return { + _tl.auth.SentCodeTypeApp: CodeType.APP, + _tl.auth.SentCodeTypeSms: CodeType.SMS, + _tl.auth.SentCodeTypeCall: CodeType.CALL, + _tl.auth.SentCodeTypeFlashCall: CodeType.FLASH_CALL, + _tl.auth.SentCodeTypeMissedCall: CodeType.MISSED_CALL, + }[type(self._code.type)] + + @property + def next_type(self): + """ + The `CodeType` which will be sent if `client.send_code_request` + is used again after `timeout` seconds have elapsed. It may be `None`. + """ + if not self._code.next_type: + return None + + return { + _tl.auth.CodeTypeSms: CodeType.SMS, + _tl.auth.CodeTypeCall: CodeType.CALL, + _tl.auth.CodeTypeFlashCall: CodeType.FLASH_CALL, + _tl.auth.CodeTypeMissedCall: CodeType.MISSED_CALL, + }[type(self._code.next_type)] + + @property + def timeout(self): + """ + How many seconds are left before `client.send_code_request` can be used to resend the code. + Resending the code before this many seconds have elapsed may or may not work. + + This value can be `None`. + + This value is a positive floating point number, and is monotically decreasing. + The value will reach zero after enough seconds have elapsed. This lets you do some work + and call sleep on the value and still wait just long enough. + + If you need the original timeout, call `round` on the value as soon as possible. + """ + if not self._code.timeout: + return None + + return max(0.0, (self._start + self._code.timeout) - asyncio.get_running_loop().time()) + + @property + def length(self): + """ + The length of the sent code. + + If the length is unknown (it could be any length), `None` is returned. + This can be true for `CodeType.FLASH_CALL`. + """ + if isinstance(self._code.type, _tl.auth.SentCodeTypeFlashCall): + return None if self._code.type.pattern in ('', '*') else len(self._code.type.pattern) + else: + return self._code.type.length + + def check(self, code): + """ + Check if the user's input code is valid. + + This can be used to implement a client-side validation before actually trying to login + (mostly useful with a graphic interface, to hint the user the code is not yet correct). + """ + if not isinstance(code, str): + raise TypeError(f'code must be str, but was {type(code)}') + + if isinstance(self._code.type, _tl.auth.SentCodeTypeFlashCall): + if self._code.type.pattern in ('', '*'): + return True + + if not all(c.isdigit() or c == '*' for c in self._code.type.pattern): + # Potentially unsafe to use this pattern in a regex + raise RuntimeError(f'Unrecognised code pattern: {self._code.type.pattern!r}') + + pattern = self._code.type.pattern.replace('*', r'\d*') + numbers = ''.join(c for c in code if c.isdigit()) + return re.match(f'^{pattern}$', numbers) is not None + + if isinstance(self._code.type, _tl.auth.SentCodeTypeMissedCall): + if not code.startswith(self._code.type.prefix): + return False + + return len(code) == self._code.type.length