Add edit_2FA function

New convenience function to allow easy editing of 2FA settings

telethon/telegram_client.py

@@ -38,12 +38,12 @@ from .errors import (
     RPCError, UnauthorizedError, PhoneCodeEmptyError, PhoneCodeExpiredError,
     PhoneCodeHashEmptyError, PhoneCodeInvalidError, LocationInvalidError,
     SessionPasswordNeededError, FileMigrateError, PhoneNumberUnoccupiedError,
-    PhoneNumberOccupiedError
+    PhoneNumberOccupiedError, EmailUnconfirmedError
 )
 from .network import ConnectionMode
 from .tl.custom import Draft, Dialog
 from .tl.functions.account import (
-    GetPasswordRequest
+    GetPasswordRequest, UpdatePasswordSettingsRequest
 )
 from .tl.functions.auth import (
     CheckPasswordRequest, LogOutRequest, SendCodeRequest, SignInRequest,
@@ -86,6 +86,7 @@ from .tl.types import (
     PhotoSizeEmpty, MessageService
 )
 from .tl.types.messages import DialogsSlice
+from .tl.types.account import PasswordInputSettings, NoPassword
 from .extensions import markdown, html
 
 __log__ = logging.getLogger(__name__)
@@ -2470,4 +2471,75 @@ class TelegramClient(TelegramBareClient):
             'Make sure you have encountered this peer before.'.format(peer)
         )
 
+    def edit_2FA(self, current_password=None, new_password=None, hint='', email=None):
+        """
+        Changes the 2FA settings of the logged in user, according to the
+        passed parameters. Take note of the parameter explanations
+
+        current_password (:obj:`str`, optional):
+            The current password, to authorize changing to ``new_password``
+            Must be set if changing existing 2FA settings.
+            Must not be set if 2FA is currently disabled.
+            Passing this by itself will remove 2FA (if correct)
+
+        new_password (:obj:`str`, optional):
+            The password to set as 2FA. If 2FA was already enabled, ``current_password``
+            must be set. Leaving this blank or ``None`` will remove the password
+
+        hint (:obj:`str`, optional):
+            Hint to be displayed by Telegram when it asks for 2FA.
+            Leaving unspecified is highly discouraged.
+            Has no effect if ``new_password`` is not set
+
+        email (:obj:`str`, optional):
+            Recovery and verification email.
+            Throws ``EmailUnconfirmedError`` if value differs from current one
+            Has no effect if ``new_password`` is not set
+
+        Returns:
+            ``True`` if successful, ``False`` otherwise
+        """
+        if new_password is None and current_password is None:
+            __log__.warn('Both new_password and current_password were not '\
+                          'specified, nothing to do')
+            return False
+        
+        pass_result = self(GetPasswordRequest())
+        if isinstance(pass_result, NoPassword) and current_password:
+            __log__.warn('You supplied current_password even though this '\
+                          'account does not have one set')
+            return False  # can do "current_password = None" and continue too
+        salt_random = os.urandom(8)
+        salt = pass_result.new_salt + salt_random
+
+        if current_password:
+            current_password = pass_result.current_salt +\
+                current_password.encode() + pass_result.current_salt
+            current_password_hash = hashlib.sha256(current_password).digest()
+        else:
+            current_password_hash = salt
+
+        if new_password:  # setting new password
+            new_password = salt + new_password.encode('utf-8') + salt
+            new_password_hash = hashlib.sha256(new_password).digest()
+            new_settings = PasswordInputSettings(
+                new_salt=salt,
+                new_password_hash=new_password_hash,
+                hint=hint
+            )
+            if email:  # if enabling 2FA or changing email
+                new_settings.email = email  # TG counts empty string as None
+            return self(UpdatePasswordSettingsRequest(
+                current_password_hash, new_settings=new_settings
+            ))
+        else:  # removing password
+            return self(UpdatePasswordSettingsRequest(
+                current_password_hash,
+                new_settings=PasswordInputSettings(
+                    new_salt=bytes(),
+                    new_password_hash=bytes(),
+                    hint=hint
+                )
+            ))
+
     # endregion