Support encrypted file reading & caching
Also don't convert palette images to JPEG when creating thumbnails.
This commit is contained in:
		
							
								
								
									
										7
									
								
								TODO.md
									
									
									
									
									
								
							
							
						
						
									
										7
									
								
								TODO.md
									
									
									
									
									
								
							| @@ -1,8 +1,6 @@ | ||||
| - Media | ||||
|   - Caching | ||||
|     - What effect will it have on GIFs? Can we set `cache:false` on them or get | ||||
|       the frame count once they're cached? | ||||
|   - Reading encrypted media | ||||
|   - Encrypt file for upload in thread | ||||
|   - Cache our own uploads | ||||
|   - 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 | ||||
|   | ||||
| @@ -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 \ | ||||
|   | ||||
| @@ -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: | ||||
|         return str(await Media(self, mxc, data=None).get()) | ||||
|     async def get_media(self, mxc: str, crypt_dict: CryptDict = None) -> str: | ||||
|         return str(await Media(self, mxc, None, crypt_dict).get()) | ||||
|  | ||||
|  | ||||
|     async def get_thumbnail(self, mxc: str, width: int, height: int) -> str: | ||||
|         return str(await Thumbnail(self, mxc, None, (width, height)).get()) | ||||
|     async def get_thumbnail( | ||||
|         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()) | ||||
|   | ||||
| @@ -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(): | ||||
|   | ||||
| @@ -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 => { | ||||
|   | ||||
| @@ -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 | ||||
|   | ||||
| @@ -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 | ||||
|   | ||||
| @@ -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) | ||||
|   | ||||
| @@ -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) | ||||
|   | ||||
		Reference in New Issue
	
	Block a user
	