From 9e5843a0dcccccba04968476b79155056bc547f9 Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Fri, 23 Feb 2018 21:06:28 +0200 Subject: [PATCH] Refactor and clean up code --- mautrix_telegram/abstract_user.py | 23 ++- mautrix_telegram/bot.py | 10 +- mautrix_telegram/commands/handler.py | 5 +- mautrix_telegram/commands/telegram.py | 61 +++--- mautrix_telegram/formatter/from_matrix.py | 14 +- mautrix_telegram/formatter/from_telegram.py | 210 ++++++++++++-------- mautrix_telegram/matrix.py | 4 +- mautrix_telegram/portal.py | 111 ++++++----- mautrix_telegram/user.py | 5 +- 9 files changed, 257 insertions(+), 186 deletions(-) diff --git a/mautrix_telegram/abstract_user.py b/mautrix_telegram/abstract_user.py index abdf917a..ed642b2e 100644 --- a/mautrix_telegram/abstract_user.py +++ b/mautrix_telegram/abstract_user.py @@ -107,9 +107,8 @@ class AbstractUser: # region Telegram update handling async def _update(self, update): - if isinstance(update, - (UpdateShortChatMessage, UpdateShortMessage, UpdateNewChannelMessage, - UpdateNewMessage, UpdateEditMessage, UpdateEditChannelMessage)): + if isinstance(update, (UpdateShortChatMessage, UpdateShortMessage, UpdateNewChannelMessage, + UpdateNewMessage, UpdateEditMessage, UpdateEditChannelMessage)): await self.update_message(update) elif isinstance(update, UpdateDeleteMessages): await self.delete_message(update) @@ -122,13 +121,9 @@ class AbstractUser: elif isinstance(update, (UpdateChatAdmins, UpdateChatParticipantAdmin)): await self.update_admin(update) elif isinstance(update, UpdateChatParticipants): - portal = po.Portal.get_by_tgid(update.participants.chat_id) - if portal and portal.mxid: - await portal.update_telegram_participants(update.participants.participants) + await self.update_participants(update) elif isinstance(update, UpdateChannelPinnedMessage): - portal = po.Portal.get_by_tgid(update.channel_id) - if portal and portal.mxid: - await portal.update_telegram_pin(self, update.id) + await self.update_pinned_messages(update) elif isinstance(update, (UpdateUserName, UpdateUserPhoto)): await self.update_others_info(update) elif isinstance(update, UpdateReadHistoryOutbox): @@ -136,6 +131,16 @@ class AbstractUser: else: self.log.debug("Unhandled update: %s", update) + async def update_pinned_messages(self, update): + portal = po.Portal.get_by_tgid(update.channel_id) + if portal and portal.mxid: + await portal.update_telegram_pin(self, update.id) + + async def update_participants(self, update): + portal = po.Portal.get_by_tgid(update.participants.chat_id) + if portal and portal.mxid: + await portal.update_telegram_participants(update.participants.participants) + async def update_read_receipt(self, update): if not isinstance(update.peer, PeerUser): self.log.debug("Unexpected read receipt peer: %s", update.peer) diff --git a/mautrix_telegram/bot.py b/mautrix_telegram/bot.py index 020c12ee..ba43bffe 100644 --- a/mautrix_telegram/bot.py +++ b/mautrix_telegram/bot.py @@ -156,12 +156,10 @@ class Bot(AbstractUser): return action = update.message.action - if isinstance(action, MessageActionChatAddUser): - if self.tgid in action.users: - self.add_chat(to_id, type) - elif isinstance(action, MessageActionChatDeleteUser): - if action.user_id == self.tgid: - self.remove_chat(to_id) + if isinstance(action, MessageActionChatAddUser) and self.tgid in action.users: + self.add_chat(to_id, type) + elif isinstance(action, MessageActionChatDeleteUser) and action.user_id == self.tgid: + self.remove_chat(to_id) def is_in_chat(self, peer_id): return peer_id in self.chats diff --git a/mautrix_telegram/commands/handler.py b/mautrix_telegram/commands/handler.py index 43655129..e3350fbf 100644 --- a/mautrix_telegram/commands/handler.py +++ b/mautrix_telegram/commands/handler.py @@ -94,8 +94,9 @@ class CommandHandler: try: await command(evt) except FloodWaitError as e: - return evt.reply(f"Flood error: Please wait {format_duration(e.seconds)}") + return await evt.reply(f"Flood error: Please wait {format_duration(e.seconds)}") except Exception: self.log.exception("Fatal error handling command " f"{evt.command} {' '.join(args)} from {sender.mxid}") - return evt.reply("Fatal error while handling command. Check logs for more details.") + return await evt.reply("Fatal error while handling command. " + "Check logs for more details.") diff --git a/mautrix_telegram/commands/telegram.py b/mautrix_telegram/commands/telegram.py index 4774501e..072ed705 100644 --- a/mautrix_telegram/commands/telegram.py +++ b/mautrix_telegram/commands/telegram.py @@ -59,8 +59,8 @@ async def search(evt): return await evt.reply("\n".join(reply)) -@command_handler() -async def pm(evt): +@command_handler(name="pm") +async def private_message(evt): if len(evt.args) == 0: return await evt.reply("**Usage:** `$cmdprefix+sp pm `") @@ -159,16 +159,7 @@ async def join(evt): return await evt.reply(f"Created room for {portal.title}") -@command_handler() -async def create(evt): - type = evt.args[0] if len(evt.args) > 0 else "group" - if type not in {"chat", "group", "supergroup", "channel"}: - return await evt.reply( - "**Usage:** `$cmdprefix+sp create ['group'/'supergroup'/'channel']`") - - if po.Portal.get_by_mxid(evt.room_id): - return await evt.reply("This is already a portal room.") - +async def _get_initial_state(evt): state = await evt.az.intent.get_room_state(evt.room_id) title = None about = None @@ -180,20 +171,41 @@ async def create(evt): about = event["content"]["topic"] elif event["type"] == "m.room.power_levels": levels = event["content"] + return title, about, levels + + +def _check_power_levels(levels, bot_mxid): + try: + if levels["users"][bot_mxid] < 100: + raise ValueError() + except (TypeError, KeyError, ValueError): + return (f"Please give [the bridge bot](https://matrix.to/#/{bot_mxid}) a power level of " + "100 before creating a Telegram chat.") + + for user, level in levels["users"].items(): + if level >= 100 and user != bot_mxid: + return (f"Please make sure only the bridge bot has power level above 99 before " + f"creating a Telegram chat.\n\n" + f"Use power level 95 instead of 100 for admins.") + + +@command_handler() +async def create(evt): + type = evt.args[0] if len(evt.args) > 0 else "group" + if type not in {"chat", "group", "supergroup", "channel"}: + return await evt.reply( + "**Usage:** `$cmdprefix+sp create ['group'/'supergroup'/'channel']`") + + if po.Portal.get_by_mxid(evt.room_id): + return await evt.reply("This is already a portal room.") + + title, about, levels = await _get_initial_state(evt) if not title: return await evt.reply("Please set a title before creating a Telegram chat.") - elif (not levels or not levels["users"] or evt.az.intent.mxid not in levels["users"] or - levels["users"][evt.az.intent.mxid] < 100): - return await evt.reply("Please give " - f"[the bridge bot](https://matrix.to/#/{evt.az.intent.mxid})" - " a power level of 100 before creating a Telegram chat.") - else: - for user, level in levels["users"].items(): - if level >= 100 and user != evt.az.intent.mxid: - return await evt.reply( - f"Please make sure only the bridge bot has power level above" - f"99 before creating a Telegram chat.\n\n" - f"Use power level 95 instead of 100 for admins.") + + power_level_error = _check_power_levels(levels, evt.az.bot_mxid) + if power_level_error: + return await evt.reply(power_level_error) supergroup = type == "supergroup" type = { @@ -207,6 +219,7 @@ async def create(evt): try: await portal.create_telegram_chat(evt.sender, supergroup=supergroup) except ValueError as e: + portal.delete() return await evt.reply(e.args[0]) return await evt.reply(f"Telegram chat created. ID: {portal.tgid}") diff --git a/mautrix_telegram/formatter/from_matrix.py b/mautrix_telegram/formatter/from_matrix.py index 845fa609..c1e78f6d 100644 --- a/mautrix_telegram/formatter/from_matrix.py +++ b/mautrix_telegram/formatter/from_matrix.py @@ -53,30 +53,28 @@ class MatrixParser(HTMLParser): mention = self.mention_regex.match(url) if mention: mxid = mention.group(1) - user = pu.Puppet.get_by_mxid(mxid, create=False) + user = (pu.Puppet.get_by_mxid(mxid, create=False) + or u.User.get_by_mxid(mxid, create=False)) if not user: - user = u.User.get_by_mxid(mxid, create=False) - if not user: - return None, None + return None, None if user.username: entity_type = MessageEntityMention url = f"@{user.username}" else: entity_type = MessageEntityMentionName args["user_id"] = user.tgid - return url, entity_type + return entity_type, url room = self.room_regex.match(url) if room: username = po.Portal.get_username_from_mx_alias(room.group(1)) portal = po.Portal.find_by_username(username) if portal and portal.username: - return f"@{portal.username}", MessageEntityMention + return MessageEntityMention, f"@{portal.username}" if url.startswith("mailto:"): return MessageEntityEmail, url[len("mailto:"):] - - if self.get_starttag_text() == url: + elif self.get_starttag_text() == url: return MessageEntityUrl, url else: args["url"] = url diff --git a/mautrix_telegram/formatter/from_telegram.py b/mautrix_telegram/formatter/from_telegram.py index e504bb65..92dd00a9 100644 --- a/mautrix_telegram/formatter/from_telegram.py +++ b/mautrix_telegram/formatter/from_telegram.py @@ -43,6 +43,74 @@ def telegram_reply_to_matrix(evt, source): return {} +async def _add_forward_header(source, text, html, fwd_from_id): + if not html: + html = escape(text) + user = u.User.get_by_tgid(fwd_from_id) + if user: + fwd_from = f"{user.mxid}" + else: + puppet = pu.Puppet.get(fwd_from_id, create=False) + if puppet and puppet.displayname: + fwd_from = f"{puppet.displayname}" + else: + user = await source.client.get_entity(fwd_from_id) + if user: + fwd_from = f"{pu.Puppet.get_displayname(user, format=False)}" + else: + fwd_from = None + if not fwd_from: + fwd_from = "Unknown user" + text = f"Forwarded from {fwd_from}:\n{text}" + html = (f"Forwarded message from {fwd_from}
" + f"
{html}
") + return text, html + + +async def _add_reply_header(source, text, html, evt, relates_to, + native_replies, message_link_in_reply, main_intent, reply_text): + space = (evt.to_id.channel_id + if isinstance(evt, Message) and isinstance(evt.to_id, PeerChannel) + else source.tgid) + + msg = DBMessage.query.get((evt.reply_to_msg_id, space)) + if not msg: + return text, html + + if native_replies: + relates_to["m.in_reply_to"] = { + "event_id": msg.mxid, + "room_id": msg.mx_room, + } + if reply_text == "Edit": + html = f"Edit: {html or escape(text)}" + text = f"Edit: {text}" + return text, html + + reply_displayname = "unknown user" + try: + event = await main_intent.get_event(msg.mx_room, msg.mxid) + content = event["content"] + body = (content["formatted_body"] + if "formatted_body" in content + else content["body"]) + sender = event['sender'] + puppet = pu.Puppet.get_by_mxid(sender, create=False) + reply_displayname = puppet.displayname if puppet else sender + reply_to_user = f"{reply_displayname}" + reply_to_msg = (("{reply_text}") + if message_link_in_reply else "Reply") + quote = f"{reply_to_msg} to {reply_to_user}
{body}
" + except (ValueError, KeyError, MatrixRequestError): + quote = f"{reply_text} to unknown user (Failed to fetch message):
" + if not html: + html = escape(text) + html = quote + html + text = f"{reply_text} to {reply_displayname}:\n{text}" + return text, html + + async def telegram_to_matrix(evt, source, native_replies=False, message_link_in_reply=False, main_intent=None, reply_text="Reply"): text = add_surrogates(evt.message) @@ -50,61 +118,11 @@ async def telegram_to_matrix(evt, source, native_replies=False, message_link_in_ relates_to = {} if evt.fwd_from: - if not html: - html = escape(text) - from_id = evt.fwd_from.from_id - user = u.User.get_by_tgid(from_id) - if user: - fwd_from = f"{user.mxid}" - else: - puppet = pu.Puppet.get(from_id, create=False) - if puppet and puppet.displayname: - fwd_from = f"{puppet.displayname}" - else: - user = await source.client.get_entity(from_id) - if user: - fwd_from = pu.Puppet.get_displayname(user, format=False) - else: - fwd_from = None - if not fwd_from: - fwd_from = "Unknown user" - html = (f"Forwarded message from {fwd_from}
" - f"
{html}
") + text, html = await _add_forward_header(source, text, html, evt.fwd_from.from_id) if evt.reply_to_msg_id: - space = (evt.to_id.channel_id - if isinstance(evt, Message) and isinstance(evt.to_id, PeerChannel) - else source.tgid) - msg = DBMessage.query.get((evt.reply_to_msg_id, space)) - if msg: - if native_replies: - relates_to["m.in_reply_to"] = { - "event_id": msg.mxid, - "room_id": msg.mx_room, - } - if reply_text == "Edit": - html = "Edit: " + (html or escape(text)) - else: - try: - event = await main_intent.get_event(msg.mx_room, msg.mxid) - content = event["content"] - body = (content["formatted_body"] - if "formatted_body" in content - else content["body"]) - sender = event['sender'] - puppet = pu.Puppet.get_by_mxid(sender, create=False) - displayname = puppet.displayname if puppet else sender - reply_to_user = f"{displayname}" - reply_to_msg = (("{reply_text}") - if message_link_in_reply else "Reply") - quote = f"{reply_to_msg} to {reply_to_user}
{body}
" - except (ValueError, KeyError, MatrixRequestError): - quote = "{reply_text} to unknown user (Failed to fetch message):
" - if html: - html = quote + html - else: - html = quote + escape(text) + text, html = await _add_reply_header(source, text, html, evt, relates_to, native_replies, + message_link_in_reply, main_intent, reply_text) if isinstance(evt, Message) and evt.post and evt.post_author: if not html: @@ -150,44 +168,15 @@ def _telegram_entities_to_matrix(text, entities): elif entity_type == MessageEntityCode: html.append(f"{entity_text}") elif entity_type == MessageEntityPre: - if entity.language: - html.append("
"
-                            f"{entity_text}"
-                            "
") - else: - html.append(f"
{entity_text}
") + skip_entity = _parse_pre(html, entity_text, entity.language) elif entity_type == MessageEntityMention: - username = entity_text[1:] - - user = u.User.find_by_username(username) or pu.Puppet.find_by_username(username) - if user: - mxid = user.mxid - else: - portal = po.Portal.find_by_username(username) - mxid = portal.alias or portal.mxid if portal else None - - if mxid: - html.append(f"{entity_text}") - else: - skip_entity = True + skip_entity = _parse_mention(html, entity_text) elif entity_type == MessageEntityMentionName: - user = u.User.get_by_tgid(entity.user_id) - if user: - mxid = user.mxid - else: - puppet = pu.Puppet.get(entity.user_id, create=False) - mxid = puppet.mxid if puppet else None - if mxid: - html.append(f"{entity_text}") - else: - skip_entity = True + skip_entity = _parse_name_mention(html, entity_text, entity.user_id) elif entity_type == MessageEntityEmail: html.append(f"{entity_text}") elif entity_type in {MessageEntityTextUrl, MessageEntityUrl}: - url = escape(entity.url) if entity_type == MessageEntityTextUrl else entity_text - if not url.startswith(("https://", "http://", "ftp://", "magnet://")): - url = "http://" + url - html.append(f"{entity_text}") + skip_entity = _parse_url(html, entity_text, entity_type, entity.url) elif entity_type == MessageEntityBotCommand: html.append(f"!{entity_text[1:]}") elif entity_type == MessageEntityHashtag: @@ -198,3 +187,52 @@ def _telegram_entities_to_matrix(text, entities): html.append(text[last_offset:]) return "".join(html) + + +def _parse_pre(html, entity_text, language): + if language: + html.append("
"
+                    f"{entity_text}"
+                    "
") + else: + html.append(f"
{entity_text}
") + return False + + +def _parse_mention(html, entity_text): + username = entity_text[1:] + + user = u.User.find_by_username(username) or pu.Puppet.find_by_username(username) + if user: + mxid = user.mxid + else: + portal = po.Portal.find_by_username(username) + mxid = portal.alias or portal.mxid if portal else None + + if mxid: + html.append(f"{entity_text}") + else: + return True + return False + + +def _parse_name_mention(html, entity_text, user_id): + user = u.User.get_by_tgid(user_id) + if user: + mxid = user.mxid + else: + puppet = pu.Puppet.get(user_id, create=False) + mxid = puppet.mxid if puppet else None + if mxid: + html.append(f"{entity_text}") + else: + return True + return False + + +def _parse_url(html, entity_text, entity_type, url): + url = escape(url) if entity_type == MessageEntityTextUrl else entity_text + if not url.startswith(("https://", "http://", "ftp://", "magnet://")): + url = "http://" + url + html.append(f"{entity_text}") + return False diff --git a/mautrix_telegram/matrix.py b/mautrix_telegram/matrix.py index 88c19c2f..24ef73b5 100644 --- a/mautrix_telegram/matrix.py +++ b/mautrix_telegram/matrix.py @@ -56,11 +56,11 @@ class MatrixHandler: members = await self.az.intent.get_room_members(room) except MatrixRequestError: members = [] - if self.az.intent.mxid not in members: + if self.az.bot_mxid not in members: if len(members) > 1: await puppet.intent.error_and_leave(room, text=None, html=( f"Please invite " - f"the bridge bot " + f"the bridge bot " f"first if you want to create a Telegram chat.")) return diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index 0faf2d73..6997bf22 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -305,7 +305,7 @@ class Portal: joined_mxids = await self.main_intent.get_room_members(self.mxid) for user in joined_mxids: - if user == self.az.intent.mxid: + if user == self.az.bot_mxid: continue puppet_id = p.Puppet.get_id_from_mxid(user) if puppet_id and puppet_id not in allowed_tgids: @@ -530,49 +530,57 @@ class Portal: # We'll just assume the user is already in the chat. pass - async def handle_matrix_message(self, sender, message, event_id): - type = message["msgtype"] - if sender.logged_in: - client = sender.client - space = self.tgid if self.peer_type == "channel" else sender.tgid - else: - client = self.bot.client - space = self.tgid if self.peer_type == "channel" else self.bot.tgid - reply_to = formatter.matrix_reply_to_telegram(message, space, room_id=self.mxid) - - if type == "m.emote": + @staticmethod + def _preprocess_matrix_message(sender, message): + if message["msgtype"] == "m.emote": if "formatted_body" in message: message["formatted_body"] = f"* {sender.displayname} {message['formatted_body']}" message["body"] = f"* {sender.displayname} {message['body']}" - type = "m.text" + message["msgtype"] = "m.text" elif not sender.logged_in: if "formatted_body" in message: message["formatted_body"] = (f"<{sender.displayname}> " f"{message['formatted_body']}") message["body"] = f"<{sender.displayname}> {message['body']}" + return type + + def _handle_matrix_text(self, client, message, reply_to): + if "format" in message and message["format"] == "org.matrix.custom.html": + message, entities = formatter.matrix_to_telegram(message["formatted_body"]) + return client.send_message(self.peer, message, entities=entities, + reply_to=reply_to) + else: + return client.send_message(self.peer, message["body"], + reply_to=reply_to) + + async def _handle_matrix_file(self, client, message, reply_to): + file = await self.main_intent.download_file(message["url"]) + + info = message["info"] + mime = info["mimetype"] + + file_name, caption = self._get_file_meta(message["body"], mime) + + attributes = [DocumentAttributeFilename(file_name=file_name)] + if "w" in info and "h" in info: + attributes.append(DocumentAttributeImageSize(w=info["w"], h=info["h"])) + + return await client.send_file(self.peer, file, mime, caption, attributes, + file_name, reply_to=reply_to) + + async def handle_matrix_message(self, sender, message, event_id): + client = sender.client if sender.logged_in else self.bot.client + space = (self.tgid if self.peer_type == "channel" # Channels have their own ID space + else (sender.tgid if sender.logged_in else self.bot.tgid)) + reply_to = formatter.matrix_reply_to_telegram(message, space, room_id=self.mxid) + + self._preprocess_matrix_message(sender, message) + type = message["msgtype"] if type == "m.text" or (self.bridge_notices and type == "m.notice"): - if "format" in message and message["format"] == "org.matrix.custom.html": - message, entities = formatter.matrix_to_telegram(message["formatted_body"]) - response = await client.send_message(self.peer, message, entities=entities, - reply_to=reply_to) - else: - response = await client.send_message(self.peer, message["body"], - reply_to=reply_to) - elif type in {"m.image", "m.file", "m.audio", "m.video"}: - file = await self.main_intent.download_file(message["url"]) - - info = message["info"] - mime = info["mimetype"] - - file_name, caption = self._get_file_meta(message["body"], mime) - - attributes = [DocumentAttributeFilename(file_name=file_name)] - if "w" in info and "h" in info: - attributes.append(DocumentAttributeImageSize(w=info["w"], h=info["h"])) - - response = await client.send_file(self.peer, file, mime, caption, attributes, - file_name, reply_to=reply_to) + response = await self._handle_matrix_text(client, message, reply_to) + elif type in ("m.image", "m.file", "m.audio", "m.video"): + response = await self._handle_matrix_file(client, message, reply_to) else: self.log.debug("Unhandled Matrix event: %s", message) return @@ -604,6 +612,8 @@ class Portal: if not mx_user or not mx_user.tgid: continue user_id = mx_user.tgid + if not user_id or user_id == sender.tgid: + continue if user not in old_users or level != old_users[user]: if self.peer_type == "chat": await sender.client(EditChatAdminRequest( @@ -682,7 +692,7 @@ class Portal: user_tgids = set() user_mxids = await self.main_intent.get_room_members(self.mxid, ("join", "invite")) for user in user_mxids: - if user == self.az.intent.mxid: + if user == self.az.bot_mxid: continue mx_user = u.User.get_by_mxid(user, create=False) if mx_user and mx_user.tgid: @@ -945,18 +955,21 @@ class Portal: self.db.add(DBMessage(tgid=evt.id, mx_room=self.mxid, mxid=mxid, tg_space=tg_space)) self.db.commit() + async def _create_room_on_action(self, source, action): + create_and_exit = (MessageActionChatCreate, MessageActionChannelCreate) + create_and_continue = (MessageActionChatAddUser, MessageActionChatJoinedByLink) + if isinstance(action, create_and_exit + create_and_continue): + await self.create_matrix_room(source, invites=[source.mxid], + update_if_exists=isinstance(action, create_and_exit)) + if not isinstance(action, create_and_continue): + return False + return True + async def handle_telegram_action(self, source, sender, update): action = update.action - if not self.mxid: - create_and_exit = (MessageActionChatCreate, MessageActionChannelCreate) - create_and_continue = (MessageActionChatAddUser, MessageActionChatJoinedByLink) - if isinstance(action, create_and_exit + create_and_continue): - await self.create_matrix_room(source, invites=[source.mxid], - update_if_exists=isinstance(action, create_and_exit)) - if not isinstance(action, create_and_continue): - return - - if self.is_duplicate_action(update): + should_ignore = (not self.mxid and not await self._create_room_on_action(source, action) + or self.is_duplicate_action(update)) + if should_ignore: return # TODO figure out how to see changes to about text / channel username @@ -1091,11 +1104,15 @@ class Portal: def delete(self): try: del self.by_tgid[self.tgid_full] + except KeyError: + pass + try: del self.by_mxid[self.mxid] except KeyError: pass - self.db.delete(self.db_instance) - self.db.commit() + if self._db_instance: + self.db.delete(self._db_instance) + self.db.commit() @classmethod def from_db(cls, db_portal): diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py index f380ee4f..13ed5420 100644 --- a/mautrix_telegram/user.py +++ b/mautrix_telegram/user.py @@ -120,8 +120,9 @@ class User(AbstractUser): del self.by_tgid[self.tgid] except KeyError: pass - self.db.delete(self.db_instance) - self.db.commit() + if self._db_instance: + self.db.delete(self._db_instance) + self.db.commit() @classmethod def from_db(cls, db_user):