From 01a23e82219dfe21335e2b36ade166ede9ca5f14 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Damir=20Jeli=C4=87?= Date: Sun, 17 Feb 2019 22:45:49 +0100 Subject: [PATCH] uploads: Add support for encrypted uploads. --- contrib/matrix_upload | 59 ++++++++++++++++++++++++++++++++++++++++-- matrix/commands.py | 5 ---- matrix/server.py | 60 +++++++++++++++++++++++++------------------ matrix/uploads.py | 59 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 151 insertions(+), 32 deletions(-) diff --git a/contrib/matrix_upload b/contrib/matrix_upload index 63ebe19..94396bc 100755 --- a/contrib/matrix_upload +++ b/contrib/matrix_upload @@ -20,9 +20,12 @@ import magic import requests import argparse from urllib.parse import urlparse +from itertools import zip_longest import urllib3 from nio import Api, UploadResponse, UploadError +from nio.crypto import encrypt_attachment + from json.decoder import JSONDecodeError urllib3.disable_warnings() @@ -95,6 +98,41 @@ class Upload(object): return self.totalsize +def chunk_bytes(iterable, n): + args = [iter(iterable)] * n + return ( + bytes( + (filter(lambda x: x is not None, chunk)) + ) for chunk in zip_longest(*args) + ) + + +class EncryptedUpload(Upload): + def __init__(self, file, chunksize=1 << 13): + super().__init__(file, chunksize) + self.source_mimetype = self.mimetype + self.mimetype = "application/octet-stream" + + with open(self.filename, "rb") as file: + self.ciphertext, self.file_keys = encrypt_attachment(file.read()) + + def send_progress(self): + message = { + "type": "progress", + "data": self.readsofar + } + to_stdout(message) + + def __iter__(self): + for chunk in chunk_bytes(self.ciphertext, self.chunksize): + self.readsofar += len(chunk) + self.send_progress() + yield chunk + + def __len__(self): + return len(self.ciphertext) + + class IterableToFileAdapter(object): def __init__(self, iterable): self.iterator = iter(iterable) @@ -109,9 +147,18 @@ class IterableToFileAdapter(object): def upload_process(args): file_path = os.path.expanduser(args.file) + thumbnail = None try: - upload = Upload(file_path, 10) + if args.encrypt: + upload = EncryptedUpload(file_path) + + if upload.source_mimetype.startswith("image"): + # TODO create a thumbnail + thumbnail = None + else: + upload = Upload(file_path) + except (FileNotFoundError, OSError, IOError) as e: error(e) @@ -153,9 +200,14 @@ def upload_process(args): "type": "status", "status": "started", "total": upload.totalsize, - "mimetype": upload.mimetype, "file_name": upload.filename, } + + if isinstance(upload, EncryptedUpload): + message["mimetype"] = upload.source_mimetype + else: + message["mimetype"] = upload.mimetype + to_stdout(message) session = requests.Session() @@ -189,6 +241,9 @@ def upload_process(args): "url": response.content_uri } + if isinstance(upload, EncryptedUpload): + message["file_keys"] = upload.file_keys + to_stdout(message) return 0 diff --git a/matrix/commands.py b/matrix/commands.py index 615a58e..56cfe71 100644 --- a/matrix/commands.py +++ b/matrix/commands.py @@ -1071,11 +1071,6 @@ def matrix_upload_command_cb(data, buffer, args): if not room_buffer: continue - if room_buffer.room.encrypted: - room_buffer.error("Uploading to encrypted rooms is " - "not yet implemented") - return W.WEECHAT_RC_OK - upload = Upload( server.name, server.config.address, diff --git a/matrix/server.py b/matrix/server.py index e6fab0c..5e2b061 100644 --- a/matrix/server.py +++ b/matrix/server.py @@ -24,7 +24,16 @@ import time import copy from collections import defaultdict, deque from atomicwrites import atomic_write -from typing import Any, Deque, Dict, Optional, List, NamedTuple, DefaultDict +from typing import ( + Any, + Deque, + Dict, + Optional, + List, + NamedTuple, + DefaultDict, + Union +) from uuid import UUID @@ -64,6 +73,7 @@ from .config import ConfigSection, Option, ServerBufferType from .globals import SCRIPT_NAME, SERVERS, W, MAX_EVENTS, TYPING_NOTICE_TIMEOUT from .utf import utf8_decode from .utils import create_server_buffer, key_from_value, server_buffer_prnt +from .uploads import Upload from .colors import Formatted, FormattedString, DEFAULT_ATTRIBUTES @@ -82,7 +92,7 @@ EncrytpionQueueItem = NamedTuple( "EncrytpionQueueItem", [ ("message_type", str), - ("formatted_message", Formatted), + ("message", Union[Formatted, Upload]), ], ) @@ -797,38 +807,25 @@ class MatrixServer(object): try: room_buffer = self.find_room_from_id(upload.room_id) except (ValueError, KeyError): - return False + return True assert self.client if room_buffer.room.encrypted: - room_buffer.error("Uploading to encrypted rooms is " - "not yet implemented") - return False + assert upload.encrypt - # TODO the content is different if the room is encrypted. - content = { - "msgtype": Api.mimetype_to_msgtype(upload.mimetype), - "body": upload.file_name, - "url": upload.content_uri, - } + content = upload.content try: uuid = self.room_send_event(upload.room_id, content) except (EncryptionError, GroupEncryptionError): - # TODO put the message in a queue to resend after group sessions - # are shared - # message = EncrytpionQueueItem(msgtype, formatted) - # self.encryption_queue[room.room_id].append(message) + message = EncrytpionQueueItem(upload.msgtype, upload) + self.encryption_queue[upload.room_id].append(message) return False - http_url = Api.mxc_to_http(upload.content_uri) - description = ("/{}".format(upload.file_name) if upload.file_name - else "") - attributes = DEFAULT_ATTRIBUTES.copy() formatted = Formatted([FormattedString( - "{url}{desc}".format(url=http_url, desc=description), + upload.render, attributes )]) @@ -1386,14 +1383,27 @@ class MatrixServer(object): room_buffer = self.room_buffers[room_id] while self.encryption_queue[room_id]: - message = self.encryption_queue[room_id].popleft() + item = self.encryption_queue[room_id].popleft() try: - if not self.room_send_message(room_buffer, - message.formatted_message, - message.message_type): + if item.message_type in [ + "m.file", + "m.video", + "m.audio", + "m.image" + ]: + ret = self.room_send_upload(item.message) + else: + ret = self.room_send_message( + room_buffer, + item.message, + item.message_type + ) + + if not ret: self.encryption_queue[room_id].pop() self.encryption_queue[room_id].appendleft(message) break + except OlmTrustError: self.encryption_queue[room_id].clear() break diff --git a/matrix/uploads.py b/matrix/uploads.py index a7b4625..a88a91b 100644 --- a/matrix/uploads.py +++ b/matrix/uploads.py @@ -21,6 +21,7 @@ from __future__ import unicode_literals import attr import time import json +from typing import Dict, Any from uuid import uuid1, UUID from enum import Enum @@ -32,6 +33,7 @@ except ImportError: from .globals import SCRIPT_NAME, SERVERS, W, UPLOADS from .utf import utf8_decode from matrix import globals as G +from nio import Api class UploadState(Enum): @@ -168,6 +170,62 @@ class Upload(object): def abort(self): pass + @property + def msgtype(self): + # type: () -> str + assert self.mimetype + return Api.mimetype_to_msgtype(self.mimetype) + + @property + def content(self): + # type: () -> Dict[Any, Any] + assert self.content_uri + + if self.encrypt: + content = { + "body": self.file_name, + "msgtype": self.msgtype, + "file": self.file_keys, + } + content["file"]["url"] = self.content_uri + content["file"]["mimetype"] = self.mimetype + + # TODO thumbnail if it's an image + + return content + + return { + "msgtype": self.msgtype, + "body": self.file_name, + "url": self.content_uri, + } + + @property + def render(self): + # type: () -> str + assert self.content_uri + + if self.encrypt: + http_url = Api.encrypted_mxc_to_plumb( + self.content_uri, + self.file_keys["key"]["k"], + self.file_keys["hashes"]["sha256"], + self.file_keys["iv"] + ) + url = http_url if http_url else self.content_uri + + description = "{}".format(self.file_name) + return ("{del_color}<{ncolor}{desc}{del_color}>{ncolor} " + "{del_color}[{ncolor}{url}{del_color}]{ncolor}").format( + del_color=W.color("chat_delimiters"), + ncolor=W.color("reset"), + desc=description, url=url) + + http_url = Api.mxc_to_http(self.content_uri) + description = ("/{}".format(self.file_name) if self.file_name + else "") + return "{url}{desc}".format(url=http_url, desc=description) + @attr.s class UploadsBuffer(object): @@ -293,6 +351,7 @@ def handle_child_message(upload, message): elif message["status"] == "done": upload.state = UploadState.finished upload.content_uri = message["url"] + upload.file_keys = message.get("file_keys", None) server = SERVERS.get(upload.server_name, None)