Add option to log in to custom puppet with shared secret

mautrix_telegram/commands/telegram/auth.py

@@ -20,7 +20,8 @@ from telethon.errors import (  # isort: skip
     AccessTokenExpiredError, AccessTokenInvalidError, FirstNameInvalidError, FloodWaitError,
     PasswordHashInvalidError, PhoneCodeExpiredError, PhoneCodeInvalidError,
     PhoneNumberAppSignupForbiddenError, PhoneNumberBannedError, PhoneNumberFloodError,
-    PhoneNumberOccupiedError, PhoneNumberUnoccupiedError, SessionPasswordNeededError)
+    PhoneNumberOccupiedError, PhoneNumberUnoccupiedError, SessionPasswordNeededError,
+    PhoneNumberInvalidError)
 
 from mautrix.types import EventID
 
@@ -84,7 +85,7 @@ async def enter_code_register(evt: CommandEvent) -> EventID:
         await evt.sender.ensure_started(even_if_no_session=True)
         first_name, last_name = evt.sender.command_status["full_name"]
         user = await evt.sender.client.sign_up(evt.args[0], first_name, last_name)
-        asyncio.ensure_future(evt.sender.post_login(user), loop=evt.loop)
+        asyncio.ensure_future(evt.sender.post_login(user, first_login=True), loop=evt.loop)
         evt.sender.command_status = None
         return await evt.reply(f"Successfully registered to Telegram.")
     except PhoneNumberOccupiedError:
@@ -166,6 +167,8 @@ async def _request_code(evt: CommandEvent, phone_number: str, next_status: Dict[
     except PhoneNumberUnoccupiedError:
         return await  evt.reply("That phone number has not been registered. "
                                 "Please register with `$cmdprefix+sp register <phone>`.")
+    except PhoneNumberInvalidError:
+        return await evt.reply("That phone number is not valid.")
     except Exception:
         evt.log.exception("Error requesting phone code")
         return await evt.reply("Unhandled exception while requesting code. "
@@ -244,7 +247,7 @@ async def _sign_in(evt: CommandEvent, **sign_in_info) -> EventID:
             await evt.reply(f"[{existing_user.displayname}]"
                             f"(https://matrix.to/#/{existing_user.mxid})"
                             " was logged out from the account.")
-        asyncio.ensure_future(evt.sender.post_login(user), loop=evt.loop)
+        asyncio.ensure_future(evt.sender.post_login(user, first_login=True), loop=evt.loop)
         evt.sender.command_status = None
         name = f"@{user.username}" if user.username else f"+{user.phone}"
         return await evt.reply(f"Successfully logged in as {name}")

mautrix_telegram/config.py

@@ -109,6 +109,7 @@ class Config(BaseBridgeConfig):
         copy("bridge.plaintext_highlights")
         copy("bridge.public_portals")
         copy("bridge.sync_with_custom_puppets")
+        copy("bridge.login_shared_secret")
         copy("bridge.telegram_link_preview")
         copy("bridge.inline_images")
         copy("bridge.image_as_file_size")

mautrix_telegram/puppet.py

@@ -424,4 +424,7 @@ def init(context: 'Context') -> Iterable[Awaitable[Any]]:
     Puppet.displayname_template = SimpleTemplate(config["bridge.displayname_template"],
                                                  "displayname")
 
+    Puppet.login_shared_secret = config["bridge.login_shared_secret"].encode("utf-8")
+    Puppet.login_device_name = "Telegram Bridge"
+
     return (puppet.try_start() for puppet in Puppet.all_with_custom_mxid())

mautrix_telegram/user.py

@@ -199,14 +199,27 @@ class User(AbstractUser, BaseUser):
             self.client.session.delete()
         return self
 
-    async def post_login(self, info: TLUser = None) -> None:
+    async def post_login(self, info: TLUser = None, first_login: bool = False) -> None:
         try:
             await self.update_info(info)
-            if not self.is_bot and config["bridge.startup_sync"]:
+        except Exception:
+            self.log.exception("Failed to update telegram account info")
+            return
+
+        try:
+            puppet = pu.Puppet.get(self.tgid)
+            if puppet.custom_mxid != self.mxid and puppet.can_auto_login(self.mxid):
+                self.log.info(f"Automatically enabling custom puppet")
+                await puppet.switch_mxid(access_token="auto", mxid=self.mxid)
+        except Exception:
+            self.log.exception("Failed to automatically enable custom puppet")
+
+        if not self.is_bot and config["bridge.startup_sync"]:
+            try:
                 await self.sync_dialogs()
                 await self.sync_contacts()
-        except Exception:
-            self.log.exception("Failed to run post-login functions for %s", self.mxid)
+            except Exception:
+                self.log.exception("Failed to run post-login sync")
 
     async def update(self, update: TypeUpdate) -> bool:
         if not self.is_bot:

mautrix_telegram/web/common/auth_api.py

@@ -119,7 +119,7 @@ class AuthAPI(abc.ABC):
         existing_user = User.get_by_tgid(user_info.id)
         if existing_user and existing_user != user:
             await existing_user.log_out()
-        asyncio.ensure_future(user.post_login(user_info), loop=self.loop)
+        asyncio.ensure_future(user.post_login(user_info, first_login=True), loop=self.loop)
         if user.command_status and user.command_status["action"] == "Login":
             user.command_status = None