Move TelegramClient extensions to own class

This commit is contained in:
Tulir Asokan
2018-02-05 17:08:00 +02:00
parent c139e24bf4
commit 11eb9277ed
4 changed files with 97 additions and 75 deletions
+8 -7
View File
@@ -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":
+1 -1
View File
@@ -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)
+83
View File
@@ -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
+5 -67
View File
@@ -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,
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: