Bin_stream: add get_u8/u16/u32/u64

3 files changed, 97 insertions, 12 deletions

diff --git a/miasm2/analysis/binary.py b/miasm2/analysis/binary.py
index 21e5c356..6dfef92a 100644
--- a/miasm2/analysis/binary.py
+++ b/miasm2/analysis/binary.py
@@ -156,7 +156,7 @@ class ContainerPE(Container):
 
         # Build the bin_stream instance and set the entry point
         try:
-            self._bin_stream = bin_stream_pe(self._executable.virt)
+            self._bin_stream = bin_stream_pe(self._executable)
             ep_detected = self._executable.Opthdr.AddressOfEntryPoint
             self._entry_point = self._executable.rva2virt(ep_detected)
         except Exception, error:
@@ -200,7 +200,7 @@ class ContainerELF(Container):
 
         # Build the bin_stream instance and set the entry point
         try:
-            self._bin_stream = bin_stream_elf(self._executable.virt)
+            self._bin_stream = bin_stream_elf(self._executable)
             self._entry_point = self._executable.Ehdr.entry + addr
         except Exception, error:
             raise ContainerParsingException('Cannot read ELF: %s' % error)
diff --git a/miasm2/core/bin_stream.py b/miasm2/core/bin_stream.py
index f8d491d2..8bd59467 100644
--- a/miasm2/core/bin_stream.py
+++ b/miasm2/core/bin_stream.py
@@ -16,6 +16,10 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 #
 
+from miasm2.core.utils import BIG_ENDIAN, LITTLE_ENDIAN
+from miasm2.core.utils import upck8le, upck16le, upck32le, upck64le
+from miasm2.core.utils import upck8be, upck16be, upck32be, upck64be
+
 
 class bin_stream(object):
 
@@ -26,7 +30,7 @@ class bin_stream(object):
     _atomic_mode = False
 
     def __init__(self, *args, **kargs):
-        pass
+        self.endianness = LITTLE_ENDIAN
 
     def __repr__(self):
         return "<%s !!>" % self.__class__.__name__
@@ -104,6 +108,55 @@ class bin_stream(object):
             start += cur_len
         return out
 
+    def get_u8(self, addr, endianness=None):
+        """
+        Return u8 from address @addr
+        endianness: Optional: LITTLE_ENDIAN/BIG_ENDIAN
+        """
+        if endianness is None:
+            endianness = self.endianness
+        data = self.getbytes(addr, 1)
+        return data
+
+    def get_u16(self, addr, endianness=None):
+        """
+        Return u16 from address @addr
+        endianness: Optional: LITTLE_ENDIAN/BIG_ENDIAN
+        """
+        if endianness is None:
+            endianness = self.endianness
+        data = self.getbytes(addr, 2)
+        if endianness == LITTLE_ENDIAN:
+            return upck16le(data)
+        else:
+            return upck32be(data)
+
+    def get_u32(self, addr, endianness=None):
+        """
+        Return u32 from address @addr
+        endianness: Optional: LITTLE_ENDIAN/BIG_ENDIAN
+        """
+        if endianness is None:
+            endianness = self.endianness
+        data = self.getbytes(addr, 4)
+        if endianness == LITTLE_ENDIAN:
+            return upck32le(data)
+        else:
+            return upck32be(data)
+
+    def get_u64(self, addr, endianness=None):
+        """
+        Return u64 from address @addr
+        endianness: Optional: LITTLE_ENDIAN/BIG_ENDIAN
+        """
+        if endianness is None:
+            endianness = self.endianness
+        data = self.getbytes(addr, 8)
+        if endianness == LITTLE_ENDIAN:
+            return upck64le(data)
+        else:
+            return upck64be(data)
+
 
 class bin_stream_str(bin_stream):
 
@@ -168,14 +221,14 @@ class bin_stream_file(bin_stream):
 
 class bin_stream_container(bin_stream):
 
