Implement base64_attachments

signalbot/api.py

@@ -1,3 +1,5 @@
+import base64
+
 import aiohttp
 import websockets
 
@@ -101,6 +103,28 @@ async def stop_typing(self, receiver: str):
         ):
             raise StopTypingError
 
+    async def get_attachment(self, attachment_id: str) -> str:
+        """Fetch attachment by given id and encode to base64."""
+        uri = f"{self._attachment_rest_uri()}/{attachment_id}"
+        try:
+            async with aiohttp.ClientSession() as session:
+                resp = await session.get(uri)
+                resp.raise_for_status()
+                content = await resp.content.read()
+        except (
+            aiohttp.ClientError,
+            aiohttp.http_exceptions.HttpProcessingError,
+        ):
+            raise GetAttachmentError
+
+        base64_bytes = base64.b64encode(content)
+        base64_string = str(base64_bytes, encoding="utf-8")
+
+        return base64_string
+
+    def _attachment_rest_uri(self):
+        return f"http://{self.signal_service}/v1/attachments"
+
     def _receive_ws_uri(self):
         return f"ws://{self.signal_service}/v1/receive/{self.phone_number}"
 
@@ -136,3 +160,7 @@ class StopTypingError(TypingError):
 
 class ReactionError(Exception):
     pass
+
+
+class GetAttachmentError(Exception):
+    pass

signalbot/bot.py

@@ -235,7 +235,7 @@ async def _produce(self, name: int) -> None:
                 logging.info(f"[Raw Message] {raw_message}")
 
                 try:
-                    message = Message.parse(raw_message)
+                    message = await Message.parse(self._signal, raw_message)
                 except UnknownMessageFormatError:
                     continue
 

signalbot/message.py

@@ -1,6 +1,8 @@
 import json
 from enum import Enum
 
+from signalbot.api import SignalAPI
+
 
 class MessageType(Enum):
     SYNC_MESSAGE = 1
@@ -25,8 +27,7 @@ def __init__(
         self.type = type
         self.text = text
         # optional
-        if base64_attachments is None:
-            self.base64_attachments = []
+        self.base64_attachments = base64_attachments or []
         self.group = group
         self.reaction = reaction
 
@@ -41,7 +42,7 @@ def recipient(self) -> str:
         return self.source
 
     @classmethod
-    def parse(cls, raw_message: str):
+    async def parse(cls, signal: SignalAPI, raw_message: str):
         try:
             raw_message = json.loads(raw_message)
         except Exception:
@@ -64,22 +65,34 @@ def parse(cls, raw_message: str):
             reaction = cls._parse_reaction(
                 raw_message["envelope"]["syncMessage"]["sentMessage"]
             )
+            base64_attachments = None
 
         # Option 2: dataMessage
         elif "dataMessage" in raw_message["envelope"]:
             type = MessageType.DATA_MESSAGE
             text = cls._parse_data_message(raw_message["envelope"]["dataMessage"])
             group = cls._parse_group_information(raw_message["envelope"]["dataMessage"])
             reaction = cls._parse_reaction(raw_message["envelope"]["dataMessage"])
+            base64_attachments = await cls._parse_attachments(
+                signal, raw_message["envelope"]["dataMessage"]
+            )
 
         else:
             raise UnknownMessageFormatError
 
-        # TODO: base64_attachments
-        base64_attachments = []
-
         return cls(source, timestamp, type, text, base64_attachments, group, reaction)
 
+    @classmethod
+    async def _parse_attachments(cls, signal: SignalAPI, data_message: dict) -> str:
+
+        if "attachments" not in data_message:
+            return []
+
+        return [
+            await signal.get_attachment(attachment["id"])
+            for attachment in data_message["attachments"]
+        ]
+
     @classmethod
     def _parse_sync_message(cls, sync_message: dict) -> str:
         try:

tests/test_message.py

@@ -1,67 +1,97 @@
+import base64
 import unittest
+from unittest.mock import AsyncMock, patch
+
 from signalbot import Message, MessageType
+from signalbot.api import SignalAPI
 
 
 class TestMessage(unittest.TestCase):
     raw_sync_message = '{"envelope":{"source":"+490123456789","sourceNumber":"+490123456789","sourceUuid":"<uuid>","sourceName":"<name>","sourceDevice":1,"timestamp":1632576001632,"syncMessage":{"sentMessage":{"timestamp":1632576001632,"message":"Uhrzeit","expiresInSeconds":0,"viewOnce":false,"mentions":[],"attachments":[],"contacts":[],"groupInfo":{"groupId":"<groupid>","type":"DELIVER"},"destination":null,"destinationNumber":null,"destinationUuid":null}}}}'  # noqa
     raw_data_message = '{"envelope":{"source":"+490123456789","sourceNumber":"+490123456789","sourceUuid":"<uuid>","sourceName":"<name>","sourceDevice":1,"timestamp":1632576001632,"dataMessage":{"timestamp":1632576001632,"message":"Uhrzeit","expiresInSeconds":0,"viewOnce":false,"mentions":[],"attachments":[],"contacts":[],"groupInfo":{"groupId":"<groupid>","type":"DELIVER"}}}}'  # noqa
     raw_reaction_message = '{"envelope":{"source":"<source>","sourceNumber":"<source>","sourceUuid":"<uuid>","sourceName":"<name>","sourceDevice":1,"timestamp":1632576001632,"syncMessage":{"sentMessage":{"timestamp":1632576001632,"message":null,"expiresInSeconds":0,"viewOnce":false,"reaction":{"emoji":"👍","targetAuthor":"<target>","targetAuthorNumber":"<target>","targetAuthorUuid":"<uuid>","targetSentTimestamp":1632576001632,"isRemove":false},"mentions":[],"attachments":[],"contacts":[],"groupInfo":{"groupId":"<groupid>","type":"DELIVER"},"destination":null,"destinationNumber":null,"destinationUuid":null}}}}'  # noqa
     raw_user_chat_message = '{"envelope":{"source":"+490123456789","sourceNumber":"+490123456789","sourceUuid":"<uuid>","sourceName":"<name>","sourceDevice":1,"timestamp":1632576001632,"dataMessage":{"timestamp":1632576001632,"message":"Uhrzeit","expiresInSeconds":0,"viewOnce":false}},"account":"+49987654321","subscription":0}'  # noqa
+    raw_attachment_message = '{"envelope":{"source":"+490123456789","sourceNumber":"+490123456789","sourceUuid":"<uuid>","sourceName":"<name>","sourceDevice":1,"timestamp":1632576001632,"dataMessage":{"timestamp":1632576001632,"message":"Uhrzeit","expiresInSeconds":0,"viewOnce":false, "attachments": [{"contentType": "image/png", "filename": "image.png", "id": "4296180834490578536","size": 12005}]}},"account":"+49987654321","subscription":0}'  # noqa
 
     expected_source = "+490123456789"
     expected_timestamp = 1632576001632
     expected_text = "Uhrzeit"
     expected_group = "<groupid>"
 
+    signal_service = "127.0.0.1:8080"
+    phone_number = "+49123456789"
+
+    group_id = "group_id1"
+    group_secret = "group.group_secret1"
+    groups = {group_id: group_secret}
+
+    def setUp(self):
+        self.signal_api = SignalAPI(
+            TestMessage.signal_service, TestMessage.phone_number
+        )
+
     # Own Message
-    def test_parse_source_own_message(self):
-        message = Message.parse(TestMessage.raw_sync_message)
+    async def test_parse_source_own_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_sync_message)
         self.assertEqual(message.timestamp, TestMessage.expected_timestamp)
 
