diff --git a/alembic/versions/d3c922a6acd2_add_decryption_info_field_for_.py b/alembic/versions/d3c922a6acd2_add_decryption_info_field_for_.py new file mode 100644 index 00000000..1469f2a8 --- /dev/null +++ b/alembic/versions/d3c922a6acd2_add_decryption_info_field_for_.py @@ -0,0 +1,26 @@ +"""Add decryption info field for reuploaded telegram files + +Revision ID: d3c922a6acd2 +Revises: 24f31fc8a72b +Create Date: 2020-03-30 20:07:17.340346 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = 'd3c922a6acd2' +down_revision = '24f31fc8a72b' +branch_labels = None +depends_on = None + + +def upgrade(): + with op.batch_alter_table("telegram_file") as batch_op: + batch_op.add_column(sa.Column("decryption_info", sa.Text(), nullable=True)) + + +def downgrade(): + with op.batch_alter_table("telegram_file") as batch_op: + batch_op.drop_column("decryption_info") diff --git a/mautrix_telegram/db/telegram_file.py b/mautrix_telegram/db/telegram_file.py index 8cc0c9c0..82f1e17c 100644 --- a/mautrix_telegram/db/telegram_file.py +++ b/mautrix_telegram/db/telegram_file.py @@ -13,15 +13,37 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -from typing import Optional +from typing import Optional, cast, Dict, Any -from sqlalchemy import Column, ForeignKey, Integer, BigInteger, String, Boolean +from sqlalchemy import (Column, ForeignKey, Integer, BigInteger, String, Boolean, Text, + TypeDecorator) from sqlalchemy.engine.result import RowProxy -from mautrix.types import ContentURI +from mautrix.types import ContentURI, EncryptedFile from mautrix.util.db import Base +class DBEncryptedFile(TypeDecorator): + impl = Text + + @property + def python_type(self): + return EncryptedFile + + def process_bind_param(self, value: EncryptedFile, dialect) -> Optional[str]: + if value is not None: + return value.json() + return None + + def process_result_value(self, value: str, dialect) -> Optional[EncryptedFile]: + if value is not None: + return EncryptedFile.parse_json(value) + return None + + def process_literal_param(self, value, dialect): + return value + + class TelegramFile(Base): __tablename__ = "telegram_file" @@ -33,12 +55,13 @@ class TelegramFile(Base): size: Optional[int] = Column(Integer, nullable=True) width: Optional[int] = Column(Integer, nullable=True) height: Optional[int] = Column(Integer, nullable=True) + decryption_info: Optional[Dict[str, Any]] = Column(DBEncryptedFile, nullable=True) thumbnail_id: str = Column("thumbnail", String, ForeignKey("telegram_file.id"), nullable=True) thumbnail: Optional['TelegramFile'] = None @classmethod def scan(cls, row: RowProxy) -> 'TelegramFile': - telegram_file: TelegramFile = super().scan(row) + telegram_file = cast(TelegramFile, super().scan(row)) if isinstance(telegram_file.thumbnail, str): telegram_file.thumbnail = cls.get(telegram_file.thumbnail) return telegram_file @@ -52,5 +75,5 @@ class TelegramFile(Base): conn.execute(self.t.insert().values( id=self.id, mxc=self.mxc, mime_type=self.mime_type, was_converted=self.was_converted, timestamp=self.timestamp, size=self.size, - width=self.width, height=self.height, + width=self.width, height=self.height, decryption_info=self.decryption_info, thumbnail=self.thumbnail.id if self.thumbnail else self.thumbnail_id)) diff --git a/mautrix_telegram/matrix.py b/mautrix_telegram/matrix.py index 9f168e7c..6d600669 100644 --- a/mautrix_telegram/matrix.py +++ b/mautrix_telegram/matrix.py @@ -54,7 +54,8 @@ class MatrixHandler(BaseMatrixHandler): self.user_id_suffix = f"{suffix}:{homeserver}" super(MatrixHandler, self).__init__(context.az, context.config, loop=context.loop, - command_processor=com.CommandProcessor(context)) + command_processor=com.CommandProcessor(context), + bridge=context.bridge) self.bot = context.bot self.previously_typing = {} diff --git a/mautrix_telegram/portal/telegram.py b/mautrix_telegram/portal/telegram.py index 61f655fd..bd782797 100644 --- a/mautrix_telegram/portal/telegram.py +++ b/mautrix_telegram/portal/telegram.py @@ -80,7 +80,8 @@ class PortalTelegram(BasePortal, ABC): async def handle_telegram_photo(self, source: 'AbstractUser', intent: IntentAPI, evt: Message, relates_to: Dict = None) -> Optional[EventID]: loc, largest_size = self._get_largest_photo_size(evt.media.photo) - file = await util.transfer_file_to_matrix(source.client, intent, loc) + file = await util.transfer_file_to_matrix(source.client, intent, loc, + encrypt=self.encrypted) if not file: return None if self.get_config("inline_images") and (evt.message @@ -98,9 +99,13 @@ class PortalTelegram(BasePortal, ABC): else largest_size.size)) name = f"image{sane_mimetypes.guess_extension(file.mime_type)}" await intent.set_typing(self.mxid, is_typing=False) - content = MediaMessageEventContent(url=file.mxc, msgtype=MessageType.IMAGE, info=info, + content = MediaMessageEventContent(msgtype=MessageType.IMAGE, info=info, body=name, relates_to=relates_to, external_url=self._get_external_url(evt)) + if file.decryption_info: + content.file = file.decryption_info + else: + content.url = file.mxc result = await self._send_message(intent, content, timestamp=evt.date) if evt.message: caption_content = await formatter.telegram_to_matrix(evt, source, self.main_intent, @@ -153,13 +158,20 @@ class PortalTelegram(BasePortal, ABC): info.width, info.height = attrs.width, attrs.height if file.thumbnail: - info.thumbnail_url = file.thumbnail.mxc + if file.thumbnail.decryption_info: + info.thumbnail_file = file.thumbnail.decryption_info + else: + info.thumbnail_url = file.thumbnail.mxc info.thumbnail_info = ThumbnailInfo(mimetype=file.thumbnail.mime_type, height=file.thumbnail.height or thumb_size.h, width=file.thumbnail.width or thumb_size.w, size=file.thumbnail.size) else: - info.thumbnail_url = file.mxc + # This is a hack for bad clients like Riot iOS that require a thumbnail + if file.decryption_info: + info.thumbnail_file = file.decryption_info + else: + info.thumbnail_url = file.mxc info.thumbnail_info = ImageInfo.deserialize(info.serialize()) return info, name @@ -186,7 +198,8 @@ class PortalTelegram(BasePortal, ABC): file = await util.transfer_file_to_matrix(source.client, intent, document, thumb_loc, is_sticker=attrs.is_sticker, tgs_convert=config["bridge.animated_sticker"], - filename=attrs.name, parallel_id=parallel_id) + filename=attrs.name, parallel_id=parallel_id, + encrypt=self.encrypted) if not file: return None @@ -199,13 +212,17 @@ class PortalTelegram(BasePortal, ABC): if attrs.is_sticker and file.mime_type.startswith("image/"): event_type = EventType.STICKER content = MediaMessageEventContent( - body=name or "unnamed file", info=info, url=file.mxc, relates_to=relates_to, + body=name or "unnamed file", info=info, relates_to=relates_to, external_url=self._get_external_url(evt), msgtype={ "video/": MessageType.VIDEO, "audio/": MessageType.AUDIO, "image/": MessageType.IMAGE, }.get(info.mimetype[:6], MessageType.FILE)) + if file.decryption_info: + content.file = file.decryption_info + else: + content.url = file.mxc return await self._send_message(intent, content, event_type=event_type, timestamp=evt.date) def handle_telegram_location(self, _: 'AbstractUser', intent: IntentAPI, evt: Message, diff --git a/mautrix_telegram/util/file_transfer.py b/mautrix_telegram/util/file_transfer.py index f9df31e5..36ae6a0c 100644 --- a/mautrix_telegram/util/file_transfer.py +++ b/mautrix_telegram/util/file_transfer.py @@ -29,11 +29,13 @@ from telethon.errors import (AuthBytesInvalidError, AuthKeyInvalidError, Locatio SecurityError, FileIdInvalidError) from mautrix.appservice import IntentAPI +from mautrix.types import EncryptedFile from ..tgclient import MautrixTelegramClient from ..db import TelegramFile as DBTelegramFile from ..util import sane_mimetypes from .parallel_file_transfer import parallel_transfer_to_matrix +from .tgs_converter import convert_tgs_to try: from PIL import Image @@ -49,7 +51,10 @@ try: except ImportError: VideoFileClip = random = string = os = mimetypes = None -from .tgs_converter import convert_tgs_to +try: + from nio.crypto import encrypt_attachment +except ImportError: + encrypt_attachment = None log: logging.Logger = logging.getLogger("mau.util") @@ -115,8 +120,8 @@ def _location_to_id(location: TypeLocation) -> str: async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: IntentAPI, - thumbnail_loc: TypeLocation, video: bytes, - mime: str) -> Optional[DBTelegramFile]: + thumbnail_loc: TypeLocation, video: bytes, mime: str, + encrypt: bool) -> Optional[DBTelegramFile]: if not Image or not VideoFileClip: return None @@ -140,11 +145,19 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In width, height = None, None mime_type = magic.from_buffer(file, mime=True) - content_uri = await intent.upload_media(file, mime_type) + decryption_info = None + upload_mime_type = mime_type + if encrypt: + file, decryption_info_dict = encrypt_attachment(file) + decryption_info = EncryptedFile.deserialize(decryption_info_dict) + upload_mime_type = "application/octet-stream" + content_uri = await intent.upload_media(file, upload_mime_type) + if decryption_info: + decryption_info.url = content_uri db_file = DBTelegramFile(id=loc_id, mxc=content_uri, mime_type=mime_type, was_converted=False, timestamp=int(time.time()), size=len(file), - width=width, height=height) + width=width, height=height, decryption_info=decryption_info) try: db_file.insert() except (IntegrityError, InvalidRequestError) as e: @@ -160,10 +173,10 @@ TypeThumbnail = Optional[Union[TypeLocation, TypePhotoSize]] async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI, - location: TypeLocation, thumbnail: TypeThumbnail = None, + location: TypeLocation, thumbnail: TypeThumbnail = None, *, is_sticker: bool = False, tgs_convert: Optional[dict] = None, - filename: Optional[str] = None, parallel_id: Optional[int] = None - ) -> Optional[DBTelegramFile]: + filename: Optional[str] = None, encrypt: bool = False, + parallel_id: Optional[int] = None) -> Optional[DBTelegramFile]: location_id = _location_to_id(location) if not location_id: return None @@ -180,14 +193,14 @@ async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentA async with lock: return await _unlocked_transfer_file_to_matrix(client, intent, location_id, location, thumbnail, is_sticker, tgs_convert, - filename, parallel_id) + filename, encrypt, parallel_id) async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI, loc_id: str, location: TypeLocation, thumbnail: TypeThumbnail, is_sticker: bool, tgs_convert: Optional[dict], filename: Optional[str], - parallel_id: Optional[int] + encrypt: bool, parallel_id: Optional[int] ) -> Optional[DBTelegramFile]: db_file = DBTelegramFile.get(loc_id) if db_file: @@ -195,7 +208,7 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten if parallel_id and isinstance(location, Document) and (not is_sticker or not tgs_convert): db_file = await parallel_transfer_to_matrix(client, intent, loc_id, location, filename, - parallel_id) + encrypt, parallel_id) mime_type = location.mime_type file = None else: @@ -228,9 +241,17 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten mime_type = new_mime_type thumbnail = None - content_uri = await intent.upload_media(file, mime_type) + decryption_info = None + upload_mime_type = mime_type + if encrypt and encrypt_attachment: + file, decryption_info_dict = encrypt_attachment(file) + decryption_info = EncryptedFile.deserialize(decryption_info_dict) + upload_mime_type = "application/octet-stream" + content_uri = await intent.upload_media(file, upload_mime_type) + if decryption_info: + decryption_info.url = content_uri - db_file = DBTelegramFile(id=loc_id, mxc=content_uri, + db_file = DBTelegramFile(id=loc_id, mxc=content_uri, decryption_info=decryption_info, mime_type=mime_type, was_converted=image_converted, timestamp=int(time.time()), size=len(file), width=width, height=height) @@ -239,7 +260,7 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten thumbnail = thumbnail.location try: db_file.thumbnail = await transfer_thumbnail_to_matrix(client, intent, thumbnail, file, - mime_type) + mime_type, encrypt) except FileIdInvalidError: log.warning(f"Failed to transfer thumbnail for {thumbnail!s}", exc_info=True) diff --git a/mautrix_telegram/util/parallel_file_transfer.py b/mautrix_telegram/util/parallel_file_transfer.py index 507646aa..5d316e48 100644 --- a/mautrix_telegram/util/parallel_file_transfer.py +++ b/mautrix_telegram/util/parallel_file_transfer.py @@ -34,11 +34,16 @@ from telethon.crypto import AuthKey from telethon import utils, helpers from mautrix.appservice import IntentAPI -from mautrix.types import ContentURI +from mautrix.types import ContentURI, EncryptedFile from ..tgclient import MautrixTelegramClient from ..db import TelegramFile as DBTelegramFile +try: + from nio.crypto import async_encrypt_attachment +except ImportError: + async_encrypt_attachment = None + log: logging.Logger = logging.getLogger("mau.util") TypeLocation = Union[Document, InputDocumentFileLocation, InputPeerPhotoFileLocation, @@ -242,18 +247,34 @@ parallel_transfer_locks: DefaultDict[int, asyncio.Lock] = defaultdict(lambda: as async def parallel_transfer_to_matrix(client: MautrixTelegramClient, intent: IntentAPI, loc_id: str, location: TypeLocation, filename: str, - parallel_id: int) -> DBTelegramFile: + encrypt: bool, parallel_id: int) -> DBTelegramFile: size = location.size mime_type = location.mime_type dc_id, location = utils.get_input_location(location) # We lock the transfers because telegram has connection count limits async with parallel_transfer_locks[parallel_id]: downloader = ParallelTransferrer(client, dc_id) - content_uri = await intent.upload_media(downloader.download(location, size), - mime_type=mime_type, filename=filename, size=size) + data = downloader.download(location, size) + decryption_info = None + up_mime_type = mime_type + if encrypt and async_encrypt_attachment: + async def encrypted(stream): + nonlocal decryption_info + async for chunk in async_encrypt_attachment(stream): + if isinstance(chunk, dict): + decryption_info = EncryptedFile.deserialize(chunk) + else: + yield chunk + + data = encrypted(data) + up_mime_type = "application/octet-stream" + content_uri = await intent.upload_media(data, mime_type=up_mime_type, filename=filename, + size=size if not encrypt else None) + if decryption_info: + decryption_info.url = content_uri return DBTelegramFile(id=loc_id, mxc=content_uri, mime_type=mime_type, was_converted=False, timestamp=int(time.time()), size=size, - width=None, height=None) + width=None, height=None, decryption_info=decryption_info) async def _internal_transfer_to_telegram(client: MautrixTelegramClient, response: ClientResponse