From c3ff27f3b9b9a3502e31999df59f412ca6657baf Mon Sep 17 00:00:00 2001
From: ChillerDragon <ChillerDragon@gmail.com>
Date: Sun, 16 Jun 2024 13:07:08 +0800
Subject: [PATCH] feat: implement to_json() for TwPacket

---
 README.md                 | 23 ++++++++++++++++++++
 tests/json_test.py        | 44 +++++++++++++++++++++++++++++++++++++++
 twnet_parser/packet.py    | 14 +++++++++++--
 twnet_parser/serialize.py |  5 +++++
 4 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 tests/json_test.py
 create mode 100644 twnet_parser/serialize.py

diff --git a/README.md b/README.md
index 740d022..7aac6ae 100644
--- a/README.md
+++ b/README.md
@@ -25,6 +25,29 @@ print(packet.header) # => <class: 'Header'>: {'flags': <class: 'PacketFlags7, 's
 print(packet.header.flags) # => <class: 'PacketFlags7'>: {'control': True, 'resend': False, 'compression': False, 'connless': False}
 for msg in packet.messages:
     print(msg.message_name) # => close
+
+print(packet.to_json())
+# {
+#   "version": "0.7",
+#   "payload_raw": "04",
+#   "payload_decompressed": "04",
+#   "header": {
+#     "flags": [
+#       "control"
+#     ],
+#     "ack": 10,
+#     "token": "cf2ede1d",
+#     "num_chunks": 0
+#   },
+#   "messages": [
+#     {
+#       "message_type": "control",
+#       "message_name": "close",
+#       "message_id": 4,
+#       "reason": null
+#     }
+#   ]
+# }
 ```
 
 More examples can be found in the [examples/](./examples/) folder:
diff --git a/tests/json_test.py b/tests/json_test.py
new file mode 100644
index 0000000..0946f87
--- /dev/null
+++ b/tests/json_test.py
@@ -0,0 +1,44 @@
+from typing import cast
+import textwrap
+
+from twnet_parser.messages6.game.cl_change_info import MsgClChangeInfo
+from twnet_parser.packet import TwPacket, NetMessage
+
+def test_change_info_packet() -> None:
+    packet = TwPacket()
+    packet.messages.append(cast(NetMessage, MsgClChangeInfo()))
+    packet = packet.to_json()
+    expected = textwrap.dedent("""\
+    {
+      "version": "0.7",
+      "payload_raw": "",
+      "payload_decompressed": "",
+      "header": {
+        "flags": [],
+        "ack": 0,
+        "token": "ffffffff",
+        "num_chunks": null
+      },
+      "messages": [
+        {
+          "message_type": "game",
+          "message_name": "cl_change_info",
+          "system_message": false,
+          "message_id": 21,
+          "header": {
+            "version": "0.6",
+            "flags": [],
+            "size": null,
+            "seq": -1
+          },
+          "name": "default",
+          "clan": "default",
+          "country": 0,
+          "skin": "default",
+          "use_custom_color": false,
+          "color_body": 0,
+          "color_feet": 0
+        }
+      ]
+    }""")
+    assert packet == expected
diff --git a/twnet_parser/packet.py b/twnet_parser/packet.py
index 6a3cbe9..6061e51 100644
--- a/twnet_parser/packet.py
+++ b/twnet_parser/packet.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
-from typing import Union, cast, Optional, Literal
+from typing import Union, cast, Optional, Literal, Iterator, Any
+import json
 
 from twnet_parser.packer import Unpacker, pack_int
 from twnet_parser.pretty_print import PrettyPrint
@@ -16,6 +17,7 @@ from twnet_parser.msg_matcher.connless6 import match_connless6
 from twnet_parser.constants import NET_MAX_SEQUENCE, NET_PACKETVERSION
 
 import twnet_parser.huffman
+import twnet_parser.serialize
 
 # TODO: what is a nice pythonic way of storing those?
 #       also does some version:: namespace thing make sense?
@@ -257,7 +259,7 @@ class TwPacket(PrettyPrint):
             raise ValueError(f"Error: invalid packet version '{self.version}'")
         self.messages: list[Union[CtrlMessage, NetMessage, ConnlessMessage]] = []
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[tuple[str, Any]]:
         yield 'version', self.version
         yield 'payload_raw', self.payload_raw
         yield 'payload_decompressed', self.payload_decompressed
@@ -265,6 +267,14 @@ class TwPacket(PrettyPrint):
 
         yield 'messages', [dict(msg) for msg in self.messages]
 
+    def to_json(self) -> str:
+        return json.dumps(
+            dict(self),
+            indent=2,
+            sort_keys=False,
+            default=twnet_parser.serialize.bytes_to_hex
+        )
+
     @property
     def version(self) -> Literal['0.6', '0.7']:
         return self._version
diff --git a/twnet_parser/serialize.py b/twnet_parser/serialize.py
new file mode 100644
index 0000000..4bb674f
--- /dev/null
+++ b/twnet_parser/serialize.py
@@ -0,0 +1,5 @@
+def bytes_to_hex(obj):
+    if isinstance(obj, bytes):
+        return obj.hex()
+    raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
+