feat: add image support (#11)

AlexanderBabel · web-flow · commit 1972b242c221 · 2023-01-06T21:50:16.000+01:00
* feat: add image support

* feat: integrate image feature into push route

* fix: set initial thumbnail size in bot.py
diff --git a/hasswebhook/bot.py b/hasswebhook/bot.py
@@ -1,18 +1,19 @@
 import asyncio
+import json
+from datetime import datetime, timedelta
+from typing import Type
 
 import pytz
+from aiohttp.web import Request, Response
 from maubot import Plugin, MessageEvent
 from maubot.handlers import command, web
 from mautrix.types import TextMessageEventContent, Format, MessageType
 from mautrix.util import markdown
 
-from typing import Type
 from .config import Config
-from aiohttp.web import Request, Response
-from .roomposter import RoomPoster, RoomPosterType
-from .setupinstructions import HassWebhookSetupInstructions
 from .db import LifetimeDatabase, LifetimeEnd
-from datetime import datetime, timedelta
+from .roomposter import RoomPoster, RoomPosterType, Image
+from .setupinstructions import HassWebhookSetupInstructions
 
 
 class HassWebhook(Plugin):
@@ -93,34 +94,61 @@ async def setup_instructions(self, evt: MessageEvent) -> None:
     @web.post("/push/{room_id}")
     async def post_data(self, req: Request) -> Response:
         room_id: str = req.match_info["room_id"]
-        self.log.info(await req.text())
+        self.log.info(f"Request for room {room_id} data: {await req.text()}")
+
         req_dict = await req.json()
         self.log.debug(req_dict)
+
         message: str = req_dict.get(self.get_message_key())
         rp_type: RoomPosterType = RoomPosterType.get_type_from_str(req_dict.get("type", "message"))
         identifier: str = req_dict.get("identifier", "")
         callback_url: str = req_dict.get("callback_url", "")
+
         lifetime: int = req_dict.get("lifetime", "")
         if lifetime == "" or int(lifetime) < 0:
             lifetime = -1
         else:
             lifetime = int(lifetime)
-        self.log.info(f"Lifetime 1: {lifetime}")
+        self.log.debug(f"Lifetime: {lifetime}")
+
+        # Image parameters
+        content: str = req_dict.get("content")
+        content_type: str = req_dict.get("contentType")
+        name: str = req_dict.get("name")
+        thumbnail_size: int = req_dict.get("thumbnailSize", 128)
+        image = None
+        self.log.info(content)
+        if not content and rp_type == RoomPosterType.IMAGE:
+            return Response(status=400, content_type="application/json", body=json.dumps(
+                {"success": False,
+                 "error": "Type is set to image. Please pass at least the 'content' property (base64 image)"}))
+
+        if content and rp_type != RoomPosterType.IMAGE:
+            rp_type = RoomPosterType.IMAGE
+
+        if rp_type == RoomPosterType.IMAGE:
+            image = Image(content=content, content_type=content_type, name=name, thumbnail_size=thumbnail_size)
+            self.log.info(f"Image content found: {content}")
+
         room_poster: RoomPoster = RoomPoster(
             hasswebhook=self,
             message=message,
             identifier=identifier,
             rp_type=rp_type,
             room_id=room_id,
             callback_url=callback_url,
-            lifetime=lifetime
+            lifetime=lifetime,
+            image=image,
         )
 
         self.log.debug(f"Received data with ID {room_id}: {req_dict}")
-        if await room_poster.post_to_room():
+
+        event_id = await room_poster.post_to_room()
+        if rp_type == RoomPosterType.MESSAGE or rp_type == RoomPosterType.IMAGE:
+            return Response(status=200, body=json.dumps({"event_id": event_id}), content_type="application/json")
+        elif event_id:
             return Response(status=200)
         else:
-            self.log.debug("I responded with 404")
             return Response(status=404)
 
     @web.get("/health")
diff --git a/hasswebhook/db.py b/hasswebhook/db.py
@@ -1,14 +1,14 @@
-from typing import Optional, Iterator, Dict, List
-from datetime import datetime
 import logging
+from datetime import datetime
+from typing import Iterator
+
 import pytz
 from attr import dataclass
-from sqlalchemy import (Column, String, Integer, Text, DateTime, ForeignKey, Table, MetaData,
+from mautrix.types import EventID, RoomID
+from sqlalchemy import (Column, String, Integer, DateTime, Table, MetaData,
                         select, and_)
 from sqlalchemy.engine.base import Engine
 
-from mautrix.types import UserID, EventID, RoomID
-
 
 @dataclass
 class LifetimeEnd:
@@ -29,10 +29,10 @@ def __init__(self, db: Engine) -> None:
         meta.bind = db
 
         self.lifetime_ends = Table("lifetime_ends", meta,
-                              Column("id", Integer, primary_key=True, autoincrement=True),
-                              Column("end_date", DateTime, nullable=False),
-                              Column("room_id", String(255), nullable=False),
-                              Column("event_id", String(255), nullable=False))
+                                   Column("id", Integer, primary_key=True, autoincrement=True),
+                                   Column("end_date", DateTime, nullable=False),
+                                   Column("room_id", String(255), nullable=False),
+                                   Column("event_id", String(255), nullable=False))
 
         meta.create_all()
 
@@ -47,7 +47,7 @@ def get_older_than(self, end_date: datetime) -> Iterator[LifetimeEnd]:
         rows = self.db.execute(select([self.lifetime_ends]).where(and_(*where)))
         for row in rows:
             yield LifetimeEnd(id=row[0], end_date=row[1].replace(tzinfo=pytz.UTC), room_id=row[2],
-                               event_id=row[3])
+                              event_id=row[3])
 
     def remove(self, lifetime_end: LifetimeEnd) -> None:
         self.db.execute(self.lifetime_ends.delete().where(
diff --git a/hasswebhook/roomposter.py b/hasswebhook/roomposter.py
@@ -1,25 +1,42 @@
-import logging
 import re
+from base64 import b64decode
 from datetime import datetime, timedelta
 from enum import Enum
+from io import BytesIO
 from typing import Optional
 
 import pytz
+from PIL import Image as pil_image
 from markdown import markdown
 from maubot import Plugin
 from maubot.matrix import MaubotMessageEvent
+from mautrix.crypto.attachments import encrypt_attachment
 from mautrix.errors.request import MForbidden
 from mautrix.types import TextMessageEventContent, Format, MessageType, RoomID, PaginationDirection, \
-    MessageEventContent, EventID
+    MessageEventContent, EventID, MediaMessageEventContent, ImageInfo, EventType, ThumbnailInfo
 
 from .db import LifetimeEnd
 
 
+class Image:
+    content: str
+    content_type: str
+    name: str
+    thumbnail_size: int
+
+    def __init__(self, content: str, content_type: str, name: str, thumbnail_size: int):
+        self.content = content
+        self.content_type = content_type
+        self.name = name
+        self.thumbnail_size = thumbnail_size
+
+
 class RoomPosterType(Enum):
     MESSAGE = 1
     EDIT = 2
     REDACTION = 3
     REACTION = 4
+    IMAGE = 5
 
     @classmethod
     def get_type_from_str(cls, mtype: str):
@@ -28,7 +45,8 @@ def get_type_from_str(cls, mtype: str):
             "message": RoomPosterType.MESSAGE,
             "redaction": RoomPosterType.REDACTION,
             "edit": RoomPosterType.EDIT,
-            "reaction": RoomPosterType.REACTION
+            "reaction": RoomPosterType.REACTION,
+            "image": RoomPosterType.IMAGE
         }
         return type_switcher.get(mtype)
 
@@ -43,22 +61,23 @@ class RoomPoster:
     lifetime: int
 
     def __init__(self, hasswebhook: Plugin, identifier: str, rp_type: RoomPosterType, room_id: str,
-                 message="", callback_url="", lifetime=-1):
+                 image: Optional[Image] = None, message="", callback_url="", lifetime=-1):
         self.rp_type = rp_type
         self.room_id = RoomID(room_id)
         self.hasswebhook = hasswebhook
         self.identifier = identifier
         self.callback_url = callback_url
         self.message = message
         self.lifetime = lifetime
