Xtao-Labs/MTPyroger
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Clean the code and add some more functionality:

Browse Source 
- The phone_code can also be passed in advance (for test numbers).
- Pass first_name and last_name for automatic account creation.
- Handle PhoneCodeInvalid errors.
This commit is contained in:
Dan 2018-01-24 17:53:30 +01:00
parent d01d852dc2
commit f65d369fd1
1 changed files with 53 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

86
pyrogram/client/client.py
View File

@ -75,10 +75,23 @@ class Client:
INVITE_LINK_RE = re.compile(r"^(?:https?://)?t\.me/joinchat/(.+)$") INVITE_LINK_RE = re.compile(r"^(?:https?://)?t\.me/joinchat/(.+)$")
DIALOGS_AT_ONCE = 100 DIALOGS_AT_ONCE = 100
def __init__(self, session_name: str, test_mode: bool = False): def __init__(self,
session_name: str,
test_mode: bool = False,
phone_number: str = None,
phone_code: str or callable = None,
password: str = None,
first_name: str = None,
last_name: str = None):
self.session_name = session_name self.session_name = session_name
self.test_mode = test_mode self.test_mode = test_mode
self.phone_number = phone_number
self.password = password
self.phone_code = phone_code
self.first_name = first_name
self.last_name = last_name
self.dc_id = None self.dc_id = None
self.auth_key = None self.auth_key = None
self.user_id = None self.user_id = None
@ -170,24 +183,27 @@ class Client:
""" """
return self.session.send(data) return self.session.send(data)
def authorize(self, phone_number=None, code_callback=None, password=None): def authorize(self):
invalid_phone_raises = phone_number is not None phone_number_invalid_raises = self.phone_number is not None
phone_code_invalid_raises = self.phone_code is not None
password_hash_invalid_raises = self.password is not None
while True: while True:
if phone_number is None: if self.phone_number is None:
phone_number = input("Enter phone number: ") self.phone_number = input("Enter phone number: ")
while True: while True:
confirm = input("Is \"{}\" correct? (y/n): ".format(phone_number)) confirm = input("Is \"{}\" correct? (y/n): ".format(self.phone_number))
if confirm in ("y", "1"): if confirm in ("y", "1"):
break break
elif confirm in ("n", "2"): elif confirm in ("n", "2"):
phone_number = input("Enter phone number: ") self.phone_number = input("Enter phone number: ")
try: try:
r = self.send( r = self.send(
functions.auth.SendCode( functions.auth.SendCode(
phone_number, self.phone_number,
self.config.api_id, self.config.api_id,
self.config.api_hash self.config.api_hash
) )
@ -203,18 +219,18 @@ class Client:
r = self.send( r = self.send(
functions.auth.SendCode( functions.auth.SendCode(
phone_number, self.phone_number,
self.config.api_id, self.config.api_id,
self.config.api_hash self.config.api_hash
) )
) )
break break
except PhoneNumberInvalid as e: except PhoneNumberInvalid as e:
if invalid_phone_raises: if phone_number_invalid_raises:
raise raise
else: else:
print(e.MESSAGE) print(e.MESSAGE)
phone_number = None self.phone_number = None
except FloodWait as e: except FloodWait as e:
print(e.MESSAGE.format(x=e.x)) print(e.MESSAGE.format(x=e.x))
time.sleep(e.x) time.sleep(e.x)
@ -225,70 +241,74 @@ class Client:
phone_registered = r.phone_registered phone_registered = r.phone_registered
phone_code_hash = r.phone_code_hash phone_code_hash = r.phone_code_hash
if not code_callback:
def code_callback():
return input("Enter phone code: ")
while True: while True:
phone_code = code_callback() self.phone_code = (
input("Enter phone code: ") if self.phone_code is None
else self.phone_code if type(self.phone_code) is str
else self.phone_code()
)
try: try:
if phone_registered: if phone_registered:
r = self.send( r = self.send(
functions.auth.SignIn( functions.auth.SignIn(
phone_number, self.phone_number,
phone_code_hash, phone_code_hash,
phone_code self.phone_code
) )
) )
else: else:
try: try:
self.send( self.send(
functions.auth.SignIn( functions.auth.SignIn(
phone_number, self.phone_number,
phone_code_hash, phone_code_hash,
phone_code self.phone_code
) )
) )
except PhoneNumberUnoccupied: except PhoneNumberUnoccupied:
pass pass
first_name = input("First name: ") self.first_name = self.first_name or input("First name: ")
last_name = input("Last name: ") self.last_name = self.last_name or input("Last name: ")
r = self.send( r = self.send(
functions.auth.SignUp( functions.auth.SignUp(
phone_number, self.phone_number,
phone_code_hash, phone_code_hash,
phone_code, self.phone_code,
first_name, self.first_name,
last_name self.last_name
) )
) )
except (PhoneCodeInvalid, PhoneCodeEmpty, PhoneCodeExpired, PhoneCodeHashEmpty) as e: except (PhoneCodeInvalid, PhoneCodeEmpty, PhoneCodeExpired, PhoneCodeHashEmpty) as e:
print(e.MESSAGE) if phone_code_invalid_raises:
raise
else:
print(e.MESSAGE)
self.phone_code = None
except SessionPasswordNeeded as e: except SessionPasswordNeeded as e:
print(e.MESSAGE) print(e.MESSAGE)
invalid_password_raises = password is not None
while True: while True:
try: try:
r = self.send(functions.account.GetPassword()) r = self.send(functions.account.GetPassword())
if password is None: if self.password is None:
print("Hint: {}".format(r.hint)) print("Hint: {}".format(r.hint))
password = input("Enter password: ") # TODO: Use getpass self.password = input("Enter password: ") # TODO: Use getpass
password = r.current_salt + password.encode() + r.current_salt self.password = r.current_salt + self.password.encode() + r.current_salt
password_hash = sha256(password).digest() password_hash = sha256(self.password).digest()
r = self.send(functions.auth.CheckPassword(password_hash)) r = self.send(functions.auth.CheckPassword(password_hash))
except PasswordHashInvalid as e: except PasswordHashInvalid as e:
if invalid_password_raises: if password_hash_invalid_raises:
raise raise
else: else:
print(e.MESSAGE) print(e.MESSAGE)
password = None self.password = None
except FloodWait as e: except FloodWait as e:
print(e.MESSAGE.format(x=e.x)) print(e.MESSAGE.format(x=e.x))
time.sleep(e.x) time.sleep(e.x)