-    def __init__(self, virt_view, offset=0L):
+    def __init__(self, binary, offset=0L):
         bin_stream.__init__(self)
-        self.bin = virt_view
-        self.l = virt_view.max_addr()
+        self.bin = binary
+        self.l = binary.virt.max_addr()
         self.offset = offset
 
     def is_addr_in(self, ad):
-        return self.bin.is_addr_in(ad)
+        return self.bin.virt.is_addr_in(ad)
 
     def getlen(self):
         return self.l
@@ -184,16 +237,16 @@ class bin_stream_container(bin_stream):
         if self.offset + l > self.l:
             raise IOError("not enough bytes")
         self.offset += l
-        return self.bin.get(self.offset - l, self.offset)
+        return self.bin.virt.get(self.offset - l, self.offset)
 
     def _getbytes(self, start, l=1):
         try:
-            return self.bin.get(start, start + l)
+            return self.bin.virt.get(start, start + l)
         except ValueError:
             raise IOError("cannot get bytes")
 
     def __str__(self):
-        out = self.bin.get(self.offset, self.offset + self.l)
+        out = self.bin.virt.get(self.offset, self.offset + self.l)
         return out
 
     def setoffset(self, val):
@@ -201,11 +254,15 @@ class bin_stream_container(bin_stream):
 
 
 class bin_stream_pe(bin_stream_container):
-    pass
+    def __init__(self, binary, *args, **kwargs):
+        super(bin_stream_pe, self).__init__(binary, *args, **kwargs)
+        self.endianness = binary._sex
 
 
 class bin_stream_elf(bin_stream_container):
-    pass
+    def __init__(self, binary, *args, **kwargs):
+        super(bin_stream_elf, self).__init__(binary, *args, **kwargs)
+        self.endianness = binary.sex
 
 
 class bin_stream_vm(bin_stream):
@@ -214,6 +271,10 @@ class bin_stream_vm(bin_stream):
         self.offset = offset
         self.base_offset = base_offset
         self.vm = vm
+        if self.vm.is_little_endian():
+            self.endianness = LITTLE_ENDIAN
+        else:
+            self.endianness = BIG_ENDIAN
 
     def getlen(self):
         return 0xFFFFFFFFFFFFFFFF
diff --git a/miasm2/core/utils.py b/miasm2/core/utils.py
index 35ddbb82..c1f48418 100644
--- a/miasm2/core/utils.py
+++ b/miasm2/core/utils.py
@@ -12,6 +12,30 @@ pck16 = lambda x: struct.pack('H', x)
 pck32 = lambda x: struct.pack('I', x)
 pck64 = lambda x: struct.pack('Q', x)
 
+# Little endian
+upck8le = lambda x: struct.unpack('<B', x)[0]
+upck16le = lambda x: struct.unpack('<H', x)[0]
+upck32le = lambda x: struct.unpack('<I', x)[0]
+upck64le = lambda x: struct.unpack('<Q', x)[0]
+pck8le = lambda x: struct.pack('<B', x)
+pck16le = lambda x: struct.pack('<H', x)
+pck32le = lambda x: struct.pack('<I', x)
+pck64le = lambda x: struct.pack('<Q', x)
+
+# Big endian
+upck8be = lambda x: struct.unpack('>B', x)[0]
+upck16be = lambda x: struct.unpack('>H', x)[0]
+upck32be = lambda x: struct.unpack('>I', x)[0]
+upck64be = lambda x: struct.unpack('>Q', x)[0]
+pck8be = lambda x: struct.pack('>B', x)
+pck16be = lambda x: struct.pack('>H', x)
+pck32be = lambda x: struct.pack('>I', x)
+pck64be = lambda x: struct.pack('>Q', x)
+
+
+LITTLE_ENDIAN = 1
+BIG_ENDIAN = 2
+
 
 pck = {8: pck8,
        16: pck16,