support html/markdown/plain messages of arbitrary length

[mtxbot.git] / src / message.py

diff --git a/src/message.py b/src/message.py
new file mode 100644 (file)
index 0000000..d66b6ef
--- /dev/null
+++ b/src/message.py
@@ -0,0 +1,74 @@
+import base64
+import html2text
+from enum import Enum
+
+class Message:
+    class Format(Enum):
+        PLAIN = 0
+        HTML = 1
+
+    def __init__(self, fragment):
+        self._next_fragment = 0
+        self._raw_content = ""
+
+        parts = fragment.split()
+
+        if len(parts) != 2:
+            raise Exception("Invalid fragment!")
+
+        if parts[0] == "PLAIN":
+            self._msg_fmt = Message.Format.PLAIN
+        elif parts[0] == "HTML":
+            self._msg_fmt = Message.Format.HTML
+        else:
+            raise Exception("Unknown message type %s!" % msg_fmt)
+
+        try:
+            self._msg_len = int(parts[1])
+        except:
+            raise Exception("Invalid message length %s!" % parts[1])
+
+        if self._msg_len < 1 or self._msg_len > 1000:
+            raise Exception("Invalid message length %d!" % self._msg_len)
+
+    def feed_fragment(self, fragment):
+        parts = fragment.split()
+
+        if len(parts) != 2:
+            raise Exception("Invalid fragment!")
+
+        try:
+            fragment_number = int(parts[0])
+        except:
+            raise Exception("Invalid fragment number %s!" % parts[0])
+
+        if fragment_number == self._next_fragment and fragment_number < self._msg_len:
+            self._next_fragment += 1
+        else:
+            raise Exception("Unexpected fragment!")
+
+        self._raw_content += parts[1]
+
+    def is_complete(self):
+        return self._next_fragment == self._msg_len
+
+    def get_content(self):
+        try:
+            msgB64Bytes = self._raw_content.encode("UTF-8")
+            msgBytes = base64.b64decode(msgB64Bytes)
+            msgStr = msgBytes.decode("UTF-8")
+        except:
+            raise Exception("Message decoding error!")
+
+        if self._msg_fmt == Message.Format.PLAIN:
+            return {
+                "msgtype": "m.text",
+                "body": msgStr,
+            }
+        elif self._msg_fmt == Message.Format.HTML:
+            return {
+                "format": "org.matrix.custom.html",
+                "msgtype": "m.text",
+                "body": html2text.html2text(msgStr),
+                "formatted_body": msgStr
+            }