From 6bf3d10e292b8b775399affbd109204b1b54b02b Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Mon, 31 Jan 2022 15:43:39 +0200 Subject: [PATCH] Improve handling of disappearing photos and files Fixes #508 --- mautrix_telegram/__main__.py | 1 + mautrix_telegram/db/__init__.py | 14 +++- mautrix_telegram/db/disappearing_message.py | 78 +++++++++++++++++++ mautrix_telegram/db/upgrade/__init__.py | 2 +- .../db/upgrade/v04_disappearing_messages.py | 32 ++++++++ mautrix_telegram/portal.py | 20 ++--- 6 files changed, 135 insertions(+), 12 deletions(-) create mode 100644 mautrix_telegram/db/disappearing_message.py create mode 100644 mautrix_telegram/db/upgrade/v04_disappearing_messages.py diff --git a/mautrix_telegram/__main__.py b/mautrix_telegram/__main__.py index 61240e83..80612c18 100644 --- a/mautrix_telegram/__main__.py +++ b/mautrix_telegram/__main__.py @@ -86,6 +86,7 @@ class TelegramBridge(Bridge): Portal.init_cls(self) self.add_startup_actions(Puppet.init_cls(self)) self.add_startup_actions(User.init_cls(self)) + self.add_startup_actions(Portal.restart_scheduled_disappearing()) if self.bot: self.add_startup_actions(self.bot.start()) if self.config["bridge.resend_bridge_info"]: diff --git a/mautrix_telegram/db/__init__.py b/mautrix_telegram/db/__init__.py index 119ebdbd..40f0706e 100644 --- a/mautrix_telegram/db/__init__.py +++ b/mautrix_telegram/db/__init__.py @@ -16,6 +16,7 @@ from mautrix.util.async_db import Database from .bot_chat import BotChat +from .disappearing_message import DisappearingMessage from .message import Message from .portal import Portal from .puppet import Puppet @@ -27,7 +28,17 @@ from .user import User def init(db: Database) -> None: - for table in (Portal, Message, Reaction, User, Puppet, TelegramFile, BotChat, PgSession): + for table in ( + Portal, + Message, + Reaction, + User, + Puppet, + TelegramFile, + BotChat, + PgSession, + DisappearingMessage, + ): table.db = db @@ -42,4 +53,5 @@ __all__ = [ "TelegramFile", "BotChat", "PgSession", + "DisappearingMessage", ] diff --git a/mautrix_telegram/db/disappearing_message.py b/mautrix_telegram/db/disappearing_message.py new file mode 100644 index 00000000..6c217173 --- /dev/null +++ b/mautrix_telegram/db/disappearing_message.py @@ -0,0 +1,78 @@ +# mautrix-telegram - A Matrix-Telegram puppeting bridge +# Copyright (C) 2021 Sumner Evans +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +import asyncpg + +from mautrix.bridge import AbstractDisappearingMessage +from mautrix.types import EventID, RoomID +from mautrix.util.async_db import Database + +fake_db = Database.create("") if TYPE_CHECKING else None + + +class DisappearingMessage(AbstractDisappearingMessage): + db: ClassVar[Database] = fake_db + + async def insert(self) -> None: + q = """ + INSERT INTO disappearing_message (room_id, event_id, expiration_seconds, expiration_ts) + VALUES ($1, $2, $3, $4) + """ + await self.db.execute( + q, self.room_id, self.event_id, self.expiration_seconds, self.expiration_ts + ) + + async def update(self) -> None: + q = "UPDATE disappearing_message SET expiration_ts=$3 WHERE room_id=$1 AND event_id=$2" + await self.db.execute(q, self.room_id, self.event_id, self.expiration_ts) + + async def delete(self) -> None: + q = "DELETE from disappearing_message WHERE room_id=$1 AND event_id=$2" + await self.db.execute(q, self.room_id, self.event_id) + + @classmethod + def _from_row(cls, row: asyncpg.Record) -> DisappearingMessage: + return cls(**row) + + @classmethod + async def get(cls, room_id: RoomID, event_id: EventID) -> DisappearingMessage | None: + q = """ + SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message + WHERE room_id=$1 AND mxid=$2 + """ + try: + return cls._from_row(await cls.db.fetchrow(q, room_id, event_id)) + except Exception: + return None + + @classmethod + async def get_all_scheduled(cls) -> list[DisappearingMessage]: + q = """ + SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message + WHERE expiration_ts IS NOT NULL + """ + return [cls._from_row(r) for r in await cls.db.fetch(q)] + + @classmethod + async def get_unscheduled_for_room(cls, room_id: RoomID) -> list[DisappearingMessage]: + q = """ + SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message + WHERE room_id = $1 AND expiration_ts IS NULL + """ + return [cls._from_row(r) for r in await cls.db.fetch(q, room_id)] diff --git a/mautrix_telegram/db/upgrade/__init__.py b/mautrix_telegram/db/upgrade/__init__.py index 18875090..33fc194e 100644 --- a/mautrix_telegram/db/upgrade/__init__.py +++ b/mautrix_telegram/db/upgrade/__init__.py @@ -2,4 +2,4 @@ from mautrix.util.async_db import UpgradeTable upgrade_table = UpgradeTable() -from . import v01_initial_revision, v02_sponsored_events, v03_reactions +from . import v01_initial_revision, v02_sponsored_events, v03_reactions, v04_disappearing_messages diff --git a/mautrix_telegram/db/upgrade/v04_disappearing_messages.py b/mautrix_telegram/db/upgrade/v04_disappearing_messages.py new file mode 100644 index 00000000..411aa8ff --- /dev/null +++ b/mautrix_telegram/db/upgrade/v04_disappearing_messages.py @@ -0,0 +1,32 @@ +# mautrix-telegram - A Matrix-Telegram puppeting bridge +# Copyright (C) 2022 Tulir Asokan +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +from asyncpg import Connection + +from . import upgrade_table + + +@upgrade_table.register(description="Add support for disappearing messages") +async def upgrade_v4(conn: Connection) -> None: + await conn.execute( + """CREATE TABLE disappearing_message ( + room_id TEXT, + event_id TEXT, + expiration_seconds BIGINT, + expiration_ts BIGINT, + + PRIMARY KEY (room_id, event_id) + )""" + ) diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index b60be3a4..31f25117 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -195,6 +195,7 @@ from mautrix.util.simple_template import SimpleTemplate from . import abstract_user as au, formatter, portal_util as putil, puppet as p, user as u, util from .config import Config from .db import ( + DisappearingMessage, Message as DBMessage, Portal as DBPortal, Reaction as DBReaction, @@ -243,6 +244,7 @@ class DocAttrs(NamedTuple): class Portal(DBPortal, BasePortal): bot: "Bot" config: Config + disappearing_msg_class = DisappearingMessage # Instance cache by_mxid: dict[RoomID, Portal] = {} @@ -321,6 +323,7 @@ class Portal(DBPortal, BasePortal): photo_id=photo_id, local_config=local_config or {}, ) + BasePortal.__init__(self) self.log = self.log.getChild(self.tgid_log if self.tgid else self.mxid) self._main_intent = None self.deleted = False @@ -2106,15 +2109,6 @@ class Portal(DBPortal, BasePortal): return f"https://t.me/c/{self.tgid}/{evt.id}" return None - async def _expire_telegram_photo(self, intent: IntentAPI, event_id: EventID, ttl: int) -> None: - try: - content = TextMessageEventContent(msgtype=MessageType.NOTICE, body="Photo has expired") - content.set_edit(event_id) - await asyncio.sleep(ttl) - await self._send_message(intent, content) - except Exception: - self.log.warning("Failed to expire Telegram photo %s", event_id, exc_info=True) - async def _handle_telegram_photo( self, source: au.AbstractUser, intent: IntentAPI, evt: Message, relates_to: RelatesTo ) -> EventID | None: @@ -2171,13 +2165,15 @@ class Portal(DBPortal, BasePortal): content.url = file.mxc result = await self._send_message(intent, content, timestamp=evt.date) if media.ttl_seconds: - asyncio.create_task(self._expire_telegram_photo(intent, result, media.ttl_seconds)) + await DisappearingMessage(self.mxid, result, media.ttl_seconds).insert() if evt.message: caption_content = await formatter.telegram_to_matrix( evt, source, self.main_intent, no_reply_fallback=True ) caption_content.external_url = content.external_url result = await self._send_message(intent, caption_content, timestamp=evt.date) + if media.ttl_seconds: + await DisappearingMessage(self.mxid, result, media.ttl_seconds).insert() return result @staticmethod @@ -2355,12 +2351,16 @@ class Portal(DBPortal, BasePortal): else: content.url = file.mxc res = await self._send_message(intent, content, event_type=event_type, timestamp=evt.date) + if evt.media.ttl_seconds: + await DisappearingMessage(self.mxid, res, evt.media.ttl_seconds).insert() if evt.message: caption_content = await formatter.telegram_to_matrix( evt, source, self.main_intent, no_reply_fallback=True ) caption_content.external_url = content.external_url res = await self._send_message(intent, caption_content, timestamp=evt.date) + if evt.media.ttl_seconds: + await DisappearingMessage(self.mxid, res, evt.media.ttl_seconds).insert() return res def _location_message_to_content(