Skip to content

Commit

Permalink
Merge pull request #95 from d0c-s4vage/hotfix/26-improper_structure_s…
Browse files Browse the repository at this point in the history
…upport

Fix structure and typedef support
  • Loading branch information
d0c-s4vage authored Jan 2, 2020
2 parents 4322633 + e9b2f1f commit 4ec0ae2
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 1 deletion.
57 changes: 56 additions & 1 deletion pfp/interp.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,49 @@ def _pfp__init(self, stream):
return new_class


def StructUnionTypeRef(curr_scope, typedef_name, refd_name, interp, node):
"""Create a typedef that resolves itself dynamically. This is needed in
situations like:
.. code-block:: c
struct MY_STRUCT {
char magic[4];
unsigned int filesize;
};
typedef struct MY_STRUCT ME;
LittleEndian();
ME s;
The typedef ``ME`` is handled before the ``MY_STRUCT`` declaration actually
occurs. The typedef value for ``ME`` should not the empty struct that is
resolved, but should be a dynamically-looked up struct definition when
a ``ME`` instance is actually declared.
"""
if isinstance(node, AST.Struct):
cls = fields.Struct
elif isinstance(node, AST.Union):
cls = fields.Union

def __new__(self, *args, **kwargs):
refd_type = curr_scope.get_type(refd_name)
if refd_type is None:
refd_node = node
else:
refd_node = refd_type._pfp__node
return StructUnionDef(typedef_name, interp, refd_node)(*args, **kwargs)

new_class = type(
typedef_name,
(cls,),
{
"__new__": __new__,
},
)
return new_class



def StructUnionDef(typedef_name, interp, node):
if isinstance(node, AST.Struct):
cls = fields.Struct
Expand Down Expand Up @@ -413,6 +456,15 @@ def add_type_class(self, name, cls):
"""
self._curr_scope["types"][name] = cls

def add_refd_struct_or_union(self, name, refd_name, interp, node):
"""Add a lazily-looked up typedef struct or union
:name: name of the typedefd struct/union
:node: the typedef node
:interp: the 010 interpreter
"""
self.add_type_class(name, StructUnionTypeRef(self, name, refd_name, interp, node))

def add_type_struct_or_union(self, name, interp, node):
"""Store the node with the name. When it is instantiated,
the node itself will be handled.
Expand Down Expand Up @@ -1624,7 +1676,10 @@ def _handle_typedef(self, node, scope, ctxt, stream):

if is_union_or_struct:
self._dlog("handling typedef struct/union '{}'".format(node.name))
scope.add_type_struct_or_union(node.name, self, node.type.type)
if node.type.type.name is None:
scope.add_type_struct_or_union(node.name, self, node.type.type)
else:
scope.add_refd_struct_or_union(node.name, node.type.type.name, self, node.type.type)
elif is_enum:
enum_cls = self._handle_node(node.type, scope, ctxt, stream)
scope.add_type_class(node.name, enum_cls)
Expand Down
77 changes: 77 additions & 0 deletions tests/test_struct_union.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,83 @@ def setUp(self):

def tearDown(self):
pass

def test_struct_vit9696_1(self):
dom = self._test_parse_build(
"\x00\x01\x02\x03\x00\x01\x02\x03",
"""
typedef struct MY_STRUCT_S {
char magic[4];
unsigned int filesize;
} MY_STRUCT;
LittleEndian();
MY_STRUCT s;
"""
)
assert dom.s.magic == "\x00\x01\x02\x03"
assert dom.s.filesize == 0x03020100

def test_struct_vit9696_2(self):
dom = self._test_parse_build(
"\x00\x01\x02\x03\x00\x01\x02\x03",
"""
struct MY_STRUCT {
char magic[4];
unsigned int filesize;
};
LittleEndian();
MY_STRUCT s;
"""
)
assert dom.s.magic == "\x00\x01\x02\x03"
assert dom.s.filesize == 0x03020100

def test_struct_vit9696_3(self):
dom = self._test_parse_build(
"\x00\x01\x02\x03\x00\x01\x02\x03",
"""
typedef struct MY_STRUCT {
char magic[4];
unsigned int filesize;
};
LittleEndian();
MY_STRUCT s;
"""
)
assert dom.s.magic == "\x00\x01\x02\x03"
assert dom.s.filesize == 0x03020100

def test_struct_vit9696_4(self):
dom = self._test_parse_build(
"\x00\x01\x02\x03\x00\x01\x02\x03",
"""
struct MY_STRUCT {
char magic[4];
unsigned int filesize;
};
LittleEndian();
struct MY_STRUCT s;
"""
)
assert dom.s.magic == "\x00\x01\x02\x03"
assert dom.s.filesize == 0x03020100

def test_struct_vit9696_5(self):
dom = self._test_parse_build(
"\x00\x01\x02\x03\x00\x01\x02\x03",
"""
struct MY_STRUCT {
char magic[4];
unsigned int filesize;
};
typedef struct MY_STRUCT ME;
LittleEndian();
ME s;
""",
debug=True,
)
assert dom.s.magic == "\x00\x01\x02\x03"
assert dom.s.filesize == 0x03020100

def test_basic_struct(self):
dom = self._test_parse_build(
Expand Down

0 comments on commit 4ec0ae2

Please sign in to comment.