Added two-step verification (fixes #4) and more info for errors

telethon/errors.py

@@ -126,6 +126,8 @@ class RPCError(Exception):
         'CHAT_ADMIN_REQUIRED': 'Chat admin privileges are required to do that in the specified chat '
                                '(for example, to send a message in a channel which is not yours).',
 
+        'PASSWORD_HASH_INVALID': 'The password (and thus its hash value) you entered is invalid.',
+
         # 401 UNAUTHORIZED
         'AUTH_KEY_UNREGISTERED': 'The key is not registered in the system.',
 
@@ -141,6 +143,8 @@ class RPCError(Exception):
 
         'AUTH_KEY_PERM_EMPTY': 'The method is unavailable for temporary authorization key, not bound to permanent.',
 
+        'SESSION_PASSWORD_NEEDED': 'Two-steps verification is enabled and a password is required.',
+
         # 420 FLOOD
         'FLOOD_WAIT_(\d+)': 'A wait of {} seconds is required.'
     }
@@ -164,6 +168,10 @@ class RPCError(Exception):
                     self.additional_data = None
                     super().__init__(self, error_msg)
 
+                    # Add another field to easily determine whether this error
+                    # should be handled as a password-required error
+                    self.password_required = message == 'SESSION_PASSWORD_NEEDED'
+
                 called_super = True
                 break
 

telethon/helpers.py

@@ -58,4 +58,24 @@ def sha1(data):
     sha.update(data)
     return sha.digest()
 
+
+def sha256(data):
+    """Calculates the SHA256 digest for the given data"""
+    sha = hashlib.sha256()
+    sha.update(data)
+    return sha.digest()
+
+
+def get_password_hash(pw, current_salt):
+    """Gets the password hash for the two-step verification.
+       curent_salt should be the byte array provided by invoking GetPasswordRequest()"""
+
+    # Passwords are encoded as UTF-8
+    # https://github.com/DrKLO/Telegram/blob/e31388/TMessagesProj/src/main/java/org/telegram/ui/LoginActivity.java#L2003
+    data = pw.encode('utf-8')
+
+    pw_hash = current_salt+data+current_salt
+    return sha256(pw_hash)
+
+
 # endregion

telethon/interactive_telegram_client.py

@@ -1,10 +1,11 @@
 from telethon.tl.types import UpdateShortChatMessage
 from telethon.tl.types import UpdateShortMessage
-from telethon import TelegramClient
+from telethon import TelegramClient, RPCError
 
 from telethon.utils import get_display_name, get_input_peer
 
 import shutil
+from getpass import getpass
 
 # Get the (current) number of lines in the terminal
 cols, rows = shutil.get_terminal_size()
@@ -51,7 +52,16 @@ class InteractiveTelegramClient(TelegramClient):
             code_ok = False
             while not code_ok:
                 code = input('Enter the code you just received: ')
-                code_ok = self.sign_in(user_phone, code)
+                try:
+                    code_ok = self.sign_in(user_phone, code)
+
+                # Two-step verification may be enabled
+                except RPCError as e:
+                    if e.password_required:
+                        pw = getpass('Two step verification is enabled. Please enter your password: ')
+                        code_ok = self.sign_in(password=pw)
+                    else:
+                        raise e
 
     def run(self):
         # Listen for updates

telethon/telegram_client.py

@@ -12,13 +12,19 @@ from telethon.tl import Session
 from telethon.tl.functions.upload import SaveBigFilePartRequest
 from telethon.tl.functions import InvokeWithLayerRequest, InitConnectionRequest
 from telethon.tl.functions.help import GetConfigRequest
-from telethon.tl.functions.auth import SendCodeRequest, SignInRequest, SignUpRequest, LogOutRequest
 from telethon.tl.functions.upload import SaveFilePartRequest, GetFileRequest
 from telethon.tl.functions.messages import \
     GetDialogsRequest, GetHistoryRequest, \
     SendMessageRequest, SendMediaRequest, \
     ReadHistoryRequest
 
+from telethon.tl.functions.auth import \
+    SendCodeRequest, CheckPasswordRequest, \
+    SignInRequest, SignUpRequest, LogOutRequest
+
+# The following is required to get the password salt
+from telethon.tl.functions.account import GetPasswordRequest
+
 # All the types we need to work with
 from telethon.tl.types import \
     InputPeerEmpty, \
@@ -41,7 +47,7 @@ from telethon.tl.all_tlobjects import layer
 class TelegramClient:
 
     # Current TelegramClient version
-    __version__ = '0.6'
+    __version__ = '0.7'
 
     # region Initialization
 
@@ -161,21 +167,32 @@ class TelegramClient:
             except InvalidDCError as error:
                 self.reconnect_to_dc(error.new_dc)
 
-    def sign_in(self, phone_number, code):
-        """Completes the authorization of a phone number by providing the received code"""
-        if phone_number not in self.phone_code_hashes:
-            raise ValueError('Please make sure you have called send_code_request first.')
+    def sign_in(self, phone_number=None, code=None, password=None):
+        """Completes the authorization of a phone number by providing the received code.
 
-        try:
-            result = self.invoke(SignInRequest(
-                phone_number, self.phone_code_hashes[phone_number], code))
+           If no phone or code is provided, then the sole password will be used. The password
+           should be used after a normal authorization attempt has happened and an RPCError
+           with `.password_required = True` was raised"""
+        if phone_number and code:
+            if phone_number not in self.phone_code_hashes:
+                raise ValueError('Please make sure you have called send_code_request first.')
 
-        except RPCError as error:
-            if error.message.startswith('PHONE_CODE_'):
-                print(error)
-                return False
-            else:
-                raise error
+            try:
+                result = self.invoke(SignInRequest(
+                    phone_number, self.phone_code_hashes[phone_number], code))
+
+            except RPCError as error:
+                if error.message.startswith('PHONE_CODE_'):
+                    print(error)
+                    return False
+                else:
+                    raise error
+        elif password:
+            salt = self.invoke(GetPasswordRequest()).current_salt
+            result = self.invoke(CheckPasswordRequest(utils.get_password_hash(password, salt)))
+        else:
+            raise ValueError('You must provide a phone_number and a code for the first time, '
+                             'and a password only if an RPCError was raised before.')
 
         # Result is an Auth.Authorization TLObject
         self.session.user = result.user