+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
+ }