uploads: Add support for encrypted uploads.

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

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,

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 "
+            assert upload.encrypt
-                              "not yet implemented")
-            return False
 
-        # TODO the content is different if the room is encrypted.
+        content = upload.content
-        content = {
-            "msgtype": Api.mimetype_to_msgtype(upload.mimetype),
-            "body": upload.file_name,
-            "url": upload.content_uri,
-        }
 
         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
+            message = EncrytpionQueueItem(upload.msgtype, upload)
-            # are shared
+            self.encryption_queue[upload.room_id].append(message)
-            # message = EncrytpionQueueItem(msgtype, formatted)
-            # self.encryption_queue[room.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,
+                    if item.message_type in [
-                                                  message.formatted_message,
+                        "m.file",
-                                                  message.message_type):
+                        "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

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)