-    def test_parse_timestamp_own_message(self):
-        message = Message.parse(TestMessage.raw_sync_message)
+    async def test_parse_timestamp_own_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_sync_message)
         self.assertEqual(message.source, TestMessage.expected_source)
 
-    def test_parse_type_own_message(self):
-        message = Message.parse(TestMessage.raw_sync_message)
+    async def test_parse_type_own_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_sync_message)
         self.assertEqual(message.type, MessageType.SYNC_MESSAGE)
 
-    def test_parse_text_own_message(self):
-        message = Message.parse(TestMessage.raw_sync_message)
+    async def test_parse_text_own_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_sync_message)
         self.assertEqual(message.text, TestMessage.expected_text)
 
-    def test_parse_group_own_message(self):
-        message = Message.parse(TestMessage.raw_sync_message)
+    async def test_parse_group_own_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_sync_message)
         self.assertEqual(message.group, TestMessage.expected_group)
 
     # Foreign Messages
-    def test_parse_source_foreign_message(self):
-        message = Message.parse(TestMessage.raw_data_message)
+    async def test_parse_source_foreign_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_data_message)
         self.assertEqual(message.timestamp, TestMessage.expected_timestamp)
 
-    def test_parse_timestamp_foreign_message(self):
-        message = Message.parse(TestMessage.raw_data_message)
+    async def test_parse_timestamp_foreign_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_data_message)
         self.assertEqual(message.source, TestMessage.expected_source)
 
-    def test_parse_type_foreign_message(self):
-        message = Message.parse(TestMessage.raw_data_message)
+    async def test_parse_type_foreign_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_data_message)
         self.assertEqual(message.type, MessageType.DATA_MESSAGE)
 
-    def test_parse_text_foreign_message(self):
-        message = Message.parse(TestMessage.raw_data_message)
+    async def test_parse_text_foreign_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_data_message)
         self.assertEqual(message.text, TestMessage.expected_text)
 
-    def test_parse_group_foreign_message(self):
-        message = Message.parse(TestMessage.raw_data_message)
+    async def test_parse_group_foreign_message(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_data_message)
         self.assertEqual(message.group, TestMessage.expected_group)
 
-    def test_read_reaction(self):
-        message = Message.parse(TestMessage.raw_reaction_message)
+    async def test_read_reaction(self):
+        message = await Message.parse(self.signal_api, TestMessage.raw_reaction_message)
         self.assertEqual(message.reaction, "👍")
 
+    @patch("aiohttp.ClientSession.get", new_callable=AsyncMock)
+    async def test_attachments(self, mock):
+        mock.return_value = b"test"
+        expected_base64_bytes = base64.b64encode(mock.return_value)
+        expected_base64_str = str(expected_base64_bytes, encoding="utf-8")
+
+        message = await Message.parse(
+            self.signal_api, TestMessage.raw_attachment_message
+        )
+        self.assertEqual(message.base64_attachments, [expected_base64_str])
+
     # User Chats
-    def test_parse_user_chat_message(self):
-        message = Message.parse(TestMessage.raw_user_chat_message)
+    async def test_parse_user_chat_message(self):
+        message = await Message.parse(
+            self.signal_api, TestMessage.raw_user_chat_message
+        )
         self.assertEqual(message.source, TestMessage.expected_source)
         self.assertEqual(message.text, TestMessage.expected_text)
         self.assertEqual(message.timestamp, TestMessage.expected_timestamp)