Encrypt uploads in a thread

This commit is contained in:
miruka 2019-11-04 15:18:01 -04:00
parent 3c1d926188
commit 82104266fe
3 changed files with 27 additions and 12 deletions

View File

@ -1,9 +1,7 @@
- Media - Media
- Encrypt file for upload in thread
- Cache our own uploads - Cache our own uploads
- Uploading progress (+local echo) - Uploading progress bar (+local echo)
- Deduplicate uploads - Image loading progress bar
- Loading progress bar
- Downloading - Downloading
- Bottom/top bar - Bottom/top bar

View File

@ -1,4 +1,5 @@
import asyncio import asyncio
import functools
import html import html
import inspect import inspect
import io import io
@ -12,20 +13,24 @@ from datetime import datetime
from functools import partial from functools import partial
from pathlib import Path from pathlib import Path
from types import ModuleType from types import ModuleType
from typing import Any, DefaultDict, Dict, Optional, Set, Tuple, Type, Union from typing import (
Any, BinaryIO, DefaultDict, Dict, Optional, Set, Tuple, Type, Union
)
from uuid import uuid4 from uuid import uuid4
import nio
from nio.crypto.attachments import encrypt_attachment
from PIL import Image as PILImage from PIL import Image as PILImage
from pymediainfo import MediaInfo from pymediainfo import MediaInfo
import nio
from . import __about__, utils from . import __about__, utils
from .html_filter import HTML_FILTER from .html_filter import HTML_FILTER
from .models.items import Account, Event, Member, Room, TypeSpecifier from .models.items import Account, Event, Member, Room, TypeSpecifier
from .models.model_store import ModelStore from .models.model_store import ModelStore
from .pyotherside_events import AlertRequested from .pyotherside_events import AlertRequested
CryptDict = Dict[str, Any]
@dataclass @dataclass
class UploadError(Exception): class UploadError(Exception):
@ -419,9 +424,19 @@ class MatrixClient(nio.AsyncClient):
self.models.pop((Member, room_id), None) self.models.pop((Member, room_id), None)
async def encrypt_attachment(self, data: bytes) -> Tuple[bytes, CryptDict]:
func = functools.partial(
nio.crypto.attachments.encrypt_attachment,
data,
)
# Run in a separate thread
return await asyncio.get_event_loop().run_in_executor(None, func)
async def upload_thumbnail( async def upload_thumbnail(
self, path: Union[Path, str], encrypt: bool = False, self, path: Union[Path, str], encrypt: bool = False,
) -> Tuple[str, Dict[str, Any], Dict[str, Any]]: ) -> Tuple[str, Dict[str, Any], CryptDict]:
png_modes = ("1", "L", "P", "RGBA") png_modes = ("1", "L", "P", "RGBA")
@ -449,7 +464,7 @@ class MatrixClient(nio.AsyncClient):
data = out.getvalue() data = out.getvalue()
if encrypt: if encrypt:
data, crypt_dict = encrypt_attachment(data) data, crypt_dict = await self.encrypt_attachment(data)
upload_mime = "application/octet-stream" upload_mime = "application/octet-stream"
else: else:
crypt_dict, upload_mime = {}, mime crypt_dict, upload_mime = {}, mime
@ -471,13 +486,15 @@ class MatrixClient(nio.AsyncClient):
async def upload_file(self, path: Union[Path, str], encrypt: bool = False, async def upload_file(self, path: Union[Path, str], encrypt: bool = False,
) -> Tuple[str, str, Dict[str, Any]]: ) -> Tuple[str, str, CryptDict]:
with open(path, "rb") as file: with open(path, "rb") as file:
mime = utils.guess_mime(file) mime = utils.guess_mime(file)
file.seek(0, 0) file.seek(0, 0)
data: Union[BinaryIO, bytes]
if encrypt: if encrypt:
data, crypt_dict = encrypt_attachment(file.read()) data, crypt_dict = await self.encrypt_attachment(file.read())
upload_mime = "application/octet-stream" upload_mime = "application/octet-stream"
else: else:
data, crypt_dict, upload_mime = file, {}, mime data, crypt_dict, upload_mime = file, {}, mime

View File

@ -12,7 +12,7 @@ from PIL import Image as PILImage
import nio import nio
from .matrix_client import MatrixClient from .matrix_client import MatrixClient, CryptDict
CryptDict = Optional[Dict[str, Any]] CryptDict = Optional[Dict[str, Any]]
Size = Tuple[int, int] Size = Tuple[int, int]