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
alembic/versions/d3c922a6acd2_add_decryption_info_field_for_.py
+26
View File
@@ -0,0 +1,26 @@
"""Add decryption info field for reuploaded telegram files
Revision ID: d3c922a6acd2
Revises: 24f31fc8a72b
Create Date: 2020-03-30 20:07:17.340346
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = 'd3c922a6acd2'
down_revision = '24f31fc8a72b'
branch_labels = None
depends_on = None
def upgrade():
with op.batch_alter_table("telegram_file") as batch_op:
batch_op.add_column(sa.Column("decryption_info", sa.Text(), nullable=True))
def downgrade():
with op.batch_alter_table("telegram_file") as batch_op:
batch_op.drop_column("decryption_info")
mautrix_telegram/db/telegram_file.py
+28 -5
View File
@@ -13,15 +13,37 @@
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from typing import Optional
from typing import Optional, cast, Dict, Any
from sqlalchemy import Column, ForeignKey, Integer, BigInteger, String, Boolean
from sqlalchemy import (Column, ForeignKey, Integer, BigInteger, String, Boolean, Text,
TypeDecorator)
from sqlalchemy.engine.result import RowProxy
from mautrix.types import ContentURI
from mautrix.types import ContentURI, EncryptedFile
from mautrix.util.db import Base
class DBEncryptedFile(TypeDecorator):
impl = Text
@property
def python_type(self):
return EncryptedFile
def process_bind_param(self, value: EncryptedFile, dialect) -> Optional[str]:
if value is not None:
return value.json()
return None
def process_result_value(self, value: str, dialect) -> Optional[EncryptedFile]:
if value is not None:
return EncryptedFile.parse_json(value)
return None
def process_literal_param(self, value, dialect):
return value
class TelegramFile(Base):
__tablename__ = "telegram_file"
@@ -33,12 +55,13 @@ class TelegramFile(Base):
size: Optional[int] = Column(Integer, nullable=True)
width: Optional[int] = Column(Integer, nullable=True)
height: Optional[int] = Column(Integer, nullable=True)
decryption_info: Optional[Dict[str, Any]] = Column(DBEncryptedFile, nullable=True)
thumbnail_id: str = Column("thumbnail", String, ForeignKey("telegram_file.id"), nullable=True)
thumbnail: Optional['TelegramFile'] = None
@classmethod
def scan(cls, row: RowProxy) -> 'TelegramFile':
telegram_file: TelegramFile = super().scan(row)
telegram_file = cast(TelegramFile, super().scan(row))
if isinstance(telegram_file.thumbnail, str):
telegram_file.thumbnail = cls.get(telegram_file.thumbnail)
return telegram_file
@@ -52,5 +75,5 @@ class TelegramFile(Base):
conn.execute(self.t.insert().values(
id=self.id, mxc=self.mxc, mime_type=self.mime_type,
was_converted=self.was_converted, timestamp=self.timestamp, size=self.size,
width=self.width, height=self.height,
width=self.width, height=self.height, decryption_info=self.decryption_info,
thumbnail=self.thumbnail.id if self.thumbnail else self.thumbnail_id))
mautrix_telegram/matrix.py
+2 -1
View File
@@ -54,7 +54,8 @@ class MatrixHandler(BaseMatrixHandler):
self.user_id_suffix = f"{suffix}:{homeserver}"
super(MatrixHandler, self).__init__(context.az, context.config, loop=context.loop,
command_processor=com.CommandProcessor(context))
command_processor=com.CommandProcessor(context),
bridge=context.bridge)
self.bot = context.bot
self.previously_typing = {}
mautrix_telegram/portal/telegram.py
+23 -6
View File
@@ -80,7 +80,8 @@ class PortalTelegram(BasePortal, ABC):
async def handle_telegram_photo(self, source: 'AbstractUser', intent: IntentAPI, evt: Message,
relates_to: Dict = None) -> Optional[EventID]:
loc, largest_size = self._get_largest_photo_size(evt.media.photo)
file = await util.transfer_file_to_matrix(source.client, intent, loc)
file = await util.transfer_file_to_matrix(source.client, intent, loc,
encrypt=self.encrypted)
if not file:
return None
if self.get_config("inline_images") and (evt.message
@@ -98,9 +99,13 @@ class PortalTelegram(BasePortal, ABC):
else largest_size.size))
name = f"image{sane_mimetypes.guess_extension(file.mime_type)}"
await intent.set_typing(self.mxid, is_typing=False)
content = MediaMessageEventContent(url=file.mxc, msgtype=MessageType.IMAGE, info=info,
content = MediaMessageEventContent(msgtype=MessageType.IMAGE, info=info,
body=name, relates_to=relates_to,
external_url=self._get_external_url(evt))
if file.decryption_info:
content.file = file.decryption_info
else:
content.url = file.mxc
result = await self._send_message(intent, content, timestamp=evt.date)
if evt.message:
caption_content = await formatter.telegram_to_matrix(evt, source, self.main_intent,
@@ -153,13 +158,20 @@ class PortalTelegram(BasePortal, ABC):
info.width, info.height = attrs.width, attrs.height
if file.thumbnail:
info.thumbnail_url = file.thumbnail.mxc
if file.thumbnail.decryption_info:
info.thumbnail_file = file.thumbnail.decryption_info
else:
info.thumbnail_url = file.thumbnail.mxc
info.thumbnail_info = ThumbnailInfo(mimetype=file.thumbnail.mime_type,
height=file.thumbnail.height or thumb_size.h,
width=file.thumbnail.width or thumb_size.w,
size=file.thumbnail.size)
else:
info.thumbnail_url = file.mxc
# This is a hack for bad clients like Riot iOS that require a thumbnail
if file.decryption_info:
info.thumbnail_file = file.decryption_info
else:
info.thumbnail_url = file.mxc
info.thumbnail_info = ImageInfo.deserialize(info.serialize())
return info, name
@@ -186,7 +198,8 @@ class PortalTelegram(BasePortal, ABC):
file = await util.transfer_file_to_matrix(source.client, intent, document, thumb_loc,
is_sticker=attrs.is_sticker,
tgs_convert=config["bridge.animated_sticker"],
filename=attrs.name, parallel_id=parallel_id)
filename=attrs.name, parallel_id=parallel_id,
encrypt=self.encrypted)
if not file:
return None
@@ -199,13 +212,17 @@ class PortalTelegram(BasePortal, ABC):
if attrs.is_sticker and file.mime_type.startswith("image/"):
event_type = EventType.STICKER
content = MediaMessageEventContent(
body=name or "unnamed file", info=info, url=file.mxc, relates_to=relates_to,
body=name or "unnamed file", info=info, relates_to=relates_to,
external_url=self._get_external_url(evt),
msgtype={
"video/": MessageType.VIDEO,
"audio/": MessageType.AUDIO,
"image/": MessageType.IMAGE,
}.get(info.mimetype[:6], MessageType.FILE))
if file.decryption_info:
content.file = file.decryption_info
else:
content.url = file.mxc
return await self._send_message(intent, content, event_type=event_type, timestamp=evt.date)
def handle_telegram_location(self, _: 'AbstractUser', intent: IntentAPI, evt: Message,
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