MemStruct: Str type

3 files changed, 86 insertions, 67 deletions

diff --git a/example/jitter/memstruct.py b/example/jitter/memstruct.py
index 038622ba..5472798d 100644
--- a/example/jitter/memstruct.py
+++ b/example/jitter/memstruct.py
@@ -6,8 +6,8 @@ as well.
 """
 
 from miasm2.analysis.machine import Machine
-from miasm2.analysis.mem import PinnedStruct, PinnedSelf, PinnedVoid, PinnedStr,\
-                                PinnedSizedArray, Ptr, Num, Array, set_allocator
+from miasm2.analysis.mem import PinnedStruct, Self, Void, Str, Array, Ptr, \
+                                Num, Array, set_allocator
 from miasm2.os_dep.common import heap
 
 # Instanciate a heap
@@ -29,10 +29,10 @@ class ListNode(PinnedStruct):
         # special marker PinnedSelf.
         # You could also set or modify ListNode.fields after the class
         # declaration and call ListNode.gen_fields()
-        ("next", Ptr("<I", PinnedSelf)),
+        ("next", Ptr("<I", Self())),
         # Ptr(_, PinnedVoid) is analogous to void*, PinnedVoid is just an empty
         # PinnedStruct type
-        ("data", Ptr("<I", PinnedVoid)),
+        ("data", Ptr("<I", Void())),
     ]
 
     def get_next(self):
@@ -135,7 +135,7 @@ class DataStr(PinnedStruct):
     fields = [
         ("valshort", Num("H")),
         # Pointer to an utf16 null terminated string
-        ("data", Ptr("<I", PinnedStr, "utf16")),
+        ("data", Ptr("<I", Str("utf16"))),
     ]
 
 
@@ -152,7 +152,7 @@ vm = jitter.vm
 # Auto-allocated by my_heap. If you allocate memory at `addr`,
 # `link = LinkedList(vm, addr)` will use this allocation.
 link = LinkedList(vm)
-# Pinnedset the struct (with '\x00' by default)
+# memset the struct (with '\x00' by default)
 link.memset()
 
 # Push three uninitialized structures
@@ -203,7 +203,7 @@ assert data.val1 == 0x22 and data.val2 == 0x11
 
 # Let's play with strings
 memstr = datastr.data.deref
-# Note that memstr is PinnedStr(..., "utf16")
+# Note that memstr is Str("utf16")
 memstr.val = 'Miams'
 
 print "Cast data.array to PinnedStr and set the string value:"
@@ -214,10 +214,10 @@ print
 raw_miams = '\x00'.join('Miams') + '\x00'*3
 raw_miams_array = [ord(c) for c in raw_miams]
 assert list(data.array)[:len(raw_miams_array)] == raw_miams_array
-assert data.array.cast(PinnedStr, "utf16") == memstr
+assert data.array.cast(Str("utf16").pinned) == memstr
 # Default is "ansi"
-assert data.array.cast(PinnedStr) != memstr
-assert data.array.cast(PinnedStr, "utf16").val == memstr.val
+assert data.array.cast(Str().pinned) != memstr
+assert data.array.cast(Str("utf16").pinned).val == memstr.val
 
 print "See that the original array has been modified:"
 print repr(data)
diff --git a/miasm2/analysis/mem.py b/miasm2/analysis/mem.py
index 7f5373c1..9787a25e 100644
--- a/miasm2/analysis/mem.py
+++ b/miasm2/analysis/mem.py
@@ -208,7 +208,6 @@ class Type(object):
         interact with this type in memory.
         @return: a PinnedType subclass.
         """
-        # FIXME: PinnedStruct field changes
         if self in DYN_MEM_STRUCT_CACHE:
             return DYN_MEM_STRUCT_CACHE[self]
         pinned_type = self._build_pinned_type()
@@ -796,9 +795,53 @@ class BitField(Union):
         return hash((super(BitField, self).__hash__(), self._num))
 
 
-# TODO
 class Str(Type):
