diff --git a/mautrix_telegram/db/message.py b/mautrix_telegram/db/message.py index 1180a1a8..ae6c3c26 100644 --- a/mautrix_telegram/db/message.py +++ b/mautrix_telegram/db/message.py @@ -152,6 +152,17 @@ class Message: rows = await cls.db.fetch(q, mx_room, tg_space, *mxids) return [cls._from_row(row) for row in rows] + @classmethod + async def find_recent( + cls, mx_room: RoomID, not_sender: TelegramID, limit: int = 20 + ) -> list[Message]: + q = f""" + SELECT {cls.columns} FROM message + WHERE mx_room=$1 AND sender<>$2 + ORDER BY tgid DESC LIMIT $3 + """ + return [cls._from_row(row) for row in await cls.db.fetch(q, mx_room, not_sender, limit)] + @classmethod async def replace_temp_mxid(cls, temp_mxid: str, mx_room: RoomID, real_mxid: EventID) -> None: q = "UPDATE message SET mxid=$1 WHERE mxid=$2 AND mx_room=$3" diff --git a/mautrix_telegram/db/puppet.py b/mautrix_telegram/db/puppet.py index d97c8f1a..b02b6710 100644 --- a/mautrix_telegram/db/puppet.py +++ b/mautrix_telegram/db/puppet.py @@ -50,6 +50,7 @@ class Puppet: avatar_set: bool is_bot: bool | None is_channel: bool + is_premium: bool custom_mxid: UserID | None access_token: str | None @@ -67,7 +68,8 @@ class Puppet: columns: ClassVar[str] = ( "id, is_registered, displayname, displayname_source, displayname_contact, " "displayname_quality, disable_updates, username, phone, photo_id, avatar_url, " - "name_set, avatar_set, is_bot, is_channel, custom_mxid, access_token, next_batch, base_url" + "name_set, avatar_set, is_bot, is_channel, is_premium, " + "custom_mxid, access_token, next_batch, base_url" ) @classmethod @@ -108,6 +110,7 @@ class Puppet: self.avatar_set, self.is_bot, self.is_channel, + self.is_premium, self.custom_mxid, self.access_token, self.next_batch, @@ -120,7 +123,7 @@ class Puppet: SET is_registered=$2, displayname=$3, displayname_source=$4, displayname_contact=$5, displayname_quality=$6, disable_updates=$7, username=$8, phone=$9, photo_id=$10, avatar_url=$11, name_set=$12, avatar_set=$13, is_bot=$14, is_channel=$15, - custom_mxid=$16, access_token=$17, next_batch=$18, base_url=$19 + is_premium=$16, custom_mxid=$17, access_token=$18, next_batch=$19, base_url=$20 WHERE id=$1 """ await self.db.execute(q, *self._values) @@ -130,8 +133,9 @@ class Puppet: INSERT INTO puppet ( id, is_registered, displayname, displayname_source, displayname_contact, displayname_quality, disable_updates, username, phone, photo_id, avatar_url, name_set, - avatar_set, is_bot, is_channel, custom_mxid, access_token, next_batch, base_url + avatar_set, is_bot, is_channel, is_premium, custom_mxid, access_token, next_batch, + base_url ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, - $19) + $19, $20) """ await self.db.execute(q, *self._values) diff --git a/mautrix_telegram/db/reaction.py b/mautrix_telegram/db/reaction.py index 9880489f..95000650 100644 --- a/mautrix_telegram/db/reaction.py +++ b/mautrix_telegram/db/reaction.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, ClassVar from asyncpg import Record from attr import dataclass +from telethon.tl.types import ReactionCustomEmoji, ReactionEmoji, TypeReaction from mautrix.types import EventID, RoomID from mautrix.util.async_db import Database @@ -58,9 +59,10 @@ class Reaction: @classmethod async def get_by_sender( cls, mxid: EventID, mx_room: RoomID, tg_sender: TelegramID - ) -> Reaction | None: + ) -> list[Reaction]: q = f"SELECT {cls.columns} FROM reaction WHERE msg_mxid=$1 AND mx_room=$2 AND tg_sender=$3" - return cls._from_row(await cls.db.fetchrow(q, mxid, mx_room, tg_sender)) + rows = await cls.db.fetch(q, mxid, mx_room, tg_sender) + return [cls._from_row(row) for row in rows] @classmethod async def get_all_by_message(cls, mxid: EventID, mx_room: RoomID) -> list[Reaction]: @@ -68,6 +70,13 @@ class Reaction: rows = await cls.db.fetch(q, mxid, mx_room) return [cls._from_row(row) for row in rows] + @property + def telegram(self) -> TypeReaction: + if self.reaction.isdecimal(): + return ReactionCustomEmoji(document_id=int(self.reaction)) + else: + return ReactionEmoji(emoticon=self.reaction) + @property def _values(self): return ( @@ -81,11 +90,11 @@ class Reaction: async def save(self) -> None: q = """ INSERT INTO reaction (mxid, mx_room, msg_mxid, tg_sender, reaction) - VALUES ($1, $2, $3, $4, $5) ON CONFLICT (msg_mxid, mx_room, tg_sender) - DO UPDATE SET mxid=$1, reaction=$5 + VALUES ($1, $2, $3, $4, $5) ON CONFLICT (msg_mxid, mx_room, tg_sender, reaction) + DO UPDATE SET mxid=excluded.mxid """ await self.db.execute(q, *self._values) async def delete(self) -> None: - q = "DELETE FROM reaction WHERE msg_mxid=$1 AND mx_room=$2 AND tg_sender=$3" - await self.db.execute(q, self.msg_mxid, self.mx_room, self.tg_sender) + q = "DELETE FROM reaction WHERE msg_mxid=$1 AND mx_room=$2 AND tg_sender=$3 AND reaction=$4" + await self.db.execute(q, self.msg_mxid, self.mx_room, self.tg_sender, self.reaction) diff --git a/mautrix_telegram/db/telegram_file.py b/mautrix_telegram/db/telegram_file.py index 1e48fb22..74a77a64 100644 --- a/mautrix_telegram/db/telegram_file.py +++ b/mautrix_telegram/db/telegram_file.py @@ -82,6 +82,11 @@ class TelegramFile: file.thumbnail = await cls.get(thumbnail_id, _thumbnail=True) return file + @classmethod + async def find_by_mxc(cls, mxc: ContentURI) -> TelegramFile | None: + q = f"SELECT {cls.columns} FROM telegram_file WHERE mxc=$1" + return cls._from_row(await cls.db.fetchrow(q, mxc)) + async def insert(self) -> None: q = ( "INSERT INTO telegram_file (id, mxc, mime_type, was_converted, size, width, height, " diff --git a/mautrix_telegram/db/upgrade/__init__.py b/mautrix_telegram/db/upgrade/__init__.py index 540c4ebf..574eb003 100644 --- a/mautrix_telegram/db/upgrade/__init__.py +++ b/mautrix_telegram/db/upgrade/__init__.py @@ -15,4 +15,5 @@ from . import ( v10_more_backfill_fields, v11_backfill_queue, v12_message_sender, + v13_multiple_reactions, ) diff --git a/mautrix_telegram/db/upgrade/v00_latest_revision.py b/mautrix_telegram/db/upgrade/v00_latest_revision.py index 3e9328c0..9ddb417c 100644 --- a/mautrix_telegram/db/upgrade/v00_latest_revision.py +++ b/mautrix_telegram/db/upgrade/v00_latest_revision.py @@ -26,6 +26,7 @@ async def create_latest_tables(conn: Connection) -> int: tg_username TEXT, tg_phone TEXT, is_bot BOOLEAN NOT NULL DEFAULT false, + is_premium BOOLEAN NOT NULL DEFAULT false, saved_contacts INTEGER NOT NULL DEFAULT 0 )""" ) @@ -78,7 +79,7 @@ async def create_latest_tables(conn: Connection) -> int: tg_sender BIGINT, reaction TEXT NOT NULL, - PRIMARY KEY (msg_mxid, mx_room, tg_sender), + PRIMARY KEY (msg_mxid, mx_room, tg_sender, reaction), UNIQUE (mxid, mx_room) )""" ) @@ -111,6 +112,7 @@ async def create_latest_tables(conn: Connection) -> int: avatar_set BOOLEAN NOT NULL DEFAULT false, is_bot BOOLEAN, is_channel BOOLEAN NOT NULL DEFAULT false, + is_premium BOOLEAN NOT NULL DEFAULT false, access_token TEXT, custom_mxid TEXT, @@ -135,6 +137,7 @@ async def create_latest_tables(conn: Connection) -> int: ON UPDATE CASCADE ON DELETE SET NULL )""" ) + await conn.execute("CREATE INDEX telegram_file_mxc_idx ON telegram_file(mxc)") await conn.execute( """CREATE TABLE bot_chat ( id BIGINT PRIMARY KEY, diff --git a/mautrix_telegram/db/upgrade/v13_multiple_reactions.py b/mautrix_telegram/db/upgrade/v13_multiple_reactions.py new file mode 100644 index 00000000..0eef72c3 --- /dev/null +++ b/mautrix_telegram/db/upgrade/v13_multiple_reactions.py @@ -0,0 +1,54 @@ +# 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 mautrix.util.async_db import Connection, Scheme + +from . import upgrade_table + + +@upgrade_table.register(description="Allow multiple reactions from the same user") +async def upgrade_v13(conn: Connection, scheme: Scheme) -> None: + await conn.execute("CREATE INDEX telegram_file_mxc_idx ON telegram_file(mxc)") + await conn.execute('ALTER TABLE "user" ADD COLUMN is_premium BOOLEAN NOT NULL DEFAULT false') + await conn.execute("ALTER TABLE puppet ADD COLUMN is_premium BOOLEAN NOT NULL DEFAULT false") + if scheme == Scheme.POSTGRES: + await conn.execute( + """ + ALTER TABLE reaction + DROP CONSTRAINT reaction_pkey, + ADD CONSTRAINT reaction_pkey PRIMARY KEY (msg_mxid, mx_room, tg_sender, reaction) + """ + ) + else: + await conn.execute( + """CREATE TABLE new_reaction ( + mxid TEXT NOT NULL, + mx_room TEXT NOT NULL, + msg_mxid TEXT NOT NULL, + tg_sender BIGINT, + reaction TEXT NOT NULL, + + PRIMARY KEY (msg_mxid, mx_room, tg_sender, reaction), + UNIQUE (mxid, mx_room) + )""" + ) + await conn.execute( + """ + INSERT INTO new_reaction (mxid, mx_room, msg_mxid, tg_sender, reaction) + SELECT mxid, mx_room, msg_mxid, tg_sender, reaction FROM reaction + """ + ) + await conn.execute("DROP TABLE reaction") + await conn.execute("ALTER TABLE new_reaction RENAME TO reaction") diff --git a/mautrix_telegram/db/user.py b/mautrix_telegram/db/user.py index a35b1d8e..c474fbfe 100644 --- a/mautrix_telegram/db/user.py +++ b/mautrix_telegram/db/user.py @@ -37,6 +37,7 @@ class User: tg_username: str | None tg_phone: str | None is_bot: bool + is_premium: bool saved_contacts: int @classmethod @@ -45,7 +46,9 @@ class User: return None return cls(**row) - columns: ClassVar[str] = "mxid, tgid, tg_username, tg_phone, is_bot, saved_contacts" + columns: ClassVar[str] = ", ".join( + ("mxid", "tgid", "tg_username", "tg_phone", "is_bot", "is_premium", "saved_contacts") + ) @classmethod async def get_by_tgid(cls, tgid: TelegramID) -> User | None: @@ -78,21 +81,23 @@ class User: self.tg_username, self.tg_phone, self.is_bot, + self.is_premium, self.saved_contacts, ) async def save(self) -> None: - q = ( - 'UPDATE "user" SET tgid=$2, tg_username=$3, tg_phone=$4, is_bot=$5, saved_contacts=$6 ' - "WHERE mxid=$1" - ) + q = """ + UPDATE "user" SET tgid=$2, tg_username=$3, tg_phone=$4, is_bot=$5, is_premium=$6, + saved_contacts=$7 + WHERE mxid=$1 + """ await self.db.execute(q, *self._values) async def insert(self) -> None: - q = ( - 'INSERT INTO "user" (mxid, tgid, tg_username, tg_phone, is_bot, saved_contacts) ' - "VALUES ($1, $2, $3, $4, $5, $6)" - ) + q = """ + INSERT INTO "user" (mxid, tgid, tg_username, tg_phone, is_bot, is_premium, saved_contacts) + VALUES ($1, $2, $3, $4, $5, $6, $7) + """ await self.db.execute(q, *self._values) async def get_contacts(self) -> list[TelegramID]: diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index c9077619..c38b05e6 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -16,6 +16,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, List, Union, cast +from collections import defaultdict from datetime import datetime from html import escape as escape_html from sqlite3 import IntegrityError @@ -52,6 +53,7 @@ from telethon.tl.functions.messages import ( EditChatTitleRequest, ExportChatInviteRequest, GetMessageReactionsListRequest, + GetMessagesReactionsRequest, MigrateChatRequest, SendReactionRequest, SetTypingRequest, @@ -102,6 +104,8 @@ from telethon.tl.types import ( Photo, PhotoEmpty, ReactionCount, + ReactionCustomEmoji, + ReactionEmoji, SendMessageCancelAction, SendMessageTypingAction, SponsoredMessage, @@ -113,11 +117,13 @@ from telethon.tl.types import ( TypeMessage, TypeMessageAction, TypePeer, + TypeReaction, TypeUser, TypeUserFull, TypeUserProfilePhoto, UpdateChannelUserTyping, UpdateChatUserTyping, + UpdateMessageReactions, UpdateNewMessage, UpdateUserTyping, User, @@ -181,6 +187,7 @@ from .db import ( Message as DBMessage, Portal as DBPortal, Reaction as DBReaction, + TelegramFile as DBTelegramFile, ) from .tgclient import MautrixTelegramClient from .types import TelegramID @@ -204,6 +211,8 @@ UpdateTyping = Union[UpdateUserTyping, UpdateChatUserTyping, UpdateChannelUserTy TypeChatPhoto = Union[ChatPhoto, ChatPhotoEmpty, Photo, PhotoEmpty] MediaHandler = Callable[["au.AbstractUser", IntentAPI, Message, RelatesTo], Awaitable[EventID]] +REACTION_POLL_MIN_INTERVAL = 20 + class BridgingError(Exception): pass @@ -261,6 +270,8 @@ class Portal(DBPortal, BasePortal): _sponsored_seen: dict[UserID, bool] _new_messages_after_sponsored: bool + _prev_reaction_poll: dict[UserID, float] + _msg_conv: putil.TelegramMessageConverter def __init__( @@ -331,6 +342,8 @@ class Portal(DBPortal, BasePortal): self._new_messages_after_sponsored = True self._bridging_blocked_at_runtime = False + self._prev_reaction_poll = defaultdict(lambda: 0.0) + self._msg_conv = putil.TelegramMessageConverter(self) # region Properties @@ -1440,8 +1453,13 @@ class Portal(DBPortal, BasePortal): await user.client.send_read_acknowledge( self.peer, max_id=message.tgid, clear_mentions=True, clear_reactions=True ) - if self.peer_type == "channel" and not self.megagroup: - asyncio.create_task(self._try_handle_read_for_sponsored_msg(user, event_id, timestamp)) + if self.peer_type == "channel": + if not self.megagroup: + asyncio.create_task( + self._try_handle_read_for_sponsored_msg(user, event_id, timestamp) + ) + else: + asyncio.create_task(self._poll_telegram_reactions(user)) async def _preproc_kick_ban( self, user: u.User | p.Puppet, source: u.User @@ -2212,12 +2230,31 @@ class Portal(DBPortal, BasePortal): await real_deleter.client.delete_messages(self.peer, [message.tgid]) async def handle_matrix_reaction( - self, user: u.User, target_event_id: EventID, reaction: str, reaction_event_id: EventID + self, user: u.User, target_event_id: EventID, emoji: str, reaction_event_id: EventID ) -> None: + emoji_id = emoji + reaction = ReactionEmoji(emoticon=variation_selector.remove(emoji)) + if emoji.startswith("mxc://"): + db_reaction = await DBTelegramFile.find_by_mxc(ContentURI(emoji)) + if not db_reaction or not db_reaction.id.isdecimal(): + self.log.debug(f"Dropping unknown reaction {emoji} by {user.mxid}") + if not self.has_bot: + await self.main_intent.redact( + self.mxid, reaction_event_id, reason="Unrecognized custom emoji" + ) + await self._send_bridge_error( + user, + Exception("Unrecognized custom emoji"), + reaction_event_id, + EventType.REACTION, + ) + return + reaction = ReactionCustomEmoji(document_id=int(db_reaction.id)) + emoji_id = db_reaction.id try: async with self.reaction_lock(target_event_id): await self._handle_matrix_reaction( - user, target_event_id, reaction, reaction_event_id + user, target_event_id, emoji_id, reaction, reaction_event_id ) except IgnoredMessageError as e: self.log.debug(str(e)) @@ -2244,7 +2281,12 @@ class Portal(DBPortal, BasePortal): asyncio.create_task(self._send_message_status(reaction_event_id, err=None)) async def _handle_matrix_reaction( - self, user: u.User, target_event_id: EventID, emoji: str, reaction_event_id: EventID + self, + user: u.User, + target_event_id: EventID, + emoji_id: str, + reaction: TypeReaction, + reaction_event_id: EventID, ) -> None: tg_space = self.tgid if self.peer_type == "channel" else user.tgid msg = await DBMessage.get_by_mxid(target_event_id, self.mxid, tg_space) @@ -2259,23 +2301,34 @@ class Portal(DBPortal, BasePortal): elif msg.edit_index != 0: raise IgnoredMessageError(f"Ignoring Matrix reaction to edit event {target_event_id}") - emoji = variation_selector.remove(emoji) - existing_react = await DBReaction.get_by_sender(msg.mxid, msg.mx_room, user.tgid) - await user.client(SendReactionRequest(peer=self.peer, msg_id=msg.tgid, reaction=emoji)) - if existing_react: - puppet = await user.get_puppet() - await puppet.intent_for(self).redact(existing_react.mx_room, existing_react.mxid) - existing_react.mxid = reaction_event_id - existing_react.reaction = emoji - await existing_react.save() - else: - await DBReaction( - mxid=reaction_event_id, - mx_room=self.mxid, - msg_mxid=msg.mxid, - tg_sender=user.tgid, - reaction=emoji, - ).save() + existing_reacts = await DBReaction.get_by_sender(msg.mxid, msg.mx_room, user.tgid) + new_tg_reactions: list[TypeReaction] = [] + reactions_to_remove: list[DBReaction] = [] + max_reactions = 3 if user.is_premium else 1 + max_reactions -= 1 # Leave one reaction of space for the new reaction + for db_reaction in existing_reacts: + if db_reaction.reaction == emoji_id: + raise IgnoredMessageError("Ignoring duplicate Matrix reaction") + if len(new_tg_reactions) < max_reactions: + new_tg_reactions.append(db_reaction.telegram) + else: + reactions_to_remove.append(db_reaction) + new_tg_reactions.append(reaction) + + await user.client( + SendReactionRequest(peer=self.peer, msg_id=msg.tgid, reaction=new_tg_reactions) + ) + puppet = await user.get_puppet() + for db_reaction in reactions_to_remove: + await db_reaction.delete() + await puppet.intent_for(self).redact(db_reaction.mx_room, db_reaction.mxid) + await DBReaction( + mxid=reaction_event_id, + mx_room=self.mxid, + msg_mxid=msg.mxid, + tg_sender=user.tgid, + reaction=emoji_id, + ).save() async def _update_telegram_power_level( self, sender: u.User, user_id: TelegramID, level: int @@ -2681,35 +2734,46 @@ class Portal(DBPortal, BasePortal): return False def _split_dm_reaction_counts(self, counts: list[ReactionCount]) -> list[MessagePeerReaction]: - if len(counts) == 1: - item = counts[0] + reactions = [] + for item in counts: if item.count == 2: - return [ + reactions += [ MessagePeerReaction(reaction=item.reaction, peer_id=PeerUser(self.tgid)), MessagePeerReaction( reaction=item.reaction, peer_id=PeerUser(self.tg_receiver) ), ] elif item.count == 1: - return [ + reactions.append( MessagePeerReaction( reaction=item.reaction, - peer_id=PeerUser(self.tg_receiver if item.chosen else self.tgid), - ), - ] - elif len(counts) == 2: - item1, item2 = counts - return [ - MessagePeerReaction( - reaction=item1.reaction, - peer_id=PeerUser(self.tg_receiver if item1.chosen else self.tgid), - ), - MessagePeerReaction( - reaction=item2.reaction, - peer_id=PeerUser(self.tg_receiver if item2.chosen else self.tgid), - ), - ] - return [] + peer_id=PeerUser(self.tg_receiver if item.chosen_order else self.tgid), + ) + ) + return reactions + + async def _poll_telegram_reactions(self, source: au.AbstractUser) -> None: + now = time.monotonic() + if self._prev_reaction_poll[source.mxid] + REACTION_POLL_MIN_INTERVAL > now: + self.log.trace( + f"Not polling reactions through {source.mxid}, " + f"last poll was less than {REACTION_POLL_MIN_INTERVAL} seconds ago" + ) + return + self._prev_reaction_poll[source.mxid] = now + self.log.debug(f"Polling reactions for recent messages through {source.mxid}") + messages = await DBMessage.find_recent(self.mxid, source.tgid) + message_ids = [message.tgid for message in messages] + updates = await source.client(GetMessagesReactionsRequest(peer=self.peer, id=message_ids)) + for user in updates.users: + user: User + puppet = await p.Puppet.get_by_tgid(TelegramID(user.id)) + await puppet.update_info(source, user) + for upd in updates.updates: + if isinstance(upd, UpdateMessageReactions): + await self.handle_telegram_reactions(source, TelegramID(upd.msg_id), upd.reactions) + else: + self.log.warning(f"Unexpected update type {type(upd)} in get reactions response") async def try_handle_telegram_reactions( self, @@ -2732,9 +2796,15 @@ class Portal(DBPortal, BasePortal): dbm: DBMessage | None = None, timestamp: datetime | None = None, ) -> None: - if self.peer_type == "channel" and not self.megagroup: + total_count = sum(item.count for item in data.results) + recent_reactions = data.recent_reactions or [] + if total_count > 0 and not recent_reactions and not data.can_see_list: # We don't know who reacted in a channel, so we can't bridge it properly either return + if self.peer_type == "channel" and not self.megagroup: + # This should never happen with the previous if + self.log.warning(f"Can see reaction list in channel ({data!s})") + # return tg_space = self.tgid if self.peer_type == "channel" else source.tgid if dbm is None: @@ -2742,69 +2812,109 @@ class Portal(DBPortal, BasePortal): if dbm is None: return - total_count = sum(item.count for item in data.results) - recent_reactions = data.recent_reactions or [] - if not recent_reactions and total_count > 0: + if not recent_reactions or len(recent_reactions) < total_count: if self.peer_type == "user": recent_reactions = self._split_dm_reaction_counts(data.results) elif source.is_bot: # Can't fetch exact reaction senders as a bot return else: - # TODO this doesn't work for some reason - return - # resp = await source.client( - # GetMessageReactionsListRequest(peer=self.peer, id=dbm.tgid, limit=20) - # ) - # recent_reactions = resp.reactions + # TODO should calls to this be limited? + resp = await source.client( + GetMessageReactionsListRequest(peer=self.peer, id=dbm.tgid, limit=100) + ) + recent_reactions = resp.reactions async with self.reaction_lock(dbm.mxid): await self._handle_telegram_reactions_locked( - dbm, recent_reactions, total_count, timestamp=timestamp + source, dbm, recent_reactions, total_count, timestamp=timestamp ) + @staticmethod + def _reactions_filter(lst: list[TypeReaction], existing: DBReaction) -> bool: + if not lst: + return False + for reaction in lst: + if isinstance(reaction, ReactionCustomEmoji) and existing.reaction == str( + reaction.document_id + ): + lst.remove(reaction) + return True + elif isinstance(reaction, ReactionEmoji) and existing.reaction == reaction.emoticon: + lst.remove(reaction) + return True + return False + + @staticmethod + async def _get_reaction_limit(sender: TelegramID) -> int: + puppet = await p.Puppet.get_by_tgid(sender, create=False) + if puppet and puppet.is_premium: + return 3 + return 1 + async def _handle_telegram_reactions_locked( self, + source: au.AbstractUser, msg: DBMessage, reaction_list: list[MessagePeerReaction], total_count: int, timestamp: datetime | None = None, ) -> None: - reactions = { - p.Puppet.get_id_from_peer(reaction.peer_id): reaction.reaction - for reaction in reaction_list - if isinstance(reaction.peer_id, (PeerUser, PeerChannel)) - } - is_full = len(reactions) == total_count + reactions: dict[TelegramID, list[TypeReaction]] = {} + custom_emoji_ids: list[int] = [] + for reaction in reaction_list: + if isinstance(reaction.peer_id, (PeerUser, PeerChannel)) and isinstance( + reaction.reaction, (ReactionEmoji, ReactionCustomEmoji) + ): + reactions.setdefault(p.Puppet.get_id_from_peer(reaction.peer_id), []).append( + reaction.reaction + ) + if isinstance(reaction.reaction, ReactionCustomEmoji): + custom_emoji_ids.append(reaction.reaction.document_id) + is_full = len(reaction_list) == total_count + custom_emojis = await util.transfer_custom_emojis_to_matrix(source, custom_emoji_ids) existing_reactions = await DBReaction.get_all_by_message(msg.mxid, msg.mx_room) removed: list[DBReaction] = [] - changed: list[tuple[DBReaction, str]] = [] for existing_reaction in existing_reactions: - new_reaction = reactions.get(existing_reaction.tg_sender) - if new_reaction is None: - if is_full: + sender_id = existing_reaction.tg_sender + new_reactions = reactions.get(sender_id) + if self._reactions_filter(new_reactions, existing_reaction): + if new_reactions is not None and len(new_reactions) == 0: + reactions.pop(sender_id) + else: + if is_full or ( + new_reactions is not None + and len(new_reactions) == await self._get_reaction_limit(sender_id) + ): removed.append(existing_reaction) # else: assume the reaction is still there, too much effort to fetch it - elif new_reaction == existing_reaction.reaction: - reactions.pop(existing_reaction.tg_sender) - else: - changed.append((existing_reaction, new_reaction)) - for sender, new_emoji in reactions.items(): - self.log.debug(f"Bridging reaction {new_emoji} by {sender} to {msg.tgid}") - puppet: p.Puppet = await p.Puppet.get_by_tgid(sender) - mxid = await puppet.intent_for(self).react( - msg.mx_room, msg.mxid, variation_selector.add(new_emoji), timestamp=timestamp - ) - await DBReaction( - mxid=mxid, - mx_room=msg.mx_room, - msg_mxid=msg.mxid, - tg_sender=sender, - reaction=new_emoji, - ).save() + new_reaction: TypeReaction + for sender, new_reactions in reactions.items(): + for new_reaction in new_reactions: + if isinstance(new_reaction, ReactionEmoji): + emoji_id = new_reaction.emoticon + matrix_reaction = variation_selector.add(new_reaction.emoticon) + elif isinstance(new_reaction, ReactionCustomEmoji): + emoji_id = str(new_reaction.document_id) + matrix_reaction = custom_emojis[new_reaction.document_id].mxc + else: + self.log.warning("Unknown reaction type %s", type(new_reaction)) + continue + self.log.debug(f"Bridging reaction {emoji_id} by {sender} to {msg.tgid}") + puppet: p.Puppet = await p.Puppet.get_by_tgid(sender) + mxid = await puppet.intent_for(self).react( + msg.mx_room, msg.mxid, matrix_reaction, timestamp=timestamp + ) + await DBReaction( + mxid=mxid, + mx_room=msg.mx_room, + msg_mxid=msg.mxid, + tg_sender=sender, + reaction=emoji_id, + ).save() for removed_reaction in removed: self.log.debug( f"Removing reaction {removed_reaction.reaction} by {removed_reaction.tg_sender} " @@ -2813,19 +2923,6 @@ class Portal(DBPortal, BasePortal): puppet = await p.Puppet.get_by_tgid(removed_reaction.tg_sender) await puppet.intent_for(self).redact(removed_reaction.mx_room, removed_reaction.mxid) await removed_reaction.delete() - for changed_reaction, new_emoji in changed: - self.log.debug( - f"Updating reaction {changed_reaction.reaction} -> {new_emoji} " - f"by {changed_reaction.tg_sender} to {msg.tgid}" - ) - puppet = await p.Puppet.get_by_tgid(changed_reaction.tg_sender) - intent = puppet.intent_for(self) - await intent.redact(changed_reaction.mx_room, changed_reaction.mxid) - changed_reaction.mxid = await intent.react( - msg.mx_room, msg.mxid, variation_selector.add(new_emoji), timestamp=timestamp - ) - changed_reaction.reaction = new_emoji - await changed_reaction.save() async def handle_telegram_message( self, source: au.AbstractUser, sender: p.Puppet | None, evt: Message diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index 37e39076..5a34d1cf 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -80,6 +80,7 @@ class Puppet(DBPuppet, BasePuppet): avatar_set: bool = False, is_bot: bool = False, is_channel: bool = False, + is_premium: bool = False, custom_mxid: UserID | None = None, access_token: str | None = None, next_batch: SyncToken | None = None, @@ -101,6 +102,7 @@ class Puppet(DBPuppet, BasePuppet): avatar_set=avatar_set, is_bot=is_bot, is_channel=is_channel, + is_premium=is_premium, custom_mxid=custom_mxid, access_token=access_token, next_batch=next_batch, @@ -255,11 +257,15 @@ class Puppet(DBPuppet, BasePuppet): async def update_info(self, source: au.AbstractUser, info: User | Channel) -> None: is_bot = False if isinstance(info, Channel) else info.bot + is_premium = False if isinstance(info, Channel) else info.premium is_channel = isinstance(info, Channel) - changed = is_bot != self.is_bot or is_channel != self.is_channel + changed = ( + is_bot != self.is_bot or is_channel != self.is_channel or is_premium != self.is_premium + ) self.is_bot = is_bot self.is_channel = is_channel + self.is_premium = is_premium if self.username != info.username: self.username = info.username diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py index c74f2c7e..5f9f1a39 100644 --- a/mautrix_telegram/user.py +++ b/mautrix_telegram/user.py @@ -92,6 +92,7 @@ class User(DBUser, AbstractUser, BaseUser): tg_username: str | None = None, tg_phone: str | None = None, is_bot: bool = False, + is_premium: bool = False, saved_contacts: int = 0, ) -> None: super().__init__( @@ -100,6 +101,7 @@ class User(DBUser, AbstractUser, BaseUser): tg_username=tg_username, tg_phone=tg_phone, is_bot=is_bot, + is_premium=is_premium, saved_contacts=saved_contacts, ) AbstractUser.__init__(self) @@ -371,6 +373,9 @@ class User(DBUser, AbstractUser, BaseUser): if self.is_bot != info.bot: self.is_bot = info.bot changed = True + if self.is_premium != info.premium: + self.is_premium = info.premium + changed = True if self.tg_username != info.username: self.tg_username = info.username changed = True diff --git a/mautrix_telegram/util/__init__.py b/mautrix_telegram/util/__init__.py index 50ca02ff..acdfb074 100644 --- a/mautrix_telegram/util/__init__.py +++ b/mautrix_telegram/util/__init__.py @@ -1,4 +1,4 @@ from .color_log import ColorFormatter -from .file_transfer import convert_image, transfer_file_to_matrix +from .file_transfer import convert_image, transfer_custom_emojis_to_matrix, transfer_file_to_matrix from .parallel_file_transfer import parallel_transfer_to_telegram from .recursive_dict import recursive_del, recursive_get, recursive_set diff --git a/requirements.txt b/requirements.txt index 8f705124..f182259e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,7 @@ aiohttp>=3,<4 yarl>=1,<2 mautrix>=0.18.1,<0.19 #telethon>=1.24,<1.25 -tulir-telethon==1.25.0a20 +tulir-telethon==1.26.0a1 asyncpg>=0.20,<0.27 mako>=1,<2 setuptools