diff --git a/scripts/generate_messages.py b/scripts/generate_messages.py index daafc12..3010023 100755 --- a/scripts/generate_messages.py +++ b/scripts/generate_messages.py @@ -53,6 +53,9 @@ class NetMessageMemberTypeJson(TypedDict): count: int member_type: ArrayMemberTypeJson + # data + size: Literal['specified_before'] + class NetMessageMemberJson(TypedDict): name: list[str] type: NetMessageMemberTypeJson @@ -172,9 +175,16 @@ def generate_msg(msg: NetMessageJson, game: Literal['game', 'system']) -> None: ftype = 'str' default = "'default'" elif member['type']['kind'] in \ - ('raw', 'sha256', 'rest', 'data'): # TODO: data has a size field + ('raw', 'sha256', 'rest'): # TODO: rest sha256 and raw ftype = 'bytes' default = "b'\\x00'" + elif member['type']['kind'] == 'data': + ftype = 'bytes' + default = "b'\\x00'" + if member['type']['size'] == 'specified_before': + args.append(' data_size: Optional[int] = None') + else: + raise ValueError(f"Error: unknown data size {member['type']}") # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}}, elif member['type']['kind'] == 'enum': ftype = 'int' @@ -248,8 +258,15 @@ def generate_msg(msg: NetMessageJson, game: Literal['game', 'system']) -> None: if member['type']['kind'] == 'string': ftype = 'str' elif member['type']['kind'] in \ - ('raw', 'sha256', 'rest', 'data'): # TODO: data has a size field + ('raw', 'sha256', 'rest'): # TODO: sha256 and raw ftype = 'bytes' + elif member['type']['kind'] == 'data': + ftype = 'bytes' + if member['type']['size'] == 'specified_before': + out_file.write(" " \ + "self.data_size: int = data_size if data_size else len(data)\n") + else: + raise ValueError(f"Error: unknown data size {member['type']}") # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}}, elif member['type']['kind'] == 'enum': ftype = 'int' @@ -307,8 +324,14 @@ def gen_unpack_members(msg: NetMessageJson) -> str: else: unpacker = 'str()' elif member['type']['kind'] in \ - ('raw', 'sha256', 'rest', 'data'): # TODO: data has a size field + ('raw', 'sha256', 'rest'): # TODO: do we need to fix size for sha256? unpacker = 'raw()' + elif member['type']['kind'] == 'data': + if member['type']['size'] == 'specified_before': + res += ' self.data_size = unpacker.get_int()\n' + unpacker = 'raw(self.data_size)' + else: + raise ValueError(f"Error: unknown data size {member['type']}") # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}}, elif member['type']['kind'] == 'enum': unpacker = 'int() # TODO: this is a enum' @@ -373,8 +396,13 @@ def get_dependencies(msg: NetMessageJson) -> str: if member['type']['disallow_cc']: packer_deps.append('SANITIZE_CC') elif member['type']['kind'] in \ - ('raw', 'sha256', 'rest', 'data'): # TODO: data has a size field + ('raw', 'sha256', 'rest'): pass + elif member['type']['kind'] == 'data': + if member['type']['size'] == 'specified_before': + typing_deps.append('Optional') + else: + raise ValueError(f"Error: unknown data size {member['type']}") # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}}, elif member['type']['kind'] == 'enum': packer_deps.append('pack_int') @@ -430,8 +458,14 @@ def pack_field(member: NetMessageMemberJson) -> str: if member['type']['kind'] == 'string': packer = 'str' elif member['type']['kind'] in \ - ('raw', 'sha256', 'rest', 'data'): # TODO: data has a size field + ('raw', 'sha256', 'rest'): # TODO: raw sha256 rest return f'self.{name}' + elif member['type']['kind'] == 'data': + if member['type']['size'] == 'specified_before': + return f'pack_int(self.data_size) + \\\n' \ + f' self.{name}' + else: + raise ValueError(f"Error: unknown data size {member['type']}") # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}}, elif member['type']['kind'] == 'enum': packer = 'int' diff --git a/twnet_parser/messages7/system/snap.py b/twnet_parser/messages7/system/snap.py index 29fbe60..4475617 100644 --- a/twnet_parser/messages7/system/snap.py +++ b/twnet_parser/messages7/system/snap.py @@ -4,6 +4,7 @@ from twnet_parser.pretty_print import PrettyPrint from twnet_parser.packer import Unpacker from twnet_parser.chunk_header import ChunkHeader from twnet_parser.packer import pack_int +from typing import Optional class MsgSnap(PrettyPrint): def __init__( @@ -13,6 +14,7 @@ class MsgSnap(PrettyPrint): num_parts: int = 0, part: int = 0, crc: int = 0, + data_size: Optional[int] = None, data: bytes = b'\x00' ) -> None: self.message_name = 'snap' @@ -24,6 +26,7 @@ class MsgSnap(PrettyPrint): self.num_parts: int = num_parts self.part: int = part self.crc: int = crc + self.data_size: int = data_size if data_size else len(data) self.data: bytes = data # first byte of data @@ -36,7 +39,8 @@ class MsgSnap(PrettyPrint): self.num_parts = unpacker.get_int() self.part = unpacker.get_int() self.crc = unpacker.get_int() - self.data = unpacker.get_raw() + self.data_size = unpacker.get_int() + self.data = unpacker.get_raw(self.data_size) return True def pack(self) -> bytes: @@ -45,4 +49,5 @@ class MsgSnap(PrettyPrint): pack_int(self.num_parts) + \ pack_int(self.part) + \ pack_int(self.crc) + \ + pack_int(self.data_size) + \ self.data \ No newline at end of file diff --git a/twnet_parser/messages7/system/snap_single.py b/twnet_parser/messages7/system/snap_single.py index cfa197a..196481e 100644 --- a/twnet_parser/messages7/system/snap_single.py +++ b/twnet_parser/messages7/system/snap_single.py @@ -4,6 +4,7 @@ from twnet_parser.pretty_print import PrettyPrint from twnet_parser.packer import Unpacker from twnet_parser.chunk_header import ChunkHeader from twnet_parser.packer import pack_int +from typing import Optional class MsgSnapSingle(PrettyPrint): def __init__( @@ -11,6 +12,7 @@ class MsgSnapSingle(PrettyPrint): tick: int = 0, delta_tick: int = 0, crc: int = 0, + data_size: Optional[int] = None, data: bytes = b'\x00' ) -> None: self.message_name = 'snap_single' @@ -20,6 +22,7 @@ class MsgSnapSingle(PrettyPrint): self.tick: int = tick self.delta_tick: int = delta_tick self.crc: int = crc + self.data_size: int = data_size if data_size else len(data) self.data: bytes = data # first byte of data @@ -30,11 +33,13 @@ class MsgSnapSingle(PrettyPrint): self.tick = unpacker.get_int() self.delta_tick = unpacker.get_int() self.crc = unpacker.get_int() - self.data = unpacker.get_raw() + self.data_size = unpacker.get_int() + self.data = unpacker.get_raw(self.data_size) return True def pack(self) -> bytes: return pack_int(self.tick) + \ pack_int(self.delta_tick) + \ pack_int(self.crc) + \ + pack_int(self.data_size) + \ self.data \ No newline at end of file