-    pass
+    def __init__(self, encoding="ansi"):
+        # TODO: encoding as lambda
+        if encoding not in ["ansi", "utf16"]:
+            raise NotImplementedError("Only 'ansi' and 'utf16' are implemented")
+        self._enc = encoding
+
+    def get(self, vm, addr):
+        """Set the string value in memory"""
+        if self._enc == "ansi":
+            get_str = get_str_ansi
+        elif self._enc == "utf16":
+            get_str = get_str_utf16
+        else:
+            raise NotImplementedError("Only 'ansi' and 'utf16' are implemented")
+        return get_str(vm, addr)
+
+    def set(self, vm, addr, s):
+        """Get the string value from memory"""
+        if self._enc == "ansi":
+            set_str = set_str_ansi
+        elif self._enc == "utf16":
+            set_str = set_str_utf16
+        else:
+            raise NotImplementedError("Only 'ansi' and 'utf16' are implemented")
+        set_str(vm, addr, s)
+
+    def size(self):
+        """This type is unsized."""
+        raise ValueError("Str is unsized")
+
+    @property
+    def enc(self):
+        return self._enc
+
+    def _get_pinned_base_class(self):
+        return PinnedStr
+
+    def __repr__(self):
+        return "%s(%s)" % (self.__class__.__name__, self.enc)
+
+    def __eq__(self, other):
+        return self.__class__ == other.__class__ and self._enc == other._enc
+
+    def __hash__(self):
+        return hash((self.__class__, self._enc))
+
 
 class Void(Type):
     """Represents the C void type."""
@@ -812,6 +855,10 @@ class Void(Type):
     def __hash__(self):
         return hash(self.__class__)
 
+class Self(Void):
+    def _build_pinned_type(self):
+        return PinnedSelf
+
 
 # PinnedType classes
 
@@ -946,7 +993,10 @@ class PinnedValue(PinnedType):
 
     @val.setter
     def val(self, value):
-        return self._type.set(self._vm, self._addr, value)
+        self._type.set(self._vm, self._addr, value)
+
+    def __repr__(self):
+        return "%r: %r" % (self.__class__, self.val)
 
 
 class PinnedStruct(PinnedType):
@@ -1138,7 +1188,6 @@ class PinnedVoid(PinnedType):
 
 
 class PinnedPtr(PinnedValue):
-
     @property
     def val(self):
         return self._type.get_val(self._vm, self._addr)
@@ -1159,10 +1208,9 @@ class PinnedPtr(PinnedValue):
         return "*%s" % hex(self.val)
 
 
-
 # This does not use _MetaPinnedStruct features, impl is custom for strings,
 # because they are unsized. The only memory field is self.value.
-class PinnedStr(PinnedType):
+class PinnedStr(PinnedValue):
     """Implements a string representation in memory.
 
     The @encoding is passed to the constructor, and is currently either null
@@ -1174,35 +1222,6 @@ class PinnedStr(PinnedType):
 
     This type is dynamically sized only (get_size is implemented, not sizeof).
     """
-    def __init__(self, vm, addr, encoding="ansi"):
-        # TODO: encoding as lambda
-        if encoding not in ["ansi", "utf16"]:
-            raise NotImplementedError("Only 'ansi' and 'utf16' are implemented")
-        # FIXME: Str type
-        super(PinnedStr, self).__init__(vm, addr, Str())
-        self._enc = encoding
-
-    @property
-    def val(self):
-        """Set the string value in memory"""
-        if self._enc == "ansi":
-            get_str = get_str_ansi
-        elif self._enc == "utf16":
-            get_str = get_str_utf16
-        else:
-            raise NotImplementedError("Only 'ansi' and 'utf16' are implemented")
-        return get_str(self._vm, self.get_addr())
-
-    @val.setter
-    def val(self, s):
-        """Get the string value from memory"""
-        if self._enc == "ansi":
-            set_str = set_str_ansi
-        elif self._enc == "utf16":
-            set_str = set_str_utf16
-        else:
-            raise NotImplementedError("Only 'ansi' and 'utf16' are implemented")
-        set_str(self._vm, self.get_addr(), s)
 
     def get_size(self):
         """This get_size implementation is quite unsafe: it reads the string
@@ -1210,9 +1229,9 @@ class PinnedStr(PinnedType):
         and provoke mem faults (analogous to strlen).
         """
         val = self.val
-        if self._enc == "ansi":
+        if self.get_type().enc == "ansi":
             return len(val) + 1
-        elif self._enc == "utf16":
+        elif self.get_type().enc == "utf16":
             # FIXME: real encoding...
             return len(val) * 2 + 2
         else:
@@ -1223,7 +1242,7 @@ class PinnedStr(PinnedType):
         return raw
 
     def __repr__(self):
-        return "%r(%s): %r" % (self.__class__, self._enc, self.val)
+        return "%r: %r" % (self.__class__, self.val)
 
 
 class PinnedArray(PinnedType):
