From 4b7cfed076fca56f0e039feecf71425cecd93697 Mon Sep 17 00:00:00 2001
From: mauricio garcia suarez <mauricio@openreplay.com>
Date: Fri, 26 Aug 2022 11:38:34 +0200
Subject: [PATCH] Modification of Message codec for python

---
 ee/connectors/msgcodec/codec.py               |  2 +
 .../ee~connectors~msgcodec~msgcodec.py.erb    | 57 +++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/ee/connectors/msgcodec/codec.py b/ee/connectors/msgcodec/codec.py
index 5aeb0e4ed..577f02f48 100644
--- a/ee/connectors/msgcodec/codec.py
+++ b/ee/connectors/msgcodec/codec.py
@@ -21,6 +21,8 @@ class Codec:
         i = 0  # n of byte (max 9 for uint64)
         while True:
             b = reader.read(1)
+            if len(b) == 0:
+                raise IndexError('bytes out of range')
             num = int.from_bytes(b, "big", signed=False)
             # print(i, x)
 
diff --git a/mobs/templates/ee~connectors~msgcodec~msgcodec.py.erb b/mobs/templates/ee~connectors~msgcodec~msgcodec.py.erb
index 6a5c5b71f..b55764ee9 100644
--- a/mobs/templates/ee~connectors~msgcodec~msgcodec.py.erb
+++ b/mobs/templates/ee~connectors~msgcodec~msgcodec.py.erb
@@ -2,6 +2,8 @@
 
 from msgcodec.codec import Codec
 from msgcodec.messages import *
+from typing import List
+import io
 
 class MessageCodec(Codec):
 
@@ -17,7 +19,62 @@ class MessageCodec(Codec):
 
     def decode(self, b: bytes) -> Message:
         reader = io.BytesIO(b)
+        return self.read_head_message(reader)
+
+    @staticmethod
+    def check_message_id(b: bytes) -> int:
+        """
+        todo: make it static and without reader. It's just the first byte
+        Read and return the first byte where the message id is encoded
+        """
+        reader = io.BytesIO(b)
+        id_ = Codec.read_uint(reader)
+
+        return id_
+
+    @staticmethod
+    def decode_key(b) -> int:
+        """
+        Decode the message key (encoded with little endian)
+        """
+        try:
+            decoded = int.from_bytes(b, "little", signed=False)
+        except Exception as e:
+            raise UnicodeDecodeError(f"Error while decoding message key (SessionID) from {b}\n{e}")
+        return decoded
+
+    def decode_detailed(self, b: bytes) -> List[Message]:
+        reader = io.BytesIO(b)
+        messages_list = list()
+        messages_list.append(self.handler(reader, 0))
+        if isinstance(messages_list[0], BatchMeta):
+            # Old BatchMeta
+            mode = 0
+        elif isinstance(messages_list[0], BatchMetadata):
+            # New BatchMeta
+            mode = 1
+        else:
+            return messages_list
+        while True:
+            try:
+                messages_list.append(self.handler(reader, mode))
+            except IndexError:
+                break
+        return messages_list
+
+    def handler(self, reader: io.BytesIO, mode=0) -> Message:
         message_id = self.read_message_id(reader)
+        if mode == 1:
+            # We skip the three bytes representing the length of message. It can be used to skip unwanted messages
+            reader.read(3)
+            return self.read_head_message(reader, message_id)
+        elif mode == 0:
+            # Old format with no bytes for message length
+            return self.read_head_message(reader, message_id)
+        else:
+            raise IOError()
+
+    def read_head_message(self, reader: io.BytesIO, message_id) -> Message:
 <% $messages.each do |msg| %>
         if message_id == <%= msg.id %>:
             return <%= msg.name %>(