diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index 219f20cd..2bfe449b 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -14,7 +14,6 @@ # # You should have received a copy of the GNU General Public License # along with this program. If not, see . -from io import BytesIO from collections import deque from datetime import datetime import asyncio @@ -23,9 +22,7 @@ import mimetypes import hashlib import logging -from PIL import Image import magic -import time from telethon.tl.functions.messages import * from telethon.tl.functions.channels import * @@ -33,8 +30,8 @@ from telethon.errors.rpc_error_list import * from telethon.tl.types import * from mautrix_appservice import MatrixRequestError, IntentError -from .db import Portal as DBPortal, Message as DBMessage, TelegramFile as DBTelegramFile -from . import puppet as p, user as u, formatter +from .db import Portal as DBPortal, Message as DBMessage +from . import puppet as p, user as u, formatter, util mimetypes.init() @@ -395,14 +392,12 @@ class Portal: async def update_avatar(self, user, photo): photo_id = f"{photo.volume_id}-{photo.local_id}" if self.photo_id != photo_id: - try: - file = await user.client.download_file_bytes(photo) - except LocationInvalidError: - return False - uploaded = await self.main_intent.upload_file(file) - await self.main_intent.set_room_avatar(self.mxid, uploaded["content_uri"]) - self.photo_id = photo_id - return True + file = await util.transfer_file_to_matrix(self.db, user.client, self.main_intent, + photo) + if file: + await self.main_intent.set_avatar(file.mxc) + self.photo_id = photo_id + return True return False async def _get_users(self, user, entity): @@ -764,7 +759,8 @@ class Portal: async def handle_telegram_photo(self, source, intent, media, relates_to=None): largest_size = self._get_largest_photo_size(media.photo) - file = await self.transfer_file_to_matrix(source.client, intent, largest_size.location) + file = await util.transfer_file_to_matrix(self.db, source.client, intent, + largest_size.location) if not file: return None info = { @@ -780,51 +776,8 @@ class Portal: return await intent.send_image(self.mxid, file.mxc, info=info, text=name, relates_to=relates_to) - def convert_webp(self, file, to="png"): - try: - image = Image.open(BytesIO(file)).convert("RGBA") - new_file = BytesIO() - image.save(new_file, to) - return f"image/{to}", new_file.getvalue() - except Exception: - self.log.exception(f"Failed to convert webp to {to}") - return "image/webp", file - - async def transfer_file_to_matrix(self, client, intent, location): - if isinstance(location, (Document, InputDocumentFileLocation)): - id = f"{location.id}-{location.version}" - elif not isinstance(location, (FileLocation, InputFileLocation)): - id = f"{location.volume_id}-{location.local_id}" - else: - return None - - db_file = DBTelegramFile.query.get(id) - if db_file: - return db_file - - try: - file = await client.download_file_bytes(location) - except LocationInvalidError: - return None - mime_type = magic.from_buffer(file, mime=True) - - image_converted = False - if mime_type == "image/webp": - mime_type, file = self.convert_webp(file, to="png") - image_converted = True - - uploaded = await intent.upload_file(file, mime_type) - - db_file = DBTelegramFile(id=id, mxc=uploaded["content_uri"], - mime_type=mime_type, was_converted=image_converted, - timestamp=int(time.time())) - self.db.add(db_file) - self.db.commit() - - return db_file - async def handle_telegram_document(self, source, intent, media, relates_to=None): - file = await self.transfer_file_to_matrix(source.client, intent, media.document) + file = await util.transfer_file_to_matrix(self.db, source.client, intent, media.document) if not file: return None name = media.caption diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index d41c9e82..a316a48d 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -18,10 +18,11 @@ from difflib import SequenceMatcher import re import logging -from telethon.tl.types import UserProfilePhoto, PeerUser +from telethon.tl.types import UserProfilePhoto from telethon.errors.rpc_error_list import LocationInvalidError from .db import Puppet as DBPuppet +from . import util config = None @@ -129,14 +130,11 @@ class Puppet: async def update_avatar(self, source, photo): photo_id = f"{photo.volume_id}-{photo.local_id}" if self.photo_id != photo_id: - try: - file = await source.client.download_file_bytes(photo) - except LocationInvalidError: - return False - uploaded = await self.intent.upload_file(file) - await self.intent.set_avatar(uploaded["content_uri"]) - self.photo_id = photo_id - return True + file = await util.transfer_file_to_matrix(self.db, source.client, self.intent, photo) + if file: + await self.intent.set_avatar(file.mxc) + self.photo_id = photo_id + return True return False @classmethod diff --git a/mautrix_telegram/util/__init__.py b/mautrix_telegram/util/__init__.py new file mode 100644 index 00000000..53d6aabd --- /dev/null +++ b/mautrix_telegram/util/__init__.py @@ -0,0 +1 @@ +from .file_transfer import transfer_file_to_matrix diff --git a/mautrix_telegram/util/file_transfer.py b/mautrix_telegram/util/file_transfer.py new file mode 100644 index 00000000..698cb874 --- /dev/null +++ b/mautrix_telegram/util/file_transfer.py @@ -0,0 +1,75 @@ +# -*- coding: future_fstrings -*- +# mautrix-telegram - A Matrix-Telegram puppeting bridge +# Copyright (C) 2018 Tulir Asokan +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU 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 General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +from io import BytesIO +import time +import logging + +import magic +from PIL import Image + +from telethon.tl.types import (Document, FileLocation, InputFileLocation, + InputDocumentFileLocation) +from telethon.errors import LocationInvalidError + +from ..db import TelegramFile as DBTelegramFile + +log = logging.getLogger("mau.util") + + +def _convert_webp(file, to="png"): + try: + image = Image.open(BytesIO(file)).convert("RGBA") + new_file = BytesIO() + image.save(new_file, to) + return f"image/{to}", new_file.getvalue() + except Exception: + log.exception(f"Failed to convert webp to {to}") + return "image/webp", file + + +async def transfer_file_to_matrix(db, client, intent, location): + if isinstance(location, (Document, InputDocumentFileLocation)): + id = f"{location.id}-{location.version}" + elif not isinstance(location, (FileLocation, InputFileLocation)): + id = f"{location.volume_id}-{location.local_id}" + else: + return None + + db_file = DBTelegramFile.query.get(id) + if db_file: + return db_file + + try: + file = await client.download_file_bytes(location) + except LocationInvalidError: + return None + mime_type = magic.from_buffer(file, mime=True) + + image_converted = False + if mime_type == "image/webp": + mime_type, file = _convert_webp(file, to="png") + image_converted = True + + uploaded = await intent.upload_file(file, mime_type) + + db_file = DBTelegramFile(id=id, mxc=uploaded["content_uri"], + mime_type=mime_type, was_converted=image_converted, + timestamp=int(time.time())) + db.add(db_file) + db.commit() + + return db_file