@@ -1260,7 +1279,7 @@ class PinnedArray(PinnedType):
 
     # just a shorthand
     def as_mem_str(self, encoding="ansi"):
-        return self.cast(PinnedStr, encoding)
+        return self.cast(Str(encoding).pinned)
 
     def raw(self):
         raise ValueError("%s is unsized, which prevents from getting its full "
diff --git a/test/analysis/mem.py b/test/analysis/mem.py
index e1a2861f..8d4a56d3 100644
--- a/test/analysis/mem.py
+++ b/test/analysis/mem.py
@@ -5,9 +5,9 @@
 import struct
 
 from miasm2.analysis.machine import Machine
-from miasm2.analysis.mem import PinnedStruct, Num, Ptr, PinnedStr, PinnedArray,\
-                                PinnedSizedArray, Array, RawStruct, Union, \
-                                BitField, PinnedSelf, PinnedVoid, Bits, \
+from miasm2.analysis.mem import PinnedStruct, Num, Ptr, Str, \
+                                Array, RawStruct, Union, \
+                                BitField, Self, Void, Bits, \
                                 set_allocator, PinnedUnion, Struct
 from miasm2.jitter.csts import PAGE_READ, PAGE_WRITE
 from miasm2.os_dep.common import heap
@@ -26,7 +26,7 @@ class MyStruct(PinnedStruct):
         # TODO: comment
         ("other", Ptr("I", OtherStruct)),
         # Ptr to a variable length String
-        ("s", Ptr("I", PinnedStr)),
+        ("s", Ptr("I", Str())),
         ("i", Ptr("I", Num("I"))),
     ]
 
@@ -117,7 +117,7 @@ assert memval == 8
 
 # Str tests
 ## Basic tests
-memstr = PinnedStr(jitter.vm, addr_str)
+memstr = Str().pinned(jitter.vm, addr_str)
 memstr.val = ""
 assert memstr.val == ""
 assert jitter.vm.get_mem(memstr.get_addr(), 1) == '\x00'
@@ -126,7 +126,7 @@ 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(PinnedStr) manipulations
+## Ptr(Str()) manipulations
 mstruct.s.val = memstr.get_addr()
 assert mstruct.s.val == addr_str
 assert mstruct.s.deref == memstr
@@ -136,13 +136,13 @@ assert mstruct.s.deref.val == "That's all folks!"
 assert memstr.val == "That's all folks!"
 
 ## Other address, same value, same encoding
-memstr2 = PinnedStr(jitter.vm, addr_str2)
+memstr2 = Str().pinned(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 = PinnedStr(jitter.vm, addr_str3, "utf16")
+memstr3 = Str("utf16").pinned(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
@@ -151,7 +151,7 @@ assert memstr3 != memstr # Encoding is different, so they are not eq
 assert memstr3.val == memstr.val # But the python value is the same
 
 
-# PinnedArray tests
+# Array tests
 # Allocate buffer manually, since memarray is unsized
 alloc_addr = my_heap.vm_alloc(jitter.vm, 0x100)
 memarray = Array(Num("I")).pinned(jitter.vm, alloc_addr)
@@ -346,15 +346,15 @@ assert bit.flags.f4_1 == 1
 # Unhealthy ideas
 class UnhealthyIdeas(PinnedStruct):
     fields = [
-        ("pastruct", Ptr("I", PinnedArray, RawStruct("=Bf"))),
-        ("apstr", Array(Ptr("I", PinnedStr), 10)),
-        ("pself", Ptr("I", PinnedSelf)),
-        ("apself", Array(Ptr("I", PinnedSelf), 2)),
-        ("ppself", Ptr("I", Ptr("I", PinnedSelf))),
-        ("pppself", Ptr("I", Ptr("I", Ptr("I", PinnedSelf)))),
+        ("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", PinnedVoid).size()
+p_size = Ptr("I", Void()).size()
 
 ideas = UnhealthyIdeas(jitter.vm)
 ideas.memset()
@@ -425,7 +425,7 @@ ms2.s2[5] = 0xab
 assert PinnedShort(jitter.vm, ms2.s2.get_addr(4)).val == 0xabcd
 
 # void* style cast
-PinnedPtrVoid = Ptr("I", PinnedVoid).pinned
+PinnedPtrVoid = Ptr("I", Void()).pinned
 PinnedPtrMyStruct = Ptr("I", MyStruct).pinned
 p = PinnedPtrVoid(jitter.vm)
 p.val = mstruct.get_addr()