From 11eb9277ed1dd2c822929b46ce82aeb678c26758 Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Mon, 5 Feb 2018 17:08:00 +0200 Subject: [PATCH] Move TelegramClient extensions to own class --- mautrix_telegram/portal.py | 15 ++++--- mautrix_telegram/puppet.py | 2 +- mautrix_telegram/tgclient.py | 83 ++++++++++++++++++++++++++++++++++++ mautrix_telegram/user.py | 72 +++---------------------------- 4 files changed, 97 insertions(+), 75 deletions(-) create mode 100644 mautrix_telegram/tgclient.py diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index ca9976fa..3ff126c0 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -252,7 +252,7 @@ class Portal: photo_id = f"{photo.volume_id}-{photo.local_id}" if self.photo_id != photo_id: try: - file = user.download_file(photo) + file = user.client.download_file_bytes(photo) except LocationInvalidError: return False uploaded = self.main_intent.upload_file(file) @@ -340,12 +340,12 @@ class Portal: reply_to = None if len(entities) > 0 and isinstance(entities[0], formatter.MessageEntityReply): reply_to = entities.pop(0).msg_id - response = sender.send_message(self.peer, message, entities=entities, - reply_to=reply_to) + response = sender.client.send_message(self.peer, message, entities=entities, + reply_to=reply_to) else: if type == "m.emote": message["body"] = "/me " + message["body"] - response = sender.send_message(self.peer, message["body"]) + response = sender.client.send_message(self.peer, message["body"]) elif type in {"m.image", "m.file", "m.audio", "m.video"}: file = self.main_intent.download_file(message["url"]) @@ -358,7 +358,8 @@ class Portal: if "w" in info and "h" in info: attributes.append(DocumentAttributeImageSize(w=info["w"], h=info["h"])) - response = sender.send_file(self.peer, file, mime, caption, attributes, file_name) + response = sender.client.send_file(self.peer, file, mime, caption, attributes, + file_name) else: self.log.debug("Unhandled Matrix event: %s", message) return @@ -528,7 +529,7 @@ class Portal: def handle_telegram_photo(self, source, sender, media): largest_size = self._get_largest_photo_size(media.photo) - file = source.download_file(largest_size.location) + file = source.client.download_file_bytes(largest_size.location) mime_type = magic.from_buffer(file, mime=True) uploaded = sender.intent.upload_file(file, mime_type) info = { @@ -554,7 +555,7 @@ class Portal: return "image/webp", file def handle_telegram_document(self, source, sender, media): - file = source.download_file(media.document) + file = source.client.download_file_bytes(media.document) mime_type = magic.from_buffer(file, mime=True) dont_change_mime = False if mime_type == "image/webp": diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index 734ea435..69251f4b 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -112,7 +112,7 @@ class Puppet: photo_id = f"{photo.volume_id}-{photo.local_id}" if self.photo_id != photo_id: try: - file = source.download_file(photo) + file = source.client.download_file_bytes(photo) except LocationInvalidError: return False uploaded = self.intent.upload_file(file) diff --git a/mautrix_telegram/tgclient.py b/mautrix_telegram/tgclient.py new file mode 100644 index 00000000..715b3b24 --- /dev/null +++ b/mautrix_telegram/tgclient.py @@ -0,0 +1,83 @@ +# -*- 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 +from telethon import TelegramClient +from telethon.tl.functions.messages import SendMessageRequest, SendMediaRequest +from telethon.tl.types import * + + +class MautrixTelegramClient(TelegramClient): + def send_message(self, entity, message, reply_to=None, entities=None, link_preview=True): + entity = self.get_input_entity(entity) + + request = SendMessageRequest( + peer=entity, + message=message, + entities=entities, + no_webpage=not link_preview, + reply_to_msg_id=self._get_reply_to(reply_to) + ) + result = self(request) + if isinstance(result, UpdateShortSentMessage): + return Message( + id=result.id, + to_id=entity, + message=message, + date=result.date, + out=result.out, + media=result.media, + entities=result.entities + ) + + return self._get_response_message(request, result) + + def send_file(self, entity, file, mime_type=None, caption=None, attributes=None, file_name=None, + reply_to=None, **kwargs): + entity = self.get_input_entity(entity) + reply_to = self._get_reply_to(reply_to) + + file_handle = self.upload_file(file, file_name=file_name, use_cache=False) + + if mime_type == "image/png": + media = InputMediaUploadedPhoto(file_handle, caption or "") + else: + attributes = attributes or [] + attr_dict = {type(attr): attr for attr in attributes} + + media = InputMediaUploadedDocument( + file=file_handle, + mime_type=mime_type or "application/octet-stream", + attributes=list(attr_dict.values()), + caption=caption or "") + + request = SendMediaRequest(entity, media, reply_to_msg_id=reply_to) + return self._get_response_message(request, self(request)) + + def download_file_bytes(self, location): + if isinstance(location, Document): + location = InputDocumentFileLocation(location.id, location.access_hash, + location.version) + elif not isinstance(location, (InputFileLocation, InputDocumentFileLocation)): + location = InputFileLocation(location.volume_id, location.local_id, location.secret) + + file = BytesIO() + + self.download_file(location, file) + + data = file.getvalue() + file.close() + return data diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py index c0add044..644557d3 100644 --- a/mautrix_telegram/user.py +++ b/mautrix_telegram/user.py @@ -14,12 +14,10 @@ # # You should have received a copy of the GNU General Public License # along with this program. If not, see . -from io import BytesIO -from telethon import TelegramClient from telethon.tl.types import * from telethon.tl.types import User as TLUser -from telethon.tl.functions.messages import SendMessageRequest, SendMediaRequest from .db import User as DBUser +from .tgclient import MautrixTelegramClient from . import portal as po, puppet as pu config = None @@ -81,10 +79,10 @@ class User: # region Telegram connection management def start(self): - self.client = TelegramClient(self.mxid, - config["telegram.api_id"], - config["telegram.api_hash"], - update_workers=2) + self.client = MautrixTelegramClient(self.mxid, + config["telegram.api_id"], + config["telegram.api_hash"], + update_workers=2) self.connected = self.client.connect() if self.logged_in: self.post_login() @@ -126,66 +124,6 @@ class User: self.save() return self.client.log_out() - def send_message(self, entity, message, reply_to=None, entities=None, link_preview=True): - entity = self.client.get_input_entity(entity) - - request = SendMessageRequest( - peer=entity, - message=message, - entities=entities, - no_webpage=not link_preview, - reply_to_msg_id=self.client._get_reply_to(reply_to) - ) - result = self.client(request) - if isinstance(result, UpdateShortSentMessage): - return Message( - id=result.id, - to_id=entity, - message=message, - date=result.date, - out=result.out, - media=result.media, - entities=result.entities - ) - - return self.client._get_response_message(request, result) - - def send_file(self, entity, file, mime_type=None, caption=None, attributes=[], file_name=None, - reply_to=None): - entity = self.client.get_input_entity(entity) - reply_to = self.client._get_reply_to(reply_to) - - file_handle = self.client.upload_file(file, file_name=file_name, use_cache=False) - - if mime_type == "image/png": - media = InputMediaUploadedPhoto(file_handle, caption or "") - else: - attr_dict = {type(attr): attr for attr in attributes} - - media = InputMediaUploadedDocument( - file=file_handle, - mime_type=mime_type or "application/octet-stream", - attributes=list(attr_dict.values()), - caption=caption or "") - - request = SendMediaRequest(entity, media, reply_to_msg_id=reply_to) - return self.client._get_response_message(request, self.client(request)) - - def download_file(self, location): - if isinstance(location, Document): - location = InputDocumentFileLocation(location.id, location.access_hash, - location.version) - elif not isinstance(location, (InputFileLocation, InputDocumentFileLocation)): - location = InputFileLocation(location.volume_id, location.local_id, location.secret) - - file = BytesIO() - - self.client.download_file(location, file) - - data = file.getvalue() - file.close() - return data - def sync_dialogs(self): dialogs = self.client.get_dialogs(limit=30) for dialog in dialogs: