From 25b1adf62617cbf0ce75e43d075491ed32c0c937 Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Sat, 23 Jun 2018 00:44:41 +0300 Subject: [PATCH] Add support for logging in with a bot. Fixes #155 --- mautrix_telegram/abstract_user.py | 27 +++++++++------ mautrix_telegram/bot.py | 1 + mautrix_telegram/db.py | 3 +- mautrix_telegram/matrix.py | 8 ++--- mautrix_telegram/portal.py | 36 +++++++++++++------- mautrix_telegram/public/__init__.py | 20 +++++++++-- mautrix_telegram/public/login.html.mako | 44 ++++++++++++++++++++----- mautrix_telegram/puppet.py | 1 - mautrix_telegram/user.py | 21 ++++++++---- 9 files changed, 116 insertions(+), 45 deletions(-) diff --git a/mautrix_telegram/abstract_user.py b/mautrix_telegram/abstract_user.py index 3ab42892..f39a023a 100644 --- a/mautrix_telegram/abstract_user.py +++ b/mautrix_telegram/abstract_user.py @@ -36,12 +36,16 @@ class AbstractUser: az = None def __init__(self): - self.connected = False self.whitelisted = False self.client = None self.tgid = None self.mxid = None self.is_relaybot = False + self.is_bot = False + + @property + def connected(self): + return self.client and self.client.is_connected() def _init_client(self): self.log.debug(f"Initializing client for {self.name}") @@ -72,6 +76,8 @@ class AbstractUser: self.log.exception("Failed to handle Telegram update") async def _get_dialogs(self, limit=None): + if self.is_bot: + return dialogs = await self.client.get_dialogs(limit=limit) return [dialog.entity for dialog in dialogs if ( not isinstance(dialog.entity, (User, ChatForbidden, ChannelForbidden)) @@ -85,32 +91,33 @@ class AbstractUser: async def is_logged_in(self): return self.client and await self.client.is_user_authorized() - async def has_full_access(self): - return await self.is_logged_in() and self.whitelisted + async def has_full_access(self, allow_bot=False): + return self.whitelisted and (not self.is_bot or allow_bot) and await self.is_logged_in() async def start(self): if not self.client: self._init_client() - self.connected = await self.client.connect() + await self.client.connect() + self.log.debug("%s connected: %s", self.mxid, self.connected) + return self async def ensure_started(self, even_if_no_session=False): if not self.whitelisted: return self - self.log.info("CONNECTING USER %s, connected=%s, even_if_no_session=%s, session_count=%s", - self.mxid, self.connected, even_if_no_session, - self.session_container.Session.query.filter( - self.session_container.Session.session_id == self.mxid).count()) + self.log.debug("ensure_started(%s, connected=%s, even_if_no_session=%s, session_count=%s)", + self.mxid, self.connected, even_if_no_session, + self.session_container.Session.query.filter( + self.session_container.Session.session_id == self.mxid).count()) should_connect = (even_if_no_session or self.session_container.Session.query.filter( self.session_container.Session.session_id == self.mxid).count() > 0) if not self.connected and should_connect: - return await self.start() + await self.start() return self def stop(self): self.client.disconnect() self.client = None - self.connected = False # region Telegram update handling diff --git a/mautrix_telegram/bot.py b/mautrix_telegram/bot.py index 505d7562..94a677fc 100644 --- a/mautrix_telegram/bot.py +++ b/mautrix_telegram/bot.py @@ -42,6 +42,7 @@ class Bot(AbstractUser): self.whitelisted = True self.username = None self.is_relaybot = True + self.is_bot = True self.chats = {chat.id: chat.type for chat in BotChat.query.all()} self.tg_whitelist = [] self.whitelist_group_admins = config["bridge.relaybot.whitelist_group_admins"] or False diff --git a/mautrix_telegram/db.py b/mautrix_telegram/db.py index 1ce3b604..2ef9525e 100644 --- a/mautrix_telegram/db.py +++ b/mautrix_telegram/db.py @@ -104,8 +104,7 @@ class Puppet(Base): class BotChat(Base): query = None __tablename__ = "bot_chat" - bot_id = Column(Integer, primary_key=True, default=0) - chat_id = Column(Integer, primary_key=True) + id = Column(Integer, primary_key=True) type = Column(String, nullable=False) diff --git a/mautrix_telegram/matrix.py b/mautrix_telegram/matrix.py index d65dbde2..cb3c84c9 100644 --- a/mautrix_telegram/matrix.py +++ b/mautrix_telegram/matrix.py @@ -113,7 +113,7 @@ class MatrixHandler: return await user.ensure_started() portal = Portal.get_by_mxid(room) - if user and await user.has_full_access() and portal: + if user and await user.has_full_access(allow_bot=True) and portal: await portal.invite_telegram(inviter, user) return @@ -218,13 +218,13 @@ class MatrixHandler: async def handle_power_levels(self, room, sender, new, old): portal = Portal.get_by_mxid(room) sender = await User.get_by_mxid(sender).ensure_started() - if await sender.has_full_access() and portal: + if await sender.has_full_access(allow_bot=True) and portal: await portal.handle_matrix_power_levels(sender, new["users"], old["users"]) async def handle_room_meta(self, type, room, sender, content): portal = Portal.get_by_mxid(room) sender = await User.get_by_mxid(sender).ensure_started() - if await sender.has_full_access() and portal: + if await sender.has_full_access(allow_bot=True) and portal: handler, content_key = { "m.room.name": (portal.handle_matrix_title, "name"), "m.room.topic": (portal.handle_matrix_about, "topic"), @@ -237,7 +237,7 @@ class MatrixHandler: async def handle_room_pin(self, room, sender, new_events, old_events): portal = Portal.get_by_mxid(room) sender = await User.get_by_mxid(sender).ensure_started() - if await sender.has_full_access() and portal: + if await sender.has_full_access(allow_bot=True) and portal: events = new_events - old_events if len(events) > 0: # New event pinned, set that as pinned in Telegram. diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index cbd124a3..bee4b7db 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -24,7 +24,6 @@ import mimetypes import unicodedata import hashlib import logging -import re import magic from sqlalchemy.exc import IntegrityError, InvalidRequestError @@ -346,12 +345,21 @@ class Portal: return None return self.alias_template.format(groupname=username) + def add_bot_chat(self, entity): + if self.bot and entity.id == self.bot.tgid: + self.bot.add_chat(self.tgid, self.peer_type) + return + + user = u.User.get_by_tgid(entity.id) + if user and user.is_bot: + user.register_portal(self) + async def sync_telegram_users(self, source, users): allowed_tgids = set() for entity in users: puppet = p.Puppet.get(entity.id) - if self.bot and puppet.tgid == self.bot.tgid: - self.bot.add_chat(self.tgid, self.peer_type) + if entity.bot: + self.add_bot_chat(entity) allowed_tgids.add(entity.id) await puppet.intent.ensure_joined(self.mxid) await puppet.update_info(source, entity) @@ -380,6 +388,9 @@ class Portal: "User had left this Telegram chat.") continue mx_user = u.User.get_by_mxid(user, create=False) + if mx_user and mx_user.is_bot and mx_user.tgid not in allowed_tgids: + mx_user.unregister_portal(self) + if mx_user and not self.has_bot and mx_user.tgid not in allowed_tgids: await self.main_intent.kick(self.mxid, mx_user.mxid, "You had left this Telegram chat.") @@ -560,7 +571,8 @@ class Portal: if p.Puppet.get_id_from_mxid(member) or member == self.main_intent.mxid: continue user = await u.User.get_by_mxid(member).ensure_started() - if (has_bot and user.relaybot_whitelisted) or await user.has_full_access(): + if (has_bot and user.relaybot_whitelisted) or await user.has_full_access( + allow_bot=True): authenticated.append(user) return authenticated @@ -607,7 +619,7 @@ class Portal: return "" async def leave_matrix(self, user, source, event_id): - if not await user.is_logged_in(): + if await user.needs_relaybot(self): async with self.require_send_lock(self.bot.tgid): response = await self.bot.client.send_message( self.peer, f"__{user.displayname} left the room.__", markdown=True) @@ -639,7 +651,7 @@ class Portal: await user.client(LeaveChannelRequest(channel=channel)) async def join_matrix(self, user, event_id): - if not await user.is_logged_in(): + if await user.needs_relaybot(self): async with self.require_send_lock(self.bot.tgid): response = await self.bot.client.send_message( self.peer, f"__{user.displayname} joined the room.__", markdown=True) @@ -654,19 +666,19 @@ class Portal: pass @staticmethod - def _preprocess_matrix_message(sender, is_logged_in, message): + def _preprocess_matrix_message(sender, use_relaybot, message): msgtype = message["msgtype"] if msgtype == "m.emote": if "formatted_body" in message: tpl = config["bridge.message_formats.m_emote.html"] tpl_args = dict(sender_display_name=sender.displayname, - message=message['formatted_body']) + message=message['formatted_body']) message["formatted_body"] = Template(tpl).safe_substitute(tpl_args) tpl = config["bridge.message_formats.m_emote.plain"] tpl_args = dict(sender_display_name=sender.displayname, message=message['body']) message["body"] = Template(tpl).safe_substitute(tpl_args) message["msgtype"] = "m.text" - elif not is_logged_in: + elif not use_relaybot: html = message["formatted_body"] if "formatted_body" in message else None text = message["body"] if msgtype == "m.text": @@ -675,7 +687,7 @@ class Portal: if not html: html = escape_html(text) tpl = config["bridge.message_formats.m_text.html"] - tpl_args = dict(sender_display_name=sender.displayname,message=html) + tpl_args = dict(sender_display_name=sender.displayname, message=html) html = Template(tpl).safe_substitute(tpl_args) tpl = config["bridge.message_formats.m_text.plain"] tpl_args = dict(sender_display_name=sender.displayname, message=text) @@ -813,7 +825,7 @@ class Portal: self.db.commit() async def handle_matrix_message(self, sender, message, event_id): - logged_in = await sender.is_logged_in() + logged_in = not await sender.needs_relaybot(self) client = sender.client if logged_in else self.bot.client sender_id = sender.tgid if logged_in else self.bot.tgid space = (self.tgid if self.peer_type == "channel" # Channels have their own ID space @@ -850,7 +862,7 @@ class Portal: pass async def handle_matrix_deletion(self, deleter, event_id): - deleter = deleter if await deleter.is_logged_in() else self.bot + deleter = deleter if await deleter.needs_relaybot(self) else self.bot space = self.tgid if self.peer_type == "channel" else deleter.tgid message = DBMessage.query.filter(DBMessage.mxid == event_id, DBMessage.tg_space == space, diff --git a/mautrix_telegram/public/__init__.py b/mautrix_telegram/public/__init__.py index e1a4533c..c9f8362d 100644 --- a/mautrix_telegram/public/__init__.py +++ b/mautrix_telegram/public/__init__.py @@ -45,15 +45,16 @@ class PublicBridgeWebsite: async def get_login(self, request): user = (User.get_by_mxid(request.rel_url.query["mxid"], create=False) if "mxid" in request.rel_url.query else None) + state = "token" if request.rel_url.query.get("mode", "") == "bot" else "request" if not user: return self.render_login( mxid=request.rel_url.query["mxid"] if "mxid" in request.rel_url.query else None, - state="request") + state=state) elif not user.whitelisted: return self.render_login(mxid=user.mxid, error="You are not whitelisted.", status=403) await user.ensure_started() if not await user.is_logged_in(): - return self.render_login(mxid=user.mxid, state="request") + return self.render_login(mxid=user.mxid, state=state) return self.render_login(mxid=user.mxid, username=user.username) @@ -62,6 +63,19 @@ class PublicBridgeWebsite: text=self.login.render(username=username, state=state, error=error, message=message, mxid=mxid)) + async def post_login_token(self, user, token): + try: + user_info = await user.client.sign_in(bot_token=token) + asyncio.ensure_future(user.post_login(user_info), loop=self.loop) + if user.command_status and user.command_status["action"] == "Login": + user.command_status = None + return self.render_login(mxid=user.mxid, state="logged-in", status=200, + username=user_info.username) + except Exception: + self.log.exception("Error sending bot token") + return self.render_login(mxid=user.mxid, state="token", status=500, + error="Internal server error while sending token.") + async def post_login_phone(self, user, phone): try: await user.client.sign_in(phone or "+123") @@ -153,6 +167,8 @@ class PublicBridgeWebsite: if "phone" in data: return await self.post_login_phone(user, data["phone"]) + elif "token" in data: + return await self.post_login_token(user, data["token"]) elif "code" in data: resp = await self.post_login_code(user, data["code"], password_in_data="password" in data) diff --git a/mautrix_telegram/public/login.html.mako b/mautrix_telegram/public/login.html.mako index a63e30e3..8c03cbdc 100644 --- a/mautrix_telegram/public/login.html.mako +++ b/mautrix_telegram/public/login.html.mako @@ -29,6 +29,25 @@ along with this program. If not, see . + +
@@ -40,6 +59,13 @@ along with this program. If not, see . You can now close this page. You should be invited to Telegram portals on Matrix momentarily.

