src/message.py

   1 import base64
   2 import html2text
   3 from enum import Enum
   4
   5 class Message:
   6     class Format(Enum):
   7         PLAIN = 0
   8         HTML = 1
   9
  10     def __init__(self, fragment):
  11         self._next_fragment = 0
  12         self._raw_content = ""
  13
  14         parts = fragment.split()
  15
  16         if len(parts) != 2:
  17             raise Exception("Invalid fragment!")
  18
  19         if parts[0] == "PLAIN":
  20             self._msg_fmt = Message.Format.PLAIN
  21         elif parts[0] == "HTML":
  22             self._msg_fmt = Message.Format.HTML
  23         else:
  24             raise Exception("Unknown message type %s!" % msg_fmt)
  25
  26         try:
  27             self._msg_len = int(parts[1])
  28         except:
  29             raise Exception("Invalid message length %s!" % parts[1])
  30
  31         if self._msg_len < 1 or self._msg_len > 1000:
  32             raise Exception("Invalid message length %d!" % self._msg_len)
  33
  34     def feed_fragment(self, fragment):
  35         parts = fragment.split()
  36
  37         if len(parts) != 2:
  38             raise Exception("Invalid fragment!")
  39
  40         try:
  41             fragment_number = int(parts[0])
  42         except:
  43             raise Exception("Invalid fragment number %s!" % parts[0])
  44
  45         if fragment_number == self._next_fragment and fragment_number < self._msg_len:
  46             self._next_fragment += 1
  47         else:
  48             raise Exception("Unexpected fragment!")
  49
  50         self._raw_content += parts[1]
  51
  52     def is_complete(self):
  53         return self._next_fragment == self._msg_len
  54
  55     def get_content(self):
  56         try:
  57             msgB64Bytes = self._raw_content.encode("UTF-8")
  58             msgBytes = base64.b64decode(msgB64Bytes)
  59             msgStr = msgBytes.decode("UTF-8")
  60         except:
  61             raise Exception("Message decoding error!")
  62
  63         if self._msg_fmt == Message.Format.PLAIN:
  64             return {
  65                 "msgtype": "m.text",
  66                 "body": msgStr,
  67             }
  68         elif self._msg_fmt == Message.Format.HTML:
  69             return {
  70                 "format": "org.matrix.custom.html",
  71                 "msgtype": "m.text",
  72                 "body": html2text.html2text(msgStr),
  73                 "formatted_body": msgStr
  74             }