Support encrypted file reading & caching

Also don't convert palette images to JPEG when creating thumbnails.

TODO.md

@@ -1,8 +1,6 @@
 - Media
-  - Caching
+  - Encrypt file for upload in thread
-    - What effect will it have on GIFs? Can we set `cache:false` on them or get
+  - Cache our own uploads
-      the frame count once they're cached?
-  - Reading encrypted media
   - Uploading progress (+local echo)
     - Deduplicate uploads
   - Loading progress bar
@@ -198,7 +196,6 @@
   - Running blocking DB function calls in executor
   - Guard against asyncio OSError Network unreachable
 
-  - downloads API
   - MatrixRoom invited members list
   - Left room events after client reboot
   - `org.matrix.room.preview_urls` events

src/python/matrix_client.py

@@ -239,7 +239,9 @@ class MatrixClient(nio.AsyncClient):
             content["url"] = url
 
         if kind == "image":
-            event_type         = nio.RoomMessageImage
+            event_type = \
+                nio.RoomEncryptedImage if encrypt else nio.RoomMessageImage
+
             content["msgtype"] = "m.image"
 
             content["info"]["w"], content["info"]["h"] = \
@@ -262,14 +264,18 @@ class MatrixClient(nio.AsyncClient):
                 content["info"]["thumbnail_info"] = thumb_info
 
         elif kind == "audio":
-            event_type                  = nio.RoomMessageAudio
+            event_type = \
+                nio.RoomEncryptedAudio if encrypt else nio.RoomMessageAudio
+
             content["msgtype"]          = "m.audio"
             content["info"]["duration"] = getattr(
                 MediaInfo.parse(path).tracks[0], "duration", 0,
             ) or 0
 
         elif kind == "video":
-            event_type         = nio.RoomMessageVideo
+            event_type = \
+                nio.RoomEncryptedVideo if encrypt else nio.RoomMessageVideo
+
             content["msgtype"] = "m.video"
 
             tracks = MediaInfo.parse(path).tracks
@@ -285,7 +291,9 @@ class MatrixClient(nio.AsyncClient):
             )
 
         else:
-            event_type          = nio.RoomMessageFile
+            event_type = \
+                nio.RoomEncryptedFile if encrypt else nio.RoomMessageFile
+
             content["msgtype"]  = "m.file"
             content["filename"] = path.name
 
@@ -414,21 +422,24 @@ class MatrixClient(nio.AsyncClient):
     async def upload_thumbnail(
         self, path: Union[Path, str], encrypt: bool = False,
     ) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:
+
+        png_modes = ("1", "L", "P", "RGBA")
+
         try:
             thumb = PILImage.open(path)
 
             small       = thumb.width <= 800 and thumb.height <= 600
             is_jpg_png  = thumb.format in ("JPEG", "PNG")
-            opaque_png = thumb.format == "PNG" and thumb.mode != "RGBA"
+            jpgable_png = thumb.format == "PNG" and thumb.mode not in png_modes
 
-            if small and is_jpg_png and not opaque_png:
+            if small and is_jpg_png and not jpgable_png:
                 raise UneededThumbnail()
 
             if not small:
                 thumb.thumbnail((800, 600), PILImage.LANCZOS)
 
             with io.BytesIO() as out:
-                if thumb.mode == "RGBA":
+                if thumb.mode in png_modes:
                     thumb.save(out, "PNG", optimize=True)
                     mime = "image/png"
                 else:
@@ -805,13 +816,16 @@ class MatrixClient(nio.AsyncClient):
 
     async def onRoomMessageMedia(self, room, ev) -> None:
         info             = ev.source["content"].get("info", {})
+        media_crypt_dict = ev.source["content"].get("file", {})
         thumb_info       = info.get("thumbnail_info", {})
