From 8a3ccb6e8c517943deeb37f9c25ac0d5b839e4c0 Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Mon, 22 Jan 2018 21:20:56 +0200 Subject: [PATCH] Implement message reply/forward bridging in both directions --- README.md | 6 +-- mautrix_telegram/commands.py | 1 - mautrix_telegram/db.py | 15 +++++-- mautrix_telegram/formatter.py | 76 ++++++++++++++++++++++++++++++++--- mautrix_telegram/matrix.py | 6 +-- mautrix_telegram/portal.py | 33 +++++++++------ mautrix_telegram/puppet.py | 4 +- mautrix_telegram/user.py | 2 +- 8 files changed, 113 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 3108ff18..67097177 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,7 @@ does not do this automatically. * [x] Formatted messages * [ ] Bot commands (!command -> /command) * [x] Mentions - * [ ] Rich quotes + * [x] Rich quotes * [ ] Locations (not implemented in Riot) * [ ] Images * [ ] Files @@ -77,8 +77,8 @@ does not do this automatically. * [x] Formatted messages * [x] Bot commands (/command -> !command) * [x] Mentions - * [ ] Replies - * [ ] Forwards + * [x] Replies + * [x] Forwards * [ ] Images * [ ] Locations * [ ] Stickers diff --git a/mautrix_telegram/commands.py b/mautrix_telegram/commands.py index 9562d00a..819b79b7 100644 --- a/mautrix_telegram/commands.py +++ b/mautrix_telegram/commands.py @@ -43,7 +43,6 @@ class CommandHandler: try: command = command_handlers[command] except KeyError: - print(sender.command_status) if sender.command_status and "next" in sender.command_status: args.insert(0, command) command = sender.command_status["next"] diff --git a/mautrix_telegram/db.py b/mautrix_telegram/db.py index a4dda1f2..7510b85f 100644 --- a/mautrix_telegram/db.py +++ b/mautrix_telegram/db.py @@ -13,9 +13,7 @@ # # You should have received a copy of the GNU General Public License # along with this program. If not, see . -from sqlalchemy import orm, \ - Column, ForeignKey, \ - Integer, String +from sqlalchemy import Column, ForeignKey, UniqueConstraint, Integer, String from sqlalchemy.orm.scoping import scoped_session from .base import Base @@ -36,6 +34,16 @@ class Portal(Base): photo_id = Column(String, nullable=True) +class Message(Base): + __tablename__ = "message" + mxid = Column(String) + mx_room = Column(String) + tgid = Column(Integer, primary_key=True) + user = Column(Integer, ForeignKey("user.tgid"), primary_key=True) + + __table_args__ = (UniqueConstraint('mxid', 'mx_room', name='_mx_id_room'), ) + + class User(Base): __tablename__ = "user" @@ -55,5 +63,6 @@ class Puppet(Base): def init(db_factory): db = scoped_session(db_factory) Portal.query = db.query_property() + Message.query = db.query_property() User.query = db.query_property() Puppet.query = db.query_property() diff --git a/mautrix_telegram/formatter.py b/mautrix_telegram/formatter.py index 879ac111..d98e7a77 100644 --- a/mautrix_telegram/formatter.py +++ b/mautrix_telegram/formatter.py @@ -19,15 +19,24 @@ from html.parser import HTMLParser from collections import deque from telethon.tl.types import * from . import user as u, puppet as p +from .db import Message as DBMessage log = None -class MatrixParser(HTMLParser): - matrix_to_regex = re.compile("https://matrix.to/#/(@.+)") +class MessageEntityReply(MessageEntityUnknown): + def __init__(self, offset=0, length=0, msg_id=0): + super().__init__(offset, length) + self.msg_id = msg_id - def __init__(self): + +class MatrixParser(HTMLParser): + mention_regex = re.compile("https://matrix.to/#/(@.+)") + reply_regex = re.compile(r"https://matrix.to/#/(!.+?)/(\$.+)") + + def __init__(self, user_id=None): super().__init__() + self._user_id = user_id self.text = "" self.entities = [] self._building_entities = {} @@ -35,6 +44,7 @@ class MatrixParser(HTMLParser): self._open_tags = deque() self._open_tags_meta = deque() self._previous_ended_line = True + self._building_reply = False def handle_starttag(self, tag, attrs): self._open_tags.appendleft(tag) @@ -63,7 +73,8 @@ class MatrixParser(HTMLParser): url = attrs["href"] except KeyError: return - mention = self.matrix_to_regex.search(url) + mention = self.mention_regex.search(url) + reply = self.reply_regex.search(url) if mention: mxid = mention.group(1) puppet_match = p.Puppet.mxid_regex.search(mxid) @@ -79,6 +90,19 @@ class MatrixParser(HTMLParser): else: EntityType = MessageEntityMentionName args["user_id"] = user.tgid + elif reply and self._user_id and ( + len(self.entities) == 0 and len(self._building_entities) == 0): + room_id = reply.group(1) + message_id = reply.group(2) + message = DBMessage.query.filter(DBMessage.mxid == message_id + and DBMessage.mx_room == room_id + and DBMessage.user == self._user_id).one_or_none() + if not message: + return + EntityType = MessageEntityReply + args["msg_id"] = message.tgid + self._building_reply = True + url = None elif url.startswith("mailto:"): url = url[len("mailto:"):] EntityType = MessageEntityEmail @@ -104,6 +128,8 @@ class MatrixParser(HTMLParser): def handle_data(self, text): text = unescape(text) + if self._building_reply: + return previous_tag = self._open_tags[0] if len(self._open_tags) > 0 else "" list_format_offset = 0 if previous_tag == "a": @@ -142,6 +168,8 @@ class MatrixParser(HTMLParser): self._open_tags_meta.popleft() except IndexError: pass + if tag == "a": + self._building_reply = False if (tag == "ul" or tag == "ol") and self.text.endswith("\n"): self.text = self.text[:-1] entity = self._building_entities.pop(tag, None) @@ -149,15 +177,50 @@ class MatrixParser(HTMLParser): self.entities.append(entity) -def matrix_to_telegram(html): +def matrix_to_telegram(html, user_id=None): try: - parser = MatrixParser() + parser = MatrixParser(user_id) parser.feed(html) return parser.text, parser.entities except: log.exception("Failed to convert Matrix format:\nhtml=%s", html) +def telegram_event_to_matrix(evt, source): + text = evt.message + html = telegram_to_matrix(evt.message, evt.entities) if evt.entities else None + + if evt.fwd_from: + if not html: + html = escape(text) + id = evt.fwd_from.from_id + user = u.User.get_by_tgid(id) + if user: + fwd_from = f"{user.mxid}" + else: + puppet = p.Puppet.get(id, create=False) + if puppet and puppet.displayname: + fwd_from = f"{puppet.displayname}" + else: + user = source.client.get_entity(id) + if user: + fwd_from = p.Puppet.get_displayname(user, format=False) + if not fwd_from: + fwd_from = "Unknown user" + html = (f"Forwarded message from {fwd_from}
" + f"
{html}
") + + if evt.reply_to_msg_id: + msg = DBMessage.query.get((evt.reply_to_msg_id, source.tgid)) + quote = f"Quote
" + if html: + html = quote + html + else: + html = quote + escape(text) + + return text, html + + def telegram_to_matrix(text, entities): try: return _telegram_to_matrix(text, entities) @@ -234,6 +297,7 @@ def _telegram_to_matrix(text, entities): skip_entity = True last_offset = entity.offset + (0 if skip_entity else entity.length) html.append(text[last_offset:]) + return "".join(html) diff --git a/mautrix_telegram/matrix.py b/mautrix_telegram/matrix.py index a807a810..93c1e5bd 100644 --- a/mautrix_telegram/matrix.py +++ b/mautrix_telegram/matrix.py @@ -61,7 +61,7 @@ class MatrixHandler: text = text[len(prefix) + 1:] return is_command, text - def handle_message(self, room, sender, message): + def handle_message(self, room, sender, message, event_id): self.log.debug(f"{sender} sent {message} to ${room}") is_command, text = self.is_command(message) @@ -69,7 +69,7 @@ class MatrixHandler: portal = Portal.get_by_mxid(room) if portal and not is_command: - portal.handle_matrix_message(sender, message) + portal.handle_matrix_message(sender, message, event_id) return if message["msgtype"] != "m.text": @@ -105,4 +105,4 @@ class MatrixHandler: elif membership == "join": pass elif type == "m.room.message": - self.handle_message(evt["room_id"], evt["sender"], content) + self.handle_message(evt["room_id"], evt["sender"], content, evt["event_id"]) diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index 50d9160b..1548fb73 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -18,7 +18,7 @@ from io import BytesIO from telethon.tl.functions.messages import GetFullChatRequest from telethon.tl.functions.channels import GetParticipantsRequest from telethon.tl.types import * -from .db import Portal as DBPortal +from .db import Portal as DBPortal, Message as DBMessage from . import puppet as p, formatter config = None @@ -117,26 +117,36 @@ class Portal: if changed: self.save() - def handle_matrix_message(self, sender, message): + def handle_matrix_message(self, sender, message, event_id): type = message["msgtype"] if type == "m.text": if "format" in message and message["format"] == "org.matrix.custom.html": - message, entities = formatter.matrix_to_telegram(message["formatted_body"]) - sender.send_message(self.peer, message, entities=entities) + message, entities = formatter.matrix_to_telegram(message["formatted_body"], + sender.tgid) + reply_to = None + if len(entities) > 0 and isinstance(entities[0], formatter.MessageEntityReply): + reply = entities.pop(0) + # message = message[:reply.offset] + message[reply.offset + reply.length:] + reply_to = reply.msg_id + response = sender.send_message(self.peer, message, entities=entities, + reply_to=reply_to) else: - sender.send_message(self.peer, message["body"]) + response = sender.send_message(self.peer, message["body"]) + self.db.add( + DBMessage(tgid=response.id, mx_room=self.mxid, mxid=event_id, user=sender.tgid)) + self.db.commit() def handle_telegram_typing(self, user, event): user.intent.set_typing(self.mxid, is_typing=True) - def handle_telegram_message(self, sender, evt): + def handle_telegram_message(self, source, sender, evt): self.log.debug("Sending %s to %s by %d", evt.message, self.mxid, sender.id) if evt.message: - if evt.entities: - html = formatter.telegram_to_matrix(evt.message, evt.entities) - sender.intent.send_text(self.mxid, evt.message, html=html) - else: - sender.intent.send_text(self.mxid, evt.message) + text, html = formatter.telegram_event_to_matrix(evt, source) + response = sender.intent.send_text(self.mxid, text, html=html) + self.db.add(DBMessage(tgid=evt.id, mx_room=self.mxid, mxid=response["event_id"], + user=source.tgid)) + self.db.commit() def update_title(self, title, intent=None): if self.title != title: @@ -146,7 +156,6 @@ class Portal: return True return False - def update_avatar(self, user, photo, intent=None): photo_id = f"{photo.volume_id}-{photo.local_id}" if self.photo_id != photo_id: diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index 3de59c08..6cf276ec 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -51,7 +51,7 @@ class Puppet: self.db.commit() @staticmethod - def get_displayname(info): + def get_displayname(info, format=True): if info.first_name or info.last_name: name = " ".join([info.first_name or "", info.last_name or ""]).strip() elif info.username: @@ -60,6 +60,8 @@ class Puppet: name = info.phone_number else: name = info.id + if not format: + return name return config.get("bridge.displayname_template", "{} (Telegram)").format(name) def update_info(self, info): diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py index 2903d4fa..5f301fcd 100644 --- a/mautrix_telegram/user.py +++ b/mautrix_telegram/user.py @@ -170,7 +170,7 @@ class User: else: self.log.debug("Handling message portal=%s sender=%s update=%s", portal, sender, update) - portal.handle_telegram_message(sender, update) + portal.handle_telegram_message(self, sender, update) @classmethod def get_by_mxid(cls, mxid, create=True):