unpack: aspack: fixup the PE header for OEP

capa/unpack/__init__.py

@@ -175,14 +175,15 @@ class AspackUnpacker(speakeasy.Speakeasy):
         with self.code_hook(until_section_hop(start, end)):
             self.emu.emu_eng.start(self.emu.get_pc())
 
-        logger.debug("OEP: 0x%x", self.emu.get_pc())
-
-        print(f"base: {self.module.get_base():x}")
-        for m in sorted(self.get_mem_maps(), key=lambda m: m.base):
-            print(f"{m.base:08x} {m.size:08x} {m.tag}")
+        oep = self.emu.get_pc()
+        logger.debug("OEP: 0x%x", oep)
 
         mm = self.get_address_map(self.module.get_base())
-        return self.mem_read(mm.get_base(), mm.get_size())
+        buf = self.mem_read(mm.get_base(), mm.get_size())
+
+        pe = pefile.PE(data=buf)
+        pe.OPTIONAL_HEADER.AddressOfEntryPoint = oep - self.module.base
+        return pe.write()
 
 
 def unpack_aspack(buf):

tests/test_unpack.py

@@ -6,8 +6,8 @@
 #  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
 import sys
-import textwrap
 
+import pefile
 import pytest
 from fixtures import *
 
@@ -42,6 +42,9 @@ def test_aspack_unpack(aspack_extractor):
         buf = f.read()
 
     unpacked = capa.unpack.unpack_pe(capa.unpack.ASPACK, buf)
+
+    pe = pefile.PE(data=unpacked)
+    assert pe.OPTIONAL_HEADER.AddressOfEntryPoint == 0x1a610
     assert b"This program cannot be run in DOS mode" in unpacked
     assert "(C) Copyright 1985-2000 Microsoft Corp.".encode("utf-16le") in unpacked
     assert "CMD.EXE has halted. %0".encode("utf-16le") in unpacked