uploads: Add support for encrypted uploads.
This commit is contained in:
Damir Jelić
parent
6a1901c0ac
commit
01a23e8221
4 changed files with 151 additions and 32 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue