""" Module for parsing binary data into structs. """ import argparse import re from collections import namedtuple from pyraknet.bitstream import BitStream, c_bit, c_float, c_double, c_int8, c_uint8, c_int16, c_uint16, c_int32, c_uint32, c_int64, c_uint64 VAR_CHARS = r"[^ \t\[\]]+" BITSTREAM_TYPES = {"bit": c_bit, "float": c_float, "double": c_double, "s8": c_int8, "u8": c_uint8, "s16": c_int16, "u16": c_uint16, "s32": c_int32, "u32": c_uint32, "s64": c_int64, "u64": c_uint64} TYPES_RE = "("+"|".join(BITSTREAM_TYPES.keys())+")" DEFINITION_SYNTAX = re.compile(r""" ^(?P\t*) # Indentation (if\ (?P.+): | while\ (?P.+): | (?Pbreak) | ((?P"""+VAR_CHARS+r""")=)? # Assign this struct a variable so the value can be back-referenced later \[ (?P"""+TYPES_RE+r""") # Struct type \] \ -\ (?P.*?) # Description for the struct (,\ expect\ (?P(.+?)))? # Expect the value to be like this expression. Struct attribute 'unexpected' will be None if no expects, True if any expects are False, or False if all expects are True. (,\ assert\ (?P(.+?)))? # Assert the value to be like this expression, will raise AssertionError if not True. )$ """, re.VERBOSE) IfStatement = namedtuple("IfStatement", ("condition",)) WhileStatement = namedtuple("WhileStatement", ("condition",)) BreakStatement = namedtuple("BreakStatement", ()) StructDefinition = namedtuple("struct_token", ("var_assign", "type", "description", "expects", "asserts")) Structure = namedtuple("Structure", ("level", "description", "value", "unexpected")) class StructParser: def __init__(self, struct_defs): """ Set up the parser with the structure definitions. Arguments: struct_defs: A string of structure definitions in my custom format (currently unnamed), see the documentation of that for details. """ self._variables = {} struct_defs = struct_defs.splitlines() struct_defs = [re.search(DEFINITION_SYNTAX, struct).groupdict() for struct in struct_defs if re.search(DEFINITION_SYNTAX, struct) is not None] # Filter out lines not matching the syntax self.defs = self._to_tree(iter(struct_defs))[0] def parse(self, data, variables=None): """ Parse the binary data, yielding structure objects. Arguments: data: The binary data to parse. variables: A dict of variables to be used in checks as defined by the structure definition, such as expects or asserts. Yields: Named structure tuples, attributes: level: The indentation level from the structure definition. description: The description from the structure definition. value: Parsed value of this structure occurrence in the binary data. The type of this is specified by the type specified in the structure definition. unexpected: None if no expects defined, True if any expects are False, False if all expects are True. Raises: AssertionError if any assert is False. """ if variables is None: variables = {} self._variables = variables if isinstance(data, BitStream): stream = data else: stream = BitStream(data) yield from self._parse_struct_occurrences(stream, self.defs) def _to_tree(self, def_iter, stack_level=0, start_def=None): current_level = [] try: if start_def is not None: def_ = start_def else: def_ = next(def_iter) while True: if len(def_["indent"]) == stack_level: def_tuple = self._to_def_tuple(def_) current_level.append((def_tuple, ())) def_ = next(def_iter) elif len(def_["indent"]) == stack_level+1: # found a child of the previous children, next_struct = self._to_tree(def_iter, stack_level+1, def_) current_level[-1] = current_level[-1][0], children if next_struct is None: raise StopIteration def_ = next_struct elif len(def_["indent"]) < stack_level: # we're at ancestor level again, done with the children return current_level, def_ except StopIteration: return current_level, None @staticmethod def _to_def_tuple(def_): if def_["if_condition"] is not None: condition = compile(def_["if_condition"], "", "eval") return IfStatement(condition) if def_["while_condition"] is not None: condition = compile(def_["while_condition"], "", "eval") return WhileStatement(condition) if def_["break"] is not None: return BreakStatement() type_ = BITSTREAM_TYPES[def_["type"]] if def_["expect"] is not None: expects = [compile("value "+i, "", "eval") for i in def_["expect"].split(" and ")] else: expects = () if def_["assert"] is not None: asserts = [compile("value "+i, "", "eval") for i in def_["assert"].split(" and ")] else: asserts = () return StructDefinition(def_["var_assign"], type_, def_["description"], expects, asserts) def _parse_struct_occurrences(self, stream, defs, stack_level=0, repeat_times=1): for _ in range(repeat_times): for def_, children in defs: if isinstance(def_, IfStatement): if children and self._eval(def_.condition): break_ = yield from self._parse_struct_occurrences(stream, children, stack_level+1) if break_: return True elif isinstance(def_, WhileStatement): if children: while self._eval(def_.condition): break_ = yield from self._parse_struct_occurrences(stream, children, stack_level+1) if break_: break elif isinstance(def_, BreakStatement): return True else: value = stream.read(def_.type) if def_.expects: for expression in def_.expects: if not self._eval(expression, value): unexpected = True break else: unexpected = False else: unexpected = None for expression in def_.asserts: assert self._eval(expression, value), (value, expression, def_) if def_.var_assign is not None: self._variables[def_.var_assign] = value yield Structure(stack_level, def_.description, value, unexpected) if children and value: break_ = yield from self._parse_struct_occurrences(stream, children, stack_level+1, value) if break_: return True def _eval(self, expression, value=None): globals_ = {"__builtins__": {}, "value": value} globals_.update(self._variables) return eval(expression, globals_) # definitely not safe, fwiw if __name__ == "__main__": argparser = argparse.ArgumentParser(description=__doc__) argparser.add_argument("filepath", help="path of binary file") argparser.add_argument("definition", help="struct definition file path to parse with") args = argparser.parse_args() with open(args.definition) as file: defs = file.read() parser = StructParser(defs) with open(args.filepath, "rb") as file: for structure in parser.parse(file.read()): print(structure)