Skip to content

Commit

Permalink
Make the Instance class dict compatible to allow for JSON serialization.
Browse files Browse the repository at this point in the history
  • Loading branch information
qkaiser committed Apr 2, 2023
1 parent 5b4e10a commit 54d283f
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 2 deletions.
4 changes: 3 additions & 1 deletion dissect/cstruct/types/enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
6 changes: 5 additions & 1 deletion dissect/cstruct/types/instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand Down
53 changes: 53 additions & 0 deletions examples/mirai_json.py
Original file line number Diff line number Diff line change
@@ -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))

0 comments on commit 54d283f

Please sign in to comment.