Merge pull request #26 from mrexodia/module-manager

Module manager

Author: mrexodia

README.md

@@ -103,7 +103,7 @@ You can get the syscall parameters from [ntsyscalls.py](https://github.com/mrexo
 
 ## Collecting the dump
 
-There is a simple [x64dbg](https://github.com/x64dbg/x64dbg) plugin available called [MiniDumpPlugin](https://github.com/mrexodia/MiniDumpPlugin/releases). To create a dump, pause execution and execute the command `MiniDump my.dmp`.
+~~There is a simple [x64dbg](https://github.com/x64dbg/x64dbg) plugin available called [MiniDumpPlugin](https://github.com/mrexodia/MiniDumpPlugin/releases)~~ The [minidump](https://help.x64dbg.com/en/latest/commands/memory-operations/minidump.html) command has been integrated into x64dbg since 2022-10-10. To create a dump, pause execution and execute the command `MiniDump my.dmp`.
 
 ## Installation
 

src/dumpulator/details.py

@@ -1,5 +1,6 @@
 import struct
 from collections import namedtuple
+from typing import List
 
 from unicorn import *
 from unicorn.x86_const import *
@@ -563,3 +564,21 @@ def __setitem__(self, index, value):
     "Reserved"
 ]
 assert len(interrupt_names) == 32
+
+def format_table(table: List[List[str]]):
+    result = ""
+    header = table[0]
+    lengths = [0] * len(header)
+    for row in table:
+        for index, col in enumerate(row):
+            lengths[index] = max(lengths[index], len(col))
+    for row in table:
+        if len(result) > 0:
+            result += "\n"
+        line = ""
+        for index, col in enumerate(row):
+            if index > 0:
+                line += " "
+            line += f"{col:>{lengths[index]}}"
+        result += line.rstrip()
+    return result

src/dumpulator/dumpulator.py

@@ -30,28 +30,15 @@ class MemoryState(Enum):
     MEM_RESERVE = 0x2000
     MEM_FREE = 0x10000
 
-class MemoryBasicInformation:
-    def __init__(self, base: int, allocation_base: int, allocation_protect: MemoryProtect):
-        self.base = base
-        self.allocation_base = allocation_base
-        self.allocation_protect = allocation_protect
-        self.region_size: int = PAGE_SIZE
-        self.state: MemoryState = None
-        self.protect: MemoryProtect = None
-        self.type: MemoryType = None
-
-    def __str__(self):
-        return f"MemoryBasicInformation(base: {hex(self.base)}, allocation_base: {hex(self.allocation_base)}, region_size: {hex(self.region_size)}, state: {self.state}, protect: {self.protect}, type: {self.type})"
-
 class MemoryRegion:
-    def __init__(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.PAGE_NOACCESS, type: MemoryType = MemoryType.MEM_PRIVATE, comment: str = ""):
+    def __init__(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.PAGE_NOACCESS, type: MemoryType = MemoryType.MEM_PRIVATE, info: Any = None):
         assert start & 0xFFF == 0
         assert size & 0xFFF == 0
         self.start = start
         self.size = size
         self.protect = protect
         self.type = type
-        self.comment = comment
+        self.info = info
 
     @property
     def end(self):
@@ -84,12 +71,12 @@ def overlaps(self, other):
 
     def __str__(self):
         result = f"{hex(self.start)}[{hex(self.size)}]"
-        if len(self.comment) > 0:
-            result += f" ({self.comment})"
+        if self.info is not None:
+            result += f" ({self.info})"
         return result
 
     def __repr__(self) -> str:
-        return f"MemoryRegion({hex(self.start)}, {hex(self.size)}, {self.protect}, {self.type}, {self.comment}"
+        return f"MemoryRegion({hex(self.start)}, {hex(self.size)}, {self.protect}, {self.type}, {repr(self.info)})"
 
     def pages(self):
         for page in range(self.start, self.end, PAGE_SIZE):
@@ -105,6 +92,20 @@ def decommit(self, addr: int, size: int) -> None:
     def protect(self, addr: int, size: int, protect: MemoryProtect) -> None:
         raise NotImplementedError()
 
+class MemoryBasicInformation:
+    def __init__(self, base: int, allocation_base: int, allocation_protect: MemoryProtect):
+        self.base = base
+        self.allocation_base = allocation_base
+        self.allocation_protect = allocation_protect
+        self.region_size: int = PAGE_SIZE
+        self.state: MemoryState = None
+        self.protect: MemoryProtect = None
+        self.type: MemoryType = None
+        self.info: Any = None
+
+    def __str__(self):
+        return f"MemoryBasicInformation(base: {hex(self.base)}, allocation_base: {hex(self.allocation_base)}, region_size: {hex(self.region_size)}, state: {self.state}, protect: {self.protect}, type: {self.type})"
+
 class MemoryManager:
     def __init__(self, page_manager: PageManager, minimum = 0x10000, maximum = 0x7fffffff0000, granularity = 0x10000):
         self._page_manager = page_manager
@@ -114,9 +115,9 @@ def __init__(self, page_manager: PageManager, minimum = 0x10000, maximum = 0x7ff
         self._regions: List[MemoryRegion] = []
         self._committed: Dict[int, MemoryRegion] = {}
 
-    def find_parent(self, region: Union[MemoryRegion, int]):
+    def find_region(self, region: Union[MemoryRegion, int]):
         if isinstance(region, int):
-            region = MemoryRegion(self.page_align(region), 0)
+            region = MemoryRegion(self.align_page(region), 0)
         index = bisect.bisect_right(self._regions, region)
         if index == 0:
             return None
@@ -127,31 +128,35 @@ def find_parent(self, region: Union[MemoryRegion, int]):
             else:
                 return None
 
-    def page_align(self, addr: int):
+    def find_commit(self, addr: int):
+        addr = self.align_page(addr)
+        return self._committed.get(addr, None)
+
+    def align_page(self, addr: int):
         mask = PAGE_SIZE - 1
         return (addr + mask) & ~mask
 
-    def allocation_align(self, addr: int):
+    def align_allocation(self, addr: int):
         mask = self._granularity - 1
         return (addr + mask) & ~mask
 
     def find_free(self, size: int):
-        assert size > 0 and self.page_align(size) == size
+        assert size > 0 and self.align_page(size) == size
         base = self._minimum
         while base < self._maximum:
             info = self.query(base)
             assert info.base == base
-            if info.state == MemoryState.MEM_FREE and info.region_size >= size and self.allocation_align(base) == base:
+            if info.state == MemoryState.MEM_FREE and info.region_size >= size and self.align_allocation(base) == base:
                 return info.base
             base += info.region_size
         return None
 
-    def reserve(self, start: int, size: int, protect: MemoryProtect, type: MemoryType = MemoryType.MEM_PRIVATE, comment: str = ""):
+    def reserve(self, start: int, size: int, protect: MemoryProtect, type: MemoryType = MemoryType.MEM_PRIVATE, info: Any = None):
         assert isinstance(protect, MemoryProtect)
         assert isinstance(type, MemoryType)
-        assert size > 0 and self.page_align(size) == size
-        assert self.allocation_align(start) == start
-        region = MemoryRegion(start, size, protect, type, comment)
+        assert size > 0 and self.align_page(size) == size
+        assert self.align_allocation(start) == start
+        region = MemoryRegion(start, size, protect, type, info)
         if region.start < self._minimum or region.end > self._maximum:
             raise KeyError(f"Requested region {region} is out of bounds")
 
@@ -170,9 +175,9 @@ def check_overlaps(index):
         self._regions.insert(index, region)
 
     def release(self, start: int):
-        assert self.allocation_align(start) == start
+        assert self.align_allocation(start) == start
 
-        parent_region = self.find_parent(start)
+        parent_region = self.find_region(start)
         if parent_region is None:
             raise KeyError(f"Could not find parent for {hex(start)}")
         if parent_region.start != start:
@@ -191,10 +196,10 @@ def release(self, start: int):
 
     def commit(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.UNDEFINED):
         assert isinstance(protect, MemoryProtect)
-        assert size > 0 and self.page_align(size) == size
-        assert self.page_align(start) == start
+        assert size > 0 and self.align_page(size) == size
+        assert self.align_page(start) == start
         region = MemoryRegion(start, size)
-        parent_region = self.find_parent(region)
+        parent_region = self.find_region(region)
         if parent_region is None:
             raise KeyError(f"Could not find parent for {region}")
 
@@ -208,17 +213,19 @@ def commit(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.U
         else:
             for page in region.pages():
                 if page in self._committed:
-                    self._page_manager.protect(page, PAGE_SIZE, protect)
-                    self._committed[page].protect = protect
+                    page_region = self._committed[page]
+                    if page_region.protect != protect:
+                        self._page_manager.protect(page, PAGE_SIZE, protect)
+                        page_region.protect = protect
                 else:
                     self._page_manager.commit(page, PAGE_SIZE, protect)
                     self._committed[page] = MemoryRegion(page, PAGE_SIZE, protect, parent_region.type)
 
     def decommit(self, start: int, size: int):
-        assert size > 0 and self.page_align(size) == size
-        assert self.page_align(start) == start
+        assert size > 0 and self.align_page(size) == size
+        assert self.align_page(start) == start
         region = MemoryRegion(start, size)
-        parent_region = self.find_parent(region)
+        parent_region = self.find_region(region)
         if parent_region is None:
             raise KeyError(f"Could not find parent for {region}")
 
@@ -234,10 +241,10 @@ def decommit(self, start: int, size: int):
 
     def protect(self, start: int, size: int, protect: MemoryProtect):
         assert isinstance(protect, MemoryProtect)
-        assert size > 0 and self.page_align(size) == size
-        assert self.page_align(start) == start
+        assert size > 0 and self.align_page(size) == size
+        assert self.align_page(start) == start
         region = MemoryRegion(start, size)
-        parent_region = self.find_parent(region)
+        parent_region = self.find_region(region)
         if parent_region is None:
             raise KeyError(f"Could not find parent for {region}")
 
@@ -250,15 +257,17 @@ def protect(self, start: int, size: int, protect: MemoryProtect):
         old_protect = self._committed[region.start].protect
         self._page_manager.protect(region.start, region.size, protect)
         for page in region.pages():
-            self._committed[page].protect = protect
+            page_region = self._committed[page]
+            if page_region.protect != protect:
+                page_region.protect = protect
 
         return old_protect
 
     def query(self, start: int):
-        assert self.page_align(start) == start
+        start = self.align_page(start)
 
         region = MemoryRegion(start, 0)
-        parent_region = self.find_parent(region)
+        parent_region = self.find_region(region)
         if parent_region is None:
             index = bisect.bisect_right(self._regions, region)
             next_start = self._maximum
@@ -279,11 +288,15 @@ def query(self, start: int):
                 continue
             elif result is None:
                 result = MemoryBasicInformation(page, parent_region.start, parent_region.protect)
+                if page == parent_region.start:
+                    result.info = parent_region.info
                 if page in self._committed:
                     result.state = MemoryState.MEM_COMMIT
                     commited_page = self._committed[page]
                     result.protect = commited_page.protect
                     result.type = commited_page.type
+                    if commited_page.info:
+                        result.info = commited_page.info
                     assert commited_page.type == parent_region.type
                 else:
                     result.state = MemoryState.MEM_RESERVE
@@ -292,6 +305,8 @@ def query(self, start: int):
             else:
                 commited_page = self._committed.get(page, None)
                 if result.state == MemoryState.MEM_RESERVE:
+                    if commited_page is not None:
+                        break
                     result.region_size += PAGE_SIZE
                 elif result.state == MemoryState.MEM_COMMIT:
                     if commited_page is not None and commited_page.type == result.type and commited_page.protect == result.protect:
@@ -301,3 +316,12 @@ def query(self, start: int):
                 else:
                     assert False  # unreachable
         return result
+
+    def map(self):
+        addr = self._minimum
+        regions: List[MemoryBasicInformation] = []
+        while addr < self._maximum:
+            info = self.query(addr)
+            regions.append(info)
+            addr += info.region_size
+        return regions

src/dumpulator/memory.py

@@ -0,0 +1,109 @@
+from typing import Dict, Optional, Type, Union, List
+
+import pefile
+from .memory import MemoryManager
+
+# TODO: support forwards
+class ModuleExport:
+    def __init__(self, address: int, ordinal: int, name: str):
+        self.address = address
+        self.ordinal = ordinal
+        self.name = name
+
+class Module:
+    def __init__(self, pe: pefile.PE, path: str):
+        self.pe = pe
+        self.path = path
+        self.name = path.split("\\")[-1]
+        self._exports_by_address: Dict[int, int] = {}
+        self._exports_by_ordinal: Dict[int, int] = {}
+        self._exports_by_name: Dict[str, int] = {}
+        self.exports: List[ModuleExport] = []
+        self._parse_pe()
+
+    def _parse_pe(self):
+        self.base: int = self.pe.OPTIONAL_HEADER.ImageBase
+        self.size: int = self.pe.OPTIONAL_HEADER.SizeOfImage
+        self.entry: int = self.base + self.pe.OPTIONAL_HEADER.AddressOfEntryPoint
+        self.pe.parse_data_directories(directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXPORT"]])
+        pe_exports = self.pe.DIRECTORY_ENTRY_EXPORT.symbols if hasattr(self.pe, "DIRECTORY_ENTRY_EXPORT") else []
+        for pe_export in pe_exports:
+            va = self.base + pe_export.address
+            if pe_export.name:
+                name = pe_export.name.decode("ascii")
+            else:
+                name = None
+            export = ModuleExport(va, pe_export.ordinal, name)
+            self._exports_by_address[export.address] = len(self.exports)
+            self._exports_by_ordinal[export.ordinal] = len(self.exports)
+            if name is not None:
+                self._exports_by_name[name] = len(self.exports)
+            self.exports.append(export)
+
+    def find_export(self, key: Union[str, int]):
+        if isinstance(key, int):
+            index = self._exports_by_ordinal.get(key, None)
+            if index is None:
+                index = self._exports_by_address.get(key, None)
+            if index is None:
+                return None
+            return self.exports[index]
+        elif isinstance(key, str):
+            index = self._exports_by_name.get(key)
+            if index is None:
+                return None
+            return self.exports[index]
+        raise TypeError()
+
+    def __repr__(self):
+        return f"Module({hex(self.base)}, {hex(self.size)}, {repr(self.path)})"
+
+    def __contains__(self, addr: int):
+        return addr >= self.base and addr < self.base + self.size
+
+class ModuleManager:
+    def __init__(self, memory: MemoryManager):
+        self._memory = memory
+        self._name_lookup: Dict[str, int] = {}
+        self._modules: Dict[int, Module] = {}
+
+    def add(self, pe: pefile.PE, path: str):
+        module = Module(pe, path)
+        self._modules[module.base] = module
+        region = self._memory.find_region(module.base)
+        assert region.start == module.base
+        assert region is not None
+        region.info = module
+        self._name_lookup[module.name] = module.base
+        self._name_lookup[module.name.lower()] = module.base
+        self._name_lookup[module.path] = module.base
+        return module
+
+    def find(self, key: Union[str, int]) -> Optional[Module]:
+        if isinstance(key, int):
+            region = self._memory.find_region(key)
+            if region.info:
+                assert isinstance(region.info, Module)
+                return region.info
+            return None
+        if isinstance(key, str):
+            base = self._name_lookup.get(key, None)
+            if base is None:
+                base = self._name_lookup.get(key.lower(), None)
+            if base is None:
+                return None
+            return self.find(base)
+        raise TypeError()
+
+    def __getitem__(self, key: Union[str, int]) -> Module:
+        module = self.find(key)
+        if module is None:
+            raise KeyError()
+        return module
+
+    def __contains__(self, key: Union[str, int]):
+        return self.find(key) is not None
+
+    def __iter__(self):
+        for base in self._modules:
+            yield self._modules[base]