iartamonov/mautrix-telegram
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Encrypt media being sent to Matrix in encrypted rooms

Browse Source
This commit is contained in:
Tulir Asokan
2020-03-30 21:29:52 +03:00
parent af285c5ffe
commit 698b56afcf
6 changed files with 140 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mautrix_telegram/util/file_transfer.py
+35 -14
View File
@@ -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)
mautrix_telegram/util/parallel_file_transfer.py
+26 -5
View File
@@ -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