+        self.image = image
 
     # Send a POST as a callback containing the event_id of the sent message
     async def callback(self, event_id: str) -> None:
         if self.callback_url:
             await self.hasswebhook.http.post(self.callback_url, json={'event_id': event_id})
 
     # Switch for each RoomPosterType
-    async def post_to_room(self) -> bool:
+    async def post_to_room(self):
         if self.rp_type == RoomPosterType.MESSAGE:
             return await self.post_message()
         if self.rp_type == RoomPosterType.REDACTION:
@@ -67,10 +86,44 @@ async def post_to_room(self) -> bool:
             return await self.post_edit()
         if self.rp_type == RoomPosterType.REACTION:
             return await self.post_reaction()
+        if self.rp_type == RoomPosterType.IMAGE:
+            return await self.post_image()
         return False
 
+    async def post_image(self) -> str:
+        media_event = MediaMessageEventContent(body=self.image.name, msgtype=MessageType.IMAGE)
+
+        upload_mime = "application/octet-stream"
+        bytes_image = b64decode(self.image.content)
+        encrypted_image, file = encrypt_attachment(bytes_image)
+        file.url = await self.hasswebhook.client.upload_media(encrypted_image, mime_type=upload_mime)
+        media_event.file = file
+
+        img = pil_image.open(BytesIO(bytes_image))
+        image_info = ImageInfo(mimetype=self.image.content_type, height=img.height, width=img.width)
+        media_event.info = image_info
+
+        byt_arr_tn = BytesIO()
+        img.thumbnail((self.image.thumbnail_size, self.image.thumbnail_size), pil_image.ANTIALIAS)
+        img.save(byt_arr_tn, format='PNG')
+
+        tn_img = pil_image.open(byt_arr_tn)
+        enc_tn, tn_file = encrypt_attachment(byt_arr_tn.getvalue())
+
+        image_info.thumbnail_info = ThumbnailInfo(mimetype="image/png", height=tn_img.height, width=tn_img.width)
+        tn_file.url = await self.hasswebhook.client.upload_media(enc_tn, mime_type=upload_mime)
+        image_info.thumbnail_file = tn_file
+
+        img.close()
+        tn_img.close()
+
+        event_id = await self.hasswebhook.client.send_message_event(self.room_id, event_type=EventType.ROOM_MESSAGE,
+                                                                    content=media_event)
+        await self.callback(event_id)
+        return event_id
+
     # Send message to room
-    async def post_message(self) -> bool:
+    async def post_message(self):
         body = "{message} by {identifier}".format(
             message=self.message, identifier=self.identifier) if self.identifier else self.message
         content = TextMessageEventContent(
@@ -87,10 +140,10 @@ async def post_message(self) -> bool:
                 end_time = datetime.now(tz=pytz.UTC) + timedelta(minutes=self.lifetime)
                 self.hasswebhook.db.insert(
                     LifetimeEnd(end_date=end_time, room_id=self.room_id, event_id=event_id_req))
+            return event_id_req
         except MForbidden:
             self.hasswebhook.log.error("Wrong Room ID")
             return False
-        return True
 
     # Redact message
     async def post_redaction(self) -> bool:
diff --git a/maubot.yaml b/maubot.yaml
@@ -1,12 +1,17 @@
 maubot: 0.1.0
 id: com.valentinriess.hasswebhook
-version: 0.0.14
+version: 0.0.15
 license: MIT
 modules:
   - hasswebhook
 main_class: HassWebhook
 config: true
 extra_files:
-- base-config.yaml
+  - base-config.yaml
 database: true
 webapp: true
+soft_dependencies:
+  - Pillow
+dependencies:
+  - Markdown
+  - pytz