+        thumb_crypt_dict = info.get("thumbnail_file", {})
 
         await self.register_nio_event(
             room,
             ev,
             content        = "",
             inline_content = ev.body,
+
             media_url        = ev.url,
             media_title      = ev.body,
             media_width      = info.get("w") or 0,
@@ -819,13 +833,21 @@ class MatrixClient(nio.AsyncClient):
             media_duration   = info.get("duration") or 0,
             media_size       = info.get("size") or 0,
             media_mime       = info.get("mimetype") or 0,
+            media_crypt_dict = media_crypt_dict,
+
+            thumbnail_url =
+                info.get("thumbnail_url") or thumb_crypt_dict.get("url") or "",
 
-            thumbnail_url    = info.get("thumbnail_url") or "",
             thumbnail_width      = thumb_info.get("w") or 0,
             thumbnail_height     = thumb_info.get("h") or 0,
+            thumbnail_crypt_dict = thumb_crypt_dict,
         )
 
 
+    async def onRoomEncryptedMedia(self, room, ev) -> None:
+        await self.onRoomMessageMedia(room, ev)
+
+
     async def onRoomCreateEvent(self, room, ev) -> None:
         co = "%1 allowed users on other matrix servers to join this room." \
              if ev.federate else \

src/python/media_cache.py

@@ -1,17 +1,20 @@
 import asyncio
+import functools
 import io
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import DefaultDict, Optional, Tuple
+from typing import Any, DefaultDict, Dict, Optional, Tuple
 from urllib.parse import urlparse
 
 import aiofiles
-import nio
 from PIL import Image as PILImage
 
+import nio
+
 from .matrix_client import MatrixClient
 
+CryptDict = Optional[Dict[str, Any]]
 Size      = Tuple[int, int]
 
 CONCURRENT_DOWNLOADS_LIMIT                   = asyncio.BoundedSemaphore(8)
@@ -29,6 +32,7 @@ class Media:
     cache:      "MediaCache"    = field()
     mxc:        str             = field()
     data:       Optional[bytes] = field(repr=False)
+    crypt_dict: CryptDict       = field(repr=False)
 
 
     def __post_init__(self) -> None:
@@ -89,7 +93,23 @@ class Media:
         if isinstance(resp, nio.DownloadError):
             raise DownloadFailed(resp.message, resp.status_code)
 
-        return resp.body
+        return await self._decrypt(resp.body)
+
+
+    async def _decrypt(self, data: bytes) -> bytes:
+        if not self.crypt_dict:
+            return data
+
+        func = functools.partial(
+            nio.crypto.attachments.decrypt_attachment,
+            data,
+            self.crypt_dict["key"]["k"],
+            self.crypt_dict["hashes"]["sha256"],
+            self.crypt_dict["iv"],
+        )
+
+        # Run in a separate thread
+        return await asyncio.get_event_loop().run_in_executor(None, func)
 
 
 @dataclass
@@ -97,6 +117,7 @@ class Thumbnail(Media):
     cache:       "MediaCache"    = field()
     mxc:         str             = field()
     data:        Optional[bytes] = field(repr=False)
+    crypt_dict:  CryptDict       = field(repr=False)
     wanted_size: Size            = field()
 
     server_size: Optional[Size] = field(init=False, repr=False, default=None)
@@ -159,6 +180,12 @@ class Thumbnail(Media):
     async def _get_remote_data(self) -> bytes:
         parsed = urlparse(self.mxc)
 
+        if self.crypt_dict:
+            resp = await self.cache.client.download(
+                server_name = parsed.netloc,
+                media_id    = parsed.path.lstrip("/"),
+            )
+        else:
             resp = await self.cache.client.thumbnail(
                 server_name = parsed.netloc,
                 media_id    = parsed.path.lstrip("/"),
@@ -166,14 +193,16 @@ class Thumbnail(Media):
                 height      = self.wanted_size[1],
             )
 
-        if isinstance(resp, nio.ThumbnailError):
+        if isinstance(resp, (nio.DownloadError, nio.ThumbnailError)):
             raise DownloadFailed(resp.message, resp.status_code)
 
-        with io.BytesIO(resp.body) as img:
+        decrypted = await self._decrypt(resp.body)
+
+        with io.BytesIO(decrypted) as img:
             # The server may return a thumbnail bigger than what we asked for
             self.server_size = PILImage.open(img).size
 
-        return resp.body
+        return decrypted
 
 
 @dataclass
@@ -190,9 +219,13 @@ class MediaCache:
         self.downloads_dir.mkdir(parents=True, exist_ok=True)
 
 
