iartamonov/mautrix-telegram
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Improve handling of disappearing photos and files

Browse Source 
Fixes #508
This commit is contained in:
Tulir Asokan
2022-01-31 15:43:39 +02:00
parent ebd2a38e56
commit 6bf3d10e29
6 changed files with 135 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mautrix_telegram/__main__.py
+1
View File
@@ -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"]:
mautrix_telegram/db/__init__.py
+13 -1
View File
@@ -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",
]
mautrix_telegram/db/disappearing_message.py
+78
View File
@@ -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 <https://www.gnu.org/licenses/>.
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)]
mautrix_telegram/db/upgrade/__init__.py
+1 -1
View File
@@ -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
mautrix_telegram/db/upgrade/v04_disappearing_messages.py
+32
View File
@@ -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 <https://www.gnu.org/licenses/>.
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)
)"""
)
mautrix_telegram/portal.py
+10 -10
View File
@@ -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(