From 79e863bf77c4a18a8fabbf13931c74a3fdd86628 Mon Sep 17 00:00:00 2001
From: Quentin Kaiser <quentin.kaiser@onekey.com>
Date: Sun, 2 Apr 2023 15:47:05 +0200
Subject: [PATCH] Make the Instance class dict compatible to allow for JSON
 serialization.

---
 dissect/cstruct/types/enum.py     |  4 ++-
 dissect/cstruct/types/instance.py |  6 +++-
 examples/mirai_json.py            | 53 +++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 examples/mirai_json.py

diff --git a/dissect/cstruct/types/enum.py b/dissect/cstruct/types/enum.py
index e4fc9e3..ada6af7 100644
--- a/dissect/cstruct/types/enum.py
+++ b/dissect/cstruct/types/enum.py
@@ -84,12 +84,14 @@ def default_array(self, count: int) -> List[EnumInstance]:
         return [self.default() for _ in range(count)]
 
 
-class EnumInstance:
+class EnumInstance(dict):
     """Implements a value instance of an Enum"""
 
     def __init__(self, enum: Enum, value: int):
         self.enum = enum
         self.value = value
+        kwargs = {self.enum.name : self.value}
+        dict.__init__(self, **kwargs)
 
     def __eq__(self, value: Union[int, EnumInstance]) -> bool:
         if isinstance(value, EnumInstance) and value.enum is not self.enum:
diff --git a/dissect/cstruct/types/instance.py b/dissect/cstruct/types/instance.py
index 0d144f4..2297056 100644
--- a/dissect/cstruct/types/instance.py
+++ b/dissect/cstruct/types/instance.py
@@ -4,12 +4,13 @@
 from dissect.cstruct.types import BaseType
 
 
-class Instance:
+class Instance(dict):
     """Holds parsed structure data."""
 
     __slots__ = ("_type", "_values", "_sizes")
 
     def __init__(self, type_: BaseType, values: Dict[str, Any], sizes: Dict[str, int] = None):
+        dict.__init__(self, **values)
         # Done in this manner to check if the attr is in the lookup
         object.__setattr__(self, "_type", type_)
         object.__setattr__(self, "_values", values)
@@ -46,6 +47,9 @@ def __bytes__(self) -> bytes:
     def _size(self, field: str) -> int:
         return self._sizes[field]
 
+    def __iter__(self):
+        yield from self._values.items()
+
     def write(self, stream: BinaryIO) -> int:
         """Write this structure to a writable file-like object.
 
diff --git a/examples/mirai_json.py b/examples/mirai_json.py
new file mode 100644
index 0000000..602c8a2
--- /dev/null
+++ b/examples/mirai_json.py
@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+from dissect.cstruct import cstruct, dumpstruct
+
+import json
+import socket
+import struct
+
+class CustomEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, bytes):
+            return obj.decode('utf-8', errors='surrogateescape')
+        return json.JSONEncoder.default(self, obj)
+
+protocol = cstruct()
+
+protocol.load(
+    """
+enum AttackType : uint8 {
+    ATK_OPT_DPORT =  7,
+    ATK_OPT_DOMAIN  =  8,
+    ATK_OPT_NUM_SOCKETS =  24,
+};
+
+struct AttackTarget {
+    DWORD   ipv4;
+    BYTE    netmask;
+};
+
+struct AttackOption {
+    AttackType   type;
+    uint8   value_length;
+    char    value[value_length];
+};
+
+struct MiraiAttack {
+    uint16  total_length;
+    uint32  duration;
+    uint8   attack_id;
+    uint8   target_count;
+    AttackTarget targets[target_count];
+    uint8   num_opts;
+    AttackOption attack_options[num_opts];
+};
+"""
+)
+
+protocol.endian = ">"
+
+if __name__ == "__main__":
+    data = b"\x000\x00\x00\x00d\n\x01\x08\x08\x08\x08 \x03\x08\x16http://www.example.com\x07\x0280\x18\x045000"
+
+    record = protocol.MiraiAttack(data)
+    print(json.dumps(record, cls=CustomEncoder))
\ No newline at end of file