+ % elif state == "bot-logged-in": +

Logged in successfully!

+

+ Logged in as @${username}. + You can now close this page. + You should be invited to Telegram portals on Matrix momentarily. +

% else:

You're already logged in!

@@ -67,24 +93,26 @@ along with this program. If not, see . + + % elif state == "token": + + + % elif state == "code": -

- -
% elif state == "password": + % endif + % if state != "request":
-
diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index 0cd198fd..ad0648bc 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -19,7 +19,6 @@ import re import logging from telethon.tl.types import UserProfilePhoto -from telethon.errors import LocationInvalidError from .db import Puppet as DBPuppet from . import util diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py index 8889d445..809b2ea8 100644 --- a/mautrix_telegram/user.py +++ b/mautrix_telegram/user.py @@ -36,10 +36,11 @@ class User(AbstractUser): by_tgid = {} def __init__(self, mxid, tgid=None, username=None, db_contacts=None, saved_contacts=0, - db_portals=None, db_instance=None): + is_bot=False, db_portals=None, db_instance=None): super().__init__() self.mxid = mxid self.tgid = tgid + self.is_bot = is_bot self.username = username self.contacts = [] self.saved_contacts = saved_contacts @@ -127,7 +128,7 @@ class User(AbstractUser): @classmethod def from_db(cls, db_user): return User(db_user.mxid, db_user.tgid, db_user.tg_username, db_user.contacts, - db_user.saved_contacts, db_user.portals, db_instance=db_user) + False, db_user.saved_contacts, db_user.portals, db_instance=db_user) # endregion # region Telegram connection management @@ -148,19 +149,23 @@ class User(AbstractUser): async def post_login(self, info=None): try: await self.update_info(info) - await self.sync_dialogs() - await self.sync_contacts() + if not self.is_bot: + await self.sync_dialogs() + await self.sync_contacts() if config["bridge.catch_up"]: await self.client.catch_up() except Exception: - self.log.exception("Failed to run post-login functions") + self.log.exception("Failed to run post-login functions for %s", self.mxid) # endregion # region Telegram actions that need custom methods - async def update_info(self, info=None): + async def update_info(self, info: User = None): info = info or await self.client.get_me() changed = False + if self.is_bot != info.bot: + self.is_bot = info.bot + changed = True if self.username != info.username: self.username = info.username changed = True @@ -252,6 +257,10 @@ class User(AbstractUser): except KeyError: pass + async def needs_relaybot(self, portal): + return not await self.is_logged_in() or ( + self.is_bot and portal.tgid_full not in self.portals) + def _hash_contacts(self): acc = 0 for id in sorted([self.saved_contacts] + [contact.id for contact in self.contacts]):