-    async def get_media(self, mxc: str) -> str:
+    async def get_media(self, mxc: str, crypt_dict: CryptDict = None) -> str:
-        return str(await Media(self, mxc, data=None).get())
+        return str(await Media(self, mxc, None, crypt_dict).get())
 
 
-    async def get_thumbnail(self, mxc: str, width: int, height: int) -> str:
+    async def get_thumbnail(
-        return str(await Thumbnail(self, mxc, None, (width, height)).get())
+        self, mxc: str, width: int, height: int, crypt_dict: CryptDict = None,
+    ) -> str:
+
+        thumb = Thumbnail(self, mxc, None, crypt_dict, (width, height))
+        return str(await thumb.get())

src/python/models/items.py

@@ -136,10 +136,12 @@ class Event(ModelItem):
     media_duration:   int            = 0
     media_size:       int            = 0
     media_mime:       str            = ""
+    media_crypt_dict: Dict[str, Any] = field(default_factory=dict)
 
     thumbnail_url:        str            = ""
     thumbnail_width:      int            = 0
     thumbnail_height:     int            = 0
+    thumbnail_crypt_dict: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         if not self.inline_content:
@@ -156,7 +158,8 @@ class Event(ModelItem):
 
     @property
     def links(self) -> List[str]:
-        if isinstance(self.source, nio.RoomMessageMedia):
+        if isinstance(self.source,
+                      (nio.RoomMessageMedia, nio.RoomEncryptedMedia)):
             return [self.media_url]
 
         if not self.content.strip():

src/qml/Base/HMxcImage.qml

@@ -21,6 +21,7 @@ HImage {
     property string mxc
     property string sourceOverride: ""
     property bool thumbnail: true
+    property var cryptDict: ({})
 
     property bool show: false
     property string cachedPath: ""
@@ -45,7 +46,8 @@ HImage {
         }
 
         let method = image.thumbnail ? "get_thumbnail" : "get_media"
-        let args = image.thumbnail ? [image.mxc, w, h] : [image.mxc]
+        let args = image.thumbnail ?
+                   [image.mxc, w, h, cryptDict] : [image.mxc, cryptDict]
 
         py.callClientCoro(
             clientUserId, "media_cache." + method, args, path => {

src/qml/Chat/Timeline/EventDelegate.qml

@@ -36,14 +36,14 @@ Column {
 
     readonly property bool smallAvatar:
         eventList.canCombine(model, nextItem) &&
-        (model.event_type == "RoomMessageEmote" ||
+        (model.event_type === "RoomMessageEmote" ||
          ! model.event_type.startsWith("RoomMessage"))
 
     readonly property bool collapseAvatar: combine
     readonly property bool hideAvatar: onRight
 
     readonly property bool hideNameLine:
-        model.event_type == "RoomMessageEmote" ||
+        model.event_type === "RoomMessageEmote" ||
         ! model.event_type.startsWith("RoomMessage") ||
         onRight ||
         combine

src/qml/Chat/Timeline/EventImage.qml

@@ -14,6 +14,9 @@ HMxcImage {
     mxc: thumbnail ?
          (loader.thumbnailMxc || loader.mediaUrl) :
          (loader.mediaUrl || loader.thumbnailMxc)
+    cryptDict: thumbnail && loader.thumbnailMxc ?
+               loader.singleMediaInfo.thumbnail_crypt_dict :
+               loader.singleMediaInfo.media_crypt_dict
 
 
     property EventMediaLoader loader

src/qml/SidePane/RoomDelegate.qml

@@ -71,6 +71,7 @@ HTileDelegate {
 
         let ev = model.data.last_event
 
+        // If it's an emote or non-message/media event
         if (ev.event_type === "RoomMessageEmote" ||
             ! ev.event_type.startsWith("RoomMessage")) {
             return Utils.processedEventText(ev)

src/qml/utils.js

@@ -109,6 +109,7 @@ function processedEventText(ev) {
     }
 
     if (ev.event_type.startsWith("RoomMessage")) { return ev.content }
+    if (ev.event_type.startsWith("RoomEncrypted")) { return ev.content }
 
     let text = qsTr(ev.content).arg(
         coloredNameHtml(ev.sender_name, ev.sender_id)