Merge pull request #284 from fmonjalet/feature_memstruct

Feature MemStruct

2 files changed, 552 insertions, 0 deletions

diff --git a/test/core/types.py b/test/core/types.py
new file mode 100644
index 00000000..bb1d5da1
--- /dev/null
+++ b/test/core/types.py
@@ -0,0 +1,549 @@
+#!/usr/bin/env python
+
+# miasm2.core.types tests
+
+import struct
+
+from miasm2.analysis.machine import Machine
+from miasm2.core.types import MemStruct, Num, Ptr, Str, \
+                              Array, RawStruct, Union, \
+                              BitField, Self, Void, Bits, \
+                              set_allocator, MemUnion, Struct
+from miasm2.jitter.csts import PAGE_READ, PAGE_WRITE
+from miasm2.os_dep.common import heap
+
+# Two structures with some fields
+class OtherStruct(MemStruct):
+    fields = [
+        ("foo", Num("H")),
+    ]
+
+class MyStruct(MemStruct):
+    fields = [
+        # Number field: just struct.pack fields with one value
+        ("num", Num("I")),
+        ("flags", Num("B")),
+        # This field is a pointer to another struct, it has a numeric
+        # value (mystruct.other.val) and can be dereferenced to get an
+        # OtherStruct instance (mystruct.other.deref)
+        ("other", Ptr("I", OtherStruct)),
+        # Ptr to a variable length String
+        ("s", Ptr("I", Str())),
+        ("i", Ptr("I", Num("I"))),
+    ]
+
+jitter = Machine("x86_32").jitter("python")
+jitter.init_stack()
+addr = 0x1000
+size = 0x1000
+addr_str = 0x1100
+addr_str2 = 0x1200
+addr_str3 = 0x1300
+# Initialize all mem with 0xaa
+jitter.vm.add_memory_page(addr, PAGE_READ | PAGE_WRITE, "\xaa"*size)
+
+
+# MemStruct tests
+## Creation
+# Use manual allocation with explicit addr for the first example
+mstruct = MyStruct(jitter.vm, addr)
+## Fields are read from the virtual memory
+assert mstruct.num == 0xaaaaaaaa
+assert mstruct.flags == 0xaa
+
+## Field assignment modifies virtual memory
+mstruct.num = 3
+assert mstruct.num == 3
+memval = struct.unpack("I", jitter.vm.get_mem(mstruct.get_addr(), 4))[0]
+assert memval == 3
+
+## Memset sets the whole structure
+mstruct.memset()
+assert mstruct.num == 0
+assert mstruct.flags == 0
+assert mstruct.other.val == 0
+assert mstruct.s.val == 0
+assert mstruct.i.val == 0
+mstruct.memset('\x11')
+assert mstruct.num == 0x11111111
+assert mstruct.flags == 0x11
+assert mstruct.other.val == 0x11111111
+assert mstruct.s.val == 0x11111111
+assert mstruct.i.val == 0x11111111
+
+
+# From now, just use heap.vm_alloc
+my_heap = heap()
+set_allocator(my_heap.vm_alloc)
+
+
+# Ptr tests
+## Setup for Ptr tests
+# the addr field can now be omited since allocator is set
+other = OtherStruct(jitter.vm)
+other.foo = 0x1234
+assert other.foo == 0x1234
+
+## Basic usage
+mstruct.other.val = other.get_addr()
+# This also works for now:
+# mstruct.other = other.get_addr()
+assert mstruct.other.val == other.get_addr()
+assert mstruct.other.deref == other
+assert mstruct.other.deref.foo == 0x1234
+
+## Deref assignment
+other2 = OtherStruct(jitter.vm)
+other2.foo = 0xbeef
+assert mstruct.other.deref != other2
+mstruct.other.deref = other2
+assert mstruct.other.deref == other2
+assert mstruct.other.deref.foo == 0xbeef
+assert mstruct.other.val == other.get_addr() # Addr did not change
+assert other.foo == 0xbeef # Deref assignment copies by value
+assert other2.foo == 0xbeef
+assert other.get_addr() != other2.get_addr() # Not the same address
+assert other == other2 # But same value
+
+## Same stuff for Ptr to MemField
+alloc_addr = my_heap.vm_alloc(jitter.vm,
+                              mstruct.get_type().get_field_type("i")
+                                     .dst_type.sizeof())
+mstruct.i = alloc_addr
+mstruct.i.deref.val = 8
+assert mstruct.i.deref.val == 8
+assert mstruct.i.val == alloc_addr
+memval = struct.unpack("I", jitter.vm.get_mem(alloc_addr, 4))[0]
+assert memval == 8
+
+
+# Str tests
+## Basic tests
+memstr = Str().lval(jitter.vm, addr_str)
+memstr.val = ""
+assert memstr.val == ""
+assert jitter.vm.get_mem(memstr.get_addr(), 1) == '\x00'
+memstr.val = "lala"
+assert jitter.vm.get_mem(memstr.get_addr(), memstr.get_size()) == 'lala\x00'
+jitter.vm.set_mem(memstr.get_addr(), 'MIAMs\x00')
+assert memstr.val == 'MIAMs'
+
+## Ptr(Str()) manipulations
+mstruct.s.val = memstr.get_addr()
+assert mstruct.s.val == addr_str
+assert mstruct.s.deref == memstr
+assert mstruct.s.deref.val == 'MIAMs'
+mstruct.s.deref.val = "That's all folks!"
+assert mstruct.s.deref.val == "That's all folks!"
+assert memstr.val == "That's all folks!"
+
+## Other address, same value, same encoding
+memstr2 = Str().lval(jitter.vm, addr_str2)
+memstr2.val = "That's all folks!"
+assert memstr2.get_addr() != memstr.get_addr()
+assert memstr2 == memstr
+
+## Same value, other encoding
+memstr3 = Str("utf16").lval(jitter.vm, addr_str3)
+memstr3.val = "That's all folks!"
+assert memstr3.get_addr() != memstr.get_addr()
+assert memstr3.get_size() != memstr.get_size() # Size is different
+assert str(memstr3) != str(memstr) # Mem representation is different
+assert memstr3 != memstr # Encoding is different, so they are not eq
+assert memstr3.val == memstr.val # But the python value is the same
+
+
+# Array tests
+# Allocate buffer manually, since memarray is unsized
+alloc_addr = my_heap.vm_alloc(jitter.vm, 0x100)
+memarray = Array(Num("I")).lval(jitter.vm, alloc_addr)
+memarray[0] = 0x02
+assert memarray[0] == 0x02
+assert jitter.vm.get_mem(memarray.get_addr(),
+                         Num("I").size) == '\x02\x00\x00\x00'
+memarray[2] = 0xbbbbbbbb
+assert memarray[2] == 0xbbbbbbbb
+assert jitter.vm.get_mem(memarray.get_addr() + 2 * Num("I").size,
+                         Num("I").size) == '\xbb\xbb\xbb\xbb'
+try:
+    s = str(memarray)
+    assert False, "Should raise"
+except (NotImplementedError, ValueError):
+    pass
+try:
+    s = len(memarray)
+    assert False, "Should raise"
+except (NotImplementedError, ValueError):
+    pass
+
+## Slice assignment
+memarray[2:4] = [3, 3]
+assert memarray[2] == 3
+assert memarray[3] == 3
+assert memarray[2:4] == [3, 3]
+try:
+    memarray[2:4] = [3, 3, 3]
+    assert False, "Should raise, mismatched sizes"
+except ValueError:
+    pass
+
+
+memsarray = Array(Num("I"), 10).lval(jitter.vm)
+# And Array(type, size).lval generates statically sized types
+assert memsarray.sizeof() == Num("I").size * 10
+memsarray.memset('\xcc')
+assert memsarray[0] == 0xcccccccc
+assert len(memsarray) == 10 * 4
+assert str(memsarray) == '\xcc' * (4 * 10)
+for val in memsarray:
+    assert val == 0xcccccccc
+assert list(memsarray) == [0xcccccccc] * 10
+memsarray[0] = 2
+assert memsarray[0] == 2
+assert str(memsarray) == '\x02\x00\x00\x00' + '\xcc' * (4 * 9)
+
+
+# Atypical fields (RawStruct and Array)
+class MyStruct2(MemStruct):
+    fields = [
+        ("s1", RawStruct("=BI")),
+        ("s2", Array(Num("B"), 10)),
+    ]
+
+ms2 = MyStruct2(jitter.vm)
+ms2.memset('\xaa')
+assert len(ms2) == 15
+
+## RawStruct
+assert len(ms2.s1) == 2
+assert ms2.s1[0] == 0xaa
+assert ms2.s1[1] == 0xaaaaaaaa
+
+## Array
+### Basic checks
+assert len(ms2.s2) == 10
+for val in ms2.s2:
+    assert val == 0xaa
+assert ms2.s2[0] == 0xaa
+assert ms2.s2[9] == 0xaa
+
+### Subscript assignment
+ms2.s2[3] = 2
+assert ms2.s2[3] == 2
+
+### Field assignment (list)
+ms2.s2 = [1] * 10
+for val in ms2.s2:
+    assert val == 1
+
+### Field assignment (MemSizedArray)
+array2 = Array(Num("B"), 10).lval(jitter.vm)
+jitter.vm.set_mem(array2.get_addr(), '\x02'*10)
+for val in array2:
+    assert val == 2
+ms2.s2 = array2
+for val in ms2.s2:
+    assert val == 2
+
+
+# Inlining a MemType tests
+class InStruct(MemStruct):
+    fields = [
+        ("foo", Num("B")),
+        ("bar", Num("B")),
+    ]
+
+class ContStruct(MemStruct):
+    fields = [
+        ("one", Num("B")),
+        # Shorthand for: ("instruct", InStruct.get_type()),
+        ("instruct", InStruct),
+        ("last", Num("B")),
+    ]
+
+cont = ContStruct(jitter.vm)
+cont.memset()
+assert len(cont) == 4
+assert len(cont.instruct) == 2
+assert cont.one == 0
+assert cont.last == 0
+assert cont.instruct.foo == 0
+assert cont.instruct.bar == 0
+cont.memset('\x11')
+assert cont.one == 0x11
+assert cont.last == 0x11
+assert cont.instruct.foo == 0x11
+assert cont.instruct.bar == 0x11
+
+cont.one = 0x01
+cont.instruct.foo = 0x02
+cont.instruct.bar = 0x03
+cont.last = 0x04
+assert cont.one == 0x01
+assert cont.instruct.foo == 0x02
+assert cont.instruct.bar == 0x03
+assert cont.last == 0x04
+assert jitter.vm.get_mem(cont.get_addr(), len(cont)) == '\x01\x02\x03\x04'
+
+
+# Union test
+class UniStruct(MemStruct):
+    fields = [
+        ("one", Num("B")),
+        ("union", Union([
+            ("instruct", InStruct),
+            ("i", Num(">I")),
+        ])),
+        ("last", Num("B")),
+    ]
+
+uni = UniStruct(jitter.vm)
+jitter.vm.set_mem(uni.get_addr(), ''.join(chr(x) for x in xrange(len(uni))))
+assert len(uni) == 6 # 1 + max(InStruct.sizeof(), 4) + 1
+assert uni.one == 0x00
+assert uni.union.instruct.foo == 0x01
+assert uni.union.instruct.bar == 0x02
+assert uni.union.i == 0x01020304
+assert uni.last == 0x05
+uni.union.instruct.foo = 0x02
+assert uni.union.i == 0x02020304
+uni.union.i = 0x11223344
+assert uni.union.instruct.foo == 0x11
+assert uni.union.instruct.bar == 0x22
+
+
+# BitField test
+class BitStruct(MemUnion):
+    fields = [
+        ("flags_num", Num("H")),
+        ("flags", BitField(Num("H"), [
+            ("f1_1", 1),
+            ("f2_5", 5),
+            ("f3_8", 8),
+            ("f4_1", 1),
+        ])),
+    ]
+
+bit = BitStruct(jitter.vm)
+bit.memset()
+assert bit.flags_num == 0
+assert bit.flags.f1_1 == 0
+assert bit.flags.f2_5 == 0
+assert bit.flags.f3_8 == 0
+assert bit.flags.f4_1 == 0
+bit.flags.f1_1 = 1
+bit.flags.f2_5 = 0b10101
+bit.flags.f3_8 = 0b10000001
+assert bit.flags_num == 0b0010000001101011
+assert bit.flags.f1_1 == 1
+assert bit.flags.f2_5 == 0b10101
+assert bit.flags.f3_8 == 0b10000001
+assert bit.flags.f4_1 == 0
+bit.flags_num = 0b1101010101011100
+assert bit.flags.f1_1 == 0
+assert bit.flags.f2_5 == 0b01110
+assert bit.flags.f3_8 == 0b01010101
+assert bit.flags.f4_1 == 1
+
+
+# Unhealthy ideas
+class UnhealthyIdeas(MemStruct):
+    fields = [
+        ("pastruct", Ptr("I", Array(RawStruct("=Bf")))),
+        ("apstr", Array(Ptr("I", Str()), 10)),
+        ("pself", Ptr("I", Self())),
+        ("apself", Array(Ptr("I", Self()), 2)),
+        ("ppself", Ptr("I", Ptr("I", Self()))),
+        ("pppself", Ptr("I", Ptr("I", Ptr("I", Self())))),
+    ]
+
+p_size = Ptr("I", Void()).size
+
+ideas = UnhealthyIdeas(jitter.vm)
+ideas.memset()
+ideas.pself = ideas.get_addr()
+assert ideas == ideas.pself.deref
+
+ideas.apself[0] = ideas.get_addr()
+assert ideas.apself[0].deref == ideas
+ideas.apself[1] = my_heap.vm_alloc(jitter.vm, UnhealthyIdeas.sizeof())
+ideas.apself[1].deref = ideas
+assert ideas.apself[1] != ideas.get_addr()
+assert ideas.apself[1].deref == ideas
+
+ideas.ppself = my_heap.vm_alloc(jitter.vm, p_size)
+ideas.ppself.deref.val = ideas.get_addr()
+assert ideas.ppself.deref.val == ideas.get_addr()
+assert ideas.ppself.deref.deref == ideas
+
+ideas.ppself.deref.val = my_heap.vm_alloc(jitter.vm, UnhealthyIdeas.sizeof())
+ideas.ppself.deref.deref = ideas
+assert ideas.ppself.deref.val != ideas.get_addr()
+assert ideas.ppself.deref.deref == ideas
+
+ideas.pppself = my_heap.vm_alloc(jitter.vm, p_size)
+ideas.pppself.deref.val = my_heap.vm_alloc(jitter.vm, p_size)
+ideas.pppself.deref.deref.val = ideas.get_addr()
+assert ideas.pppself.deref.deref.deref == ideas
+
+
+# Circular dependencies
+class A(MemStruct):
+    pass
+
+class B(MemStruct):
+    fields = [("a", Ptr("I", A)),]
+
+# Gen A's fields after declaration
+A.gen_fields([("b", Ptr("I", B)),])
+
+a = A(jitter.vm)
+b = B(jitter.vm)
+a.b.val = b.get_addr()
+b.a.val = a.get_addr()
+assert a.b.deref == b
+assert b.a.deref == a
+
+
+# Cast tests
+# MemStruct cast
+MemInt = Num("I").lval
+MemShort = Num("H").lval
+dword = MemInt(jitter.vm)
+dword.val = 0x12345678
+assert isinstance(dword.cast(MemShort), MemShort)
+assert dword.cast(MemShort).val == 0x5678
+
+# Field cast
+ms2.s2[0] = 0x34
+ms2.s2[1] = 0x12
+assert ms2.cast_field("s2", MemShort).val == 0x1234
+
+# Other method
+assert MemShort(jitter.vm, ms2.get_addr("s2")).val == 0x1234
+
+# Manual cast inside an Array
+ms2.s2[4] = 0xcd
+ms2.s2[5] = 0xab
+assert MemShort(jitter.vm, ms2.s2.get_addr(4)).val == 0xabcd
+
+# void* style cast
+MemPtrVoid = Ptr("I", Void()).lval
+p = MemPtrVoid(jitter.vm)
+p.val = mstruct.get_addr()
+assert p.deref.cast(MyStruct) == mstruct
+assert p.cast(Ptr("I", MyStruct)).deref == mstruct
+
+# Field equality tests
+assert RawStruct("IH") == RawStruct("IH")
+assert RawStruct("I") != RawStruct("IH")
+assert Num("I") == Num("I")
+assert Num(">I") != Num("<I")
+assert Ptr("I", MyStruct) == Ptr("I", MyStruct)
+assert Ptr(">I", MyStruct) != Ptr("<I", MyStruct)
+assert Ptr("I", MyStruct) != Ptr("I", MyStruct2)
+assert MyStruct.get_type() == MyStruct.get_type()
+assert MyStruct.get_type() != MyStruct2.get_type()
+assert Array(Num("H"), 12) == Array(Num("H"), 12)
+assert Array(Num("H"), 11) != Array(Num("H"), 12)
+assert Array(Num("I"), 12) != Array(Num("H"), 12)
+assert Struct("a", [("f1", Num("B")), ("f2", Num("H"))]) == \
+        Struct("a", [("f1", Num("B")), ("f2", Num("H"))])
+assert Struct("a", [("f2", Num("B")), ("f2", Num("H"))]) != \
+        Struct("a", [("f1", Num("B")), ("f2", Num("H"))])
+assert Struct("a", [("f1", Num("B")), ("f2", Num("H"))]) != \
+        Struct("a", [("f1", Num("I")), ("f2", Num("H"))])
+assert Struct("a", [("f1", Num("B")), ("f2", Num("H"))]) != \
+        Struct("b", [("f1", Num("B")), ("f2", Num("H"))])
+assert Union([("f1", Num("B")), ("f2", Num("H"))]) == \
+        Union([("f1", Num("B")), ("f2", Num("H"))])
+assert Union([("f2", Num("B")), ("f2", Num("H"))]) != \
+        Union([("f1", Num("B")), ("f2", Num("H"))])
+assert Union([("f1", Num("B")), ("f2", Num("H"))]) != \
+        Union([("f1", Num("I")), ("f2", Num("H"))])
+assert Bits(Num("I"), 3, 8) == Bits(Num("I"), 3, 8)
+assert Bits(Num("I"), 3, 8) != Bits(Num("I"), 3, 8)
+assert Bits(Num("H"), 2, 8) != Bits(Num("I"), 3, 8)
+assert Bits(Num("I"), 3, 7) != Bits(Num("I"), 3, 8)
+assert BitField(Num("B"), [("f1", 2), ("f2", 4), ("f3", 1)]) == \
+        BitField(Num("B"), [("f1", 2), ("f2", 4), ("f3", 1)])
+assert BitField(Num("H"), [("f1", 2), ("f2", 4), ("f3", 1)]) != \
+        BitField(Num("B"), [("f1", 2), ("f2", 4), ("f3", 1)])
+assert BitField(Num("B"), [("f2", 2), ("f2", 4), ("f3", 1)]) != \
+        BitField(Num("B"), [("f1", 2), ("f2", 4), ("f3", 1)])
+assert BitField(Num("B"), [("f1", 1), ("f2", 4), ("f3", 1)]) != \
+        BitField(Num("B"), [("f1", 2), ("f2", 4), ("f3", 1)])
+
+
+# Quick MemField.lval/MemField hash test
+assert Num("f").lval(jitter.vm, addr) == Num("f").lval(jitter.vm, addr)
+# Types are cached
+assert Num("f").lval == Num("f").lval
+assert Num("d").lval != Num("f").lval
+assert Union([("f1", Num("I")), ("f2", Num("H"))]).lval == \
+        Union([("f1", Num("I")), ("f2", Num("H"))]).lval
+assert Array(Num("B")).lval == Array(Num("B")).lval
+assert Array(Num("I")).lval != Array(Num("B")).lval
+assert Array(Num("B"), 20).lval == Array(Num("B"), 20).lval
+assert Array(Num("B"), 19).lval != Array(Num("B"), 20).lval
+
+# MemStruct unicity test
+assert MyStruct == Struct(MyStruct.__name__, MyStruct.fields).lval
+assert MyStruct.get_type() == Struct(MyStruct.__name__, MyStruct.fields)
+
+# Anonymous Unions
+class Anon(MemStruct):
+    fields = [
+        ("a", Num("B")),
+        # If a field name evaluates to False ("" or None for example) and the
+        # field type is a Struct subclass (Struct, Union, BitField), the field
+        # is considered as an anonymous struct or union. Therefore, Anon will
+        # have b1, b2 and c1, c2 attributes in that case.
+        ("", Union([("b1", Num("B")), ("b2", Num("H"))])),
+        ("", Struct("", [("c1", Num("B")), ("c2", Num("B"))])),
+        ("d", Num("B")),
+    ]
+
+anon = Anon(jitter.vm)
+anon.memset()
+anon.a = 0x07
+anon.b2 = 0x0201
+anon.c1 = 0x55
+anon.c2 = 0x77
+anon.d = 0x33
+assert anon.a == 0x07
+assert anon.b1 == 0x01
+assert anon.b2 == 0x0201
+assert anon.c1 == 0x55
+assert anon.c2 == 0x77
+assert anon.d == 0x33
+
+# get_offset
+for field, off in (("a", 0), ("b1", 1), ("b2", 1), ("c1", 3), ("c2", 4),
+                   ("d", 5)):
+    assert Anon.get_offset(field) == Anon.get_type().get_offset(field)
+    assert Anon.get_offset(field) == off
+
+arr_t = Array(Num("H"))
+for idx, off in ((0, 0), (1, 2), (30, 60)):
+    assert arr_t.get_offset(idx) == arr_t.lval.get_offset(idx)
+    assert arr_t.get_offset(idx) == off
+
+
+# Repr tests
+
+print "Some struct reprs:\n"
+print repr(mstruct), '\n'
+print repr(ms2), '\n'
+print repr(cont), '\n'
+print repr(uni), '\n'
+print repr(bit), '\n'
+print repr(ideas), '\n'
+print repr(Array(MyStruct2.get_type(), 2).lval(jitter.vm, addr)), '\n'
+print repr(Num("f").lval(jitter.vm, addr)), '\n'
+print repr(memarray)
+print repr(memsarray)
+print repr(memstr)
+print repr(memstr3)
+
+print "\nOk" # That's all folks!
diff --git a/test/test_all.py b/test/test_all.py
index bc019104..6a619948 100644
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -191,6 +191,7 @@ for script in ["interval.py",
                "parse_asm.py",
                "utils.py",
                "sembuilder.py",
+               "types.py",
                ]:
     testset += RegressionTest([script], base_dir="core")
 ## Expression
@@ -482,6 +483,8 @@ for script, dep in [(["x86_32.py", Example.get_sample("x86_32_sc.bin")], []),
         testset += ExampleJitter(script + ["--jitter", jitter], depends=dep,
                                  tags=tags)
 
+testset += ExampleJitter(["types.py"])
+
 
 if __name__ == "__main__":
     # Argument parsing