Refactor and clean up code

This commit is contained in:
Tulir Asokan
2018-02-23 21:06:28 +02:00
parent 2aa48f37a9
commit 9e5843a0dc
9 changed files with 257 additions and 186 deletions
+14 -9
View File
@@ -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)
+4 -6
View File
@@ -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
+3 -2
View File
@@ -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.")
+37 -24
View File
@@ -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 <user identifier>`")
@@ -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}")
+6 -8
View File
@@ -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
+124 -86
View File
@@ -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"<a href='https://matrix.to/#/{user.mxid}'>{user.mxid}</a>"
else:
puppet = pu.Puppet.get(fwd_from_id, create=False)
if puppet and puppet.displayname:
fwd_from = f"<a href='https://matrix.to/#/{puppet.mxid}'>{puppet.displayname}</a>"
else:
user = await source.client.get_entity(fwd_from_id)
if user:
fwd_from = f"<b>{pu.Puppet.get_displayname(user, format=False)}</b>"
else:
fwd_from = None
if not fwd_from:
fwd_from = "<b>Unknown user</b>"
text = f"Forwarded from {fwd_from}:\n{text}"
html = (f"Forwarded message from {fwd_from}<br/>"
f"<blockquote>{html}</blockquote>")
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"<u>Edit:</u> {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"<a href='https://matrix.to/#/{sender}'>{reply_displayname}</a>"
reply_to_msg = (("<a href='https://matrix.to/#/"
f"{msg.mx_room}/{msg.mxid}'>{reply_text}</a>")
if message_link_in_reply else "Reply")
quote = f"{reply_to_msg} to {reply_to_user}<blockquote>{body}</blockquote>"
except (ValueError, KeyError, MatrixRequestError):
quote = f"{reply_text} to unknown user <em>(Failed to fetch message)</em>:<br/>"
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"<a href='https://matrix.to/#/{user.mxid}'>{user.mxid}</a>"
else:
puppet = pu.Puppet.get(from_id, create=False)
if puppet and puppet.displayname:
fwd_from = f"<a href='https://matrix.to/#/{puppet.mxid}'>{puppet.displayname}</a>"
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 <b>{fwd_from}</b><br/>"
f"<blockquote>{html}</blockquote>")
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 = "<u>Edit:</u> " + (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"<a href='https://matrix.to/#/{sender}'>{displayname}</a>"
reply_to_msg = (("<a href='https://matrix.to/#/"
f"{msg.mx_room}/{msg.mxid}'>{reply_text}</a>")
if message_link_in_reply else "Reply")
quote = f"{reply_to_msg} to {reply_to_user}<blockquote>{body}</blockquote>"
except (ValueError, KeyError, MatrixRequestError):
quote = "{reply_text} to unknown user <em>(Failed to fetch message)</em>:<br/>"
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"<code>{entity_text}</code>")
elif entity_type == MessageEntityPre:
if entity.language:
html.append("<pre>"
f"<code class='language-{entity.language}'>{entity_text}</code>"
"</pre>")
else:
html.append(f"<pre><code>{entity_text}</code></pre>")
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"<a href='https://matrix.to/#/{mxid}'>{entity_text}</a>")
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"<a href='https://matrix.to/#/{mxid}'>{entity_text}</a>")
else:
skip_entity = True
skip_entity = _parse_name_mention(html, entity_text, entity.user_id)
elif entity_type == MessageEntityEmail:
html.append(f"<a href='mailto:{entity_text}'>{entity_text}</a>")
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"<a href='{url}'>{entity_text}</a>")
skip_entity = _parse_url(html, entity_text, entity_type, entity.url)
elif entity_type == MessageEntityBotCommand:
html.append(f"<font color='blue'>!{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("<pre>"
f"<code class='language-{language}'>{entity_text}</code>"
"</pre>")
else:
html.append(f"<pre><code>{entity_text}</code></pre>")
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"<a href='https://matrix.to/#/{mxid}'>{entity_text}</a>")
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"<a href='https://matrix.to/#/{mxid}'>{entity_text}</a>")
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"<a href='{url}'>{entity_text}</a>")
return False
+2 -2
View File
@@ -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"<a href='https://matrix.to/#/{self.az.intent.mxid}'>the bridge bot</a> "
f"<a href='https://matrix.to/#/{self.az.bot_mxid}'>the bridge bot</a> "
f"first if you want to create a Telegram chat."))
return
+64 -47
View File
@@ -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"&lt;{sender.displayname}&gt; "
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):
+3 -2
View File
@@ -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):