memdump-based pe parsing with pefile, locate pe resources

.gitignore

@@ -3,4 +3,5 @@ roach.egg-info/
 .cache/
 .coverage
 dist/
+.pytest_cache/
 *.pyc

roach/__init__.py

@@ -4,7 +4,7 @@
 
 from roach.disasm import disasm
 from roach.hash.sha import md5, sha1, sha224, sha384, sha256, sha512
-from roach.short import aes, rc4, pe, aplib, procmem, pad, insn
+from roach.short import aes, rc4, pe, aplib, procmem, procmempe, pad, insn
 from roach.string.ops import asciiz
 
 from roach.procmem import (

roach/pe.py

@@ -4,10 +4,32 @@
 
 import pefile
 
+from roach.procmem import ProcessMemoryPE, ProcessMemory
+
+class OurSectionStructure(pefile.SectionStructure):
+    def get_data(self, start=None, length=None):
+        if not isinstance(self.pe.__data__, ProcessMemoryPE):
+            return OrigSectionStructure.get_data(self, start, length)
+        data = self.pe.__data__
+        return data.readv(data.imgbase + start, length)
+
+OrigSectionStructure = pefile.SectionStructure
+pefile.SectionStructure = OurSectionStructure
+
 class PE(object):
-    def __init__(self, data):
+    """Wrapper around pefile.PE; accepts either a string (raw file contents) or
+    a ProcessMemoryPE instance."""
+
+    def __init__(self, data, fast_load=True):
+        if data.__class__ == ProcessMemory:
+            raise RuntimeError("procmem parameter should be procmempe!")
+
+        if data.__class__ == ProcessMemoryPE:
+            fast_load = False
+            data.parent = self
+
         self.data = data
-        self.pe = pefile.PE(data=data, fast_load=True)
+        self.pe = pefile.PE(data=data, fast_load=fast_load)
 
     @property
     def dos_header(self):
@@ -43,3 +65,21 @@ def section(self, name):
         for section in self.pe.sections:
             if section.Name.rstrip("\x00") == name:
                 return section
+
+    def resource(self, name):
+        # TODO Implement search by type.
+        name_str = lambda e1, e2, e3: e1.name and e1.name.string == name
+        name_int = lambda e1, e2, e3: e2.struct.Name == name
+
+        if isinstance(name, basestring):
+            compare = name_str
+        else:
+            compare = name_int
+
+        for e1 in self.pe.DIRECTORY_ENTRY_RESOURCE.entries:
+            for e2 in e1.directory.entries:
+                for e3 in e2.directory.entries:
+                    if compare(e1, e2, e3):
+                        return self.pe.get_data(
+                            e3.data.struct.OffsetToData, e3.data.struct.Size
+                        )

roach/procmem.py

@@ -130,7 +130,10 @@ def readv(self, addr, length):
         """Reads a continuous buffer with address and length."""
         ret = []
         while length:
-            a, l = self.addr_range(addr)
+            r = self.addr_range(addr)
+            if not r:
+                break
+            a, l = r
             l = min(a + l - addr, length)
             ret.append(self.read(self.v2p(addr), l))
             addr, length = addr + l, length - l
@@ -203,3 +206,48 @@ def regexv(self, query):
 
     def disasmv(self, addr, size):
         return disasm(self.readv(addr, size), addr)
+
+    def findmz(self, addr):
+        """Locates the MZ header based on an address."""
+        addr &= ~0xffff
+        while True:
+            buf = self.readv(addr, 2)
+            if not buf:
+                return
+            if buf == "MZ":
+                return addr
+            addr -= 0x10000
+
+class ProcessMemoryPE(ProcessMemory):
+    """Wrapper around ProcessMemory for reading in-memory PE files."""
+
+    def __init__(self, p, imgbase, load=True):
+        self.imgbase = imgbase
+
+        if p.__class__ == ProcessMemoryPE:
+            raise RuntimeError("object already procmempe!")
+
+        # Upgrade existing ProcessMemory instance.
+        if p.__class__ == ProcessMemory:
+            self.f = p.f
+            self.m = p.m
+            self.load = p.load
+            self._regions = p.regions
+        else:
+            ProcessMemory.__init__(p, load)
+
+    def __len__(self):
+        r = self.regions[-1]
+        return r.addr + r.size
+
+    def __getitem__(self, item):
+        if isinstance(item, (int, long)):
+            return self.readv(self.imgbase + item, 1)
+
+        if item.start is None:
+            start = self.imgbase
+            stop = item.stop
+        else:
+            start = self.imgbase + item.start
+            stop = item.stop - item.start
+        return self.readv(start, stop)

roach/short.py

@@ -7,7 +7,7 @@
 from roach.crypto.rc import RC4
 from roach.disasm import Instruction
 from roach.pe import PE
-from roach.procmem import ProcessMemory
+from roach.procmem import ProcessMemory, ProcessMemoryPE
 from roach.string.ops import Padding
 
 class aes(object):
@@ -38,5 +38,6 @@ def rc4(key, data):
 pe = PE
 aplib = aPLib()
 procmem = ProcessMemory
+procmempe = ProcessMemoryPE
 pad = Padding("pkcs7")
 insn = Instruction

setup.py

@@ -23,6 +23,7 @@
     include_package_data=True,
     install_requires=[
         "cryptography>=2.1",
+        "pefile2==1.2.11",
         "pycrypto",
     ],
     extras_require={

tests/test_pe.py

@@ -28,3 +28,10 @@ def test_pe_header():
     assert img.is32bit is True
     assert img.is64bit is False
     assert img.section(".text").VirtualAddress == 0x1000
+
+def test_calc_exe():
+    p = pe(open("tests/files/calc.exe", "rb").read(), fast_load=False)
+    assert p.is32bit is True
+    data = p.resource("WEVT_TEMPLATE")
+    assert data.startswith("CRIM")
+    assert len(data) == 4750

tests/test_procmem.py

@@ -6,7 +6,7 @@
 import struct
 import tempfile
 
-from roach import procmem, pad, insn, PAGE_READWRITE
+from roach import procmem, procmempe, pad, pe, insn, PAGE_READWRITE
 
 def test_procmem_dummy_dmp():
     p = procmem("tests/files/dummy.dmp")
@@ -52,6 +52,22 @@ def test_procmem_dummy_dmp():
     assert p.uint64v(0x41410ffe) == 0x4242424242424141
     assert p.p2v(p.v2p(0x41411414)) == 0x41411414
 
+def test_calc_dmp():
+    p = procmem("tests/files/calc.dmp")
+    assert p.findmz(0x129abc) == 0xd0000
+    p = procmempe(p, 0xd0000)
+    assert p[0] == "M" and p[1] == "Z" and p[:2] == "MZ"
+    # Old/regular method with PE header.
+    assert pe(p.readv(p.imgbase, 0x1000)).dos_header.e_lfanew == 0xd8
+    assert p[0xd8:0xdc] == "PE\x00\x00"
+
+    assert pe(p).is32bit is True
+    d = pe(p).optional_header.DATA_DIRECTORY[2]
+    assert d.VirtualAddress == 0x59000 and d.Size == 0x62798
+    data = pe(p).resource("WEVT_TEMPLATE")
+    assert data.startswith("CRIM")
+    assert len(data) == 4750
+
 def test_methods():
     fd, filepath = tempfile.mkstemp()
     os.write(fd, "".join((