Move TelegramClient extensions to own class

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,
+                response = sender.client.send_message(self.peer, message, entities=entities,
-                                               reply_to=reply_to)
+                                                      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":

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)

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 <http://www.gnu.org/licenses/>.
+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

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 <http://www.gnu.org/licenses/>.
-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,
+        self.client = MautrixTelegramClient(self.mxid,
-                                     config["telegram.api_id"],
+                                            config["telegram.api_id"],
-                                     config["telegram.api_hash"],
+                                            config["telegram.api_hash"],
-                                     update_workers=2)
+                                            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: