From 74dd6b977c7da2c1621b52c5b9e21c069b754fcf Mon Sep 17 00:00:00 2001
From: Salad Dais <SaladDais@users.noreply.github.com>
Date: Thu, 29 Jul 2021 10:26:42 +0000
Subject: [PATCH] Add extended to_dict() format for Message class

This will allow proper import / export of message logs.
---
 hippolyzer/lib/base/message/message.py | 25 ++++++++++++++++++++++++-
 tests/base/test_message_wrapper.py     |  6 ++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/hippolyzer/lib/base/message/message.py b/hippolyzer/lib/base/message/message.py
index 7bc9460..53b91f9 100644
--- a/hippolyzer/lib/base/message/message.py
+++ b/hippolyzer/lib/base/message/message.py
@@ -291,7 +291,7 @@ class Message:
         if self.raw_body and self.deserializer():
             self.deserializer().parse_message_body(self)
 
-    def to_dict(self):
+    def to_dict(self, extended=False):
         """ A dict representation of a message.
 
         This is the form used for templated messages sent via EQ.
@@ -307,6 +307,18 @@ class Message:
                     new_vars[var_name] = val
                 dict_blocks.append(new_vars)
 
+        if extended:
+            base_repr.update({
+                "packet_id": self.packet_id,
+                "meta": self.meta.copy(),
+                "dropped": self.dropped,
+                "injected": self.injected,
+                "direction": self.direction.name,
+                "send_flags": int(self.send_flags),
+                "extra": self.extra,
+                "acks": self.acks,
+            })
+
         return base_repr
 
     @classmethod
@@ -316,6 +328,17 @@ class Message:
             msg.create_block_list(block_type)
             for block in blocks:
                 msg.add_block(Block(block_type, **block))
+
+        if 'packet_id' in dict_val:
+            # extended format
+            msg.packet_id = dict_val['packet_id']
+            msg.meta = dict_val['meta']
+            msg.dropped = dict_val['dropped']
+            msg.injected = dict_val['injected']
+            msg.direction = Direction[dict_val['direction']]
+            msg.send_flags = dict_val['send_flags']
+            msg.extra = dict_val['extra']
+            msg.acks = dict_val['acks']
         return msg
 
     def invalidate_caches(self):
diff --git a/tests/base/test_message_wrapper.py b/tests/base/test_message_wrapper.py
index 1f98983..0f3dd1d 100644
--- a/tests/base/test_message_wrapper.py
+++ b/tests/base/test_message_wrapper.py
@@ -146,6 +146,12 @@ class TestMessage(unittest.TestCase):
         new_msg = Message.from_dict(self.chat_msg.to_dict())
         self.assertEqual(pickle.dumps(self.chat_msg), pickle.dumps(new_msg))
 
+    def test_todict_extended(self):
+        self.chat_msg.packet_id = 5
+        new_msg = Message.from_dict(self.chat_msg.to_dict(extended=True))
+        self.assertEqual(5, new_msg.packet_id)
+        self.assertEqual(pickle.dumps(self.chat_msg), pickle.dumps(new_msg))
+
     def test_todict_multiple_blocks(self):
         chat_msg = self.chat_msg
         # If we dupe the ChatData block it should survive to_dict()