Trying harder: copy-on-write on MemoryBuffer level and creating new r…

…eferences to memoryview on every from_memory call

malduck/extractor/extract_manager.py

@@ -211,7 +211,8 @@ def push_procmem(
 
         family = self._extract_procmem(p, matches)
         for binary in binaries:
-            family = self._extract_procmem(binary, matches) or family
+            with binary:
+                family = self._extract_procmem(binary, matches) or family
         return family
 
     @property

malduck/procmem/binmem.py

@@ -81,7 +81,7 @@ def load_binaries_from_memory(cls: Type[T], procmem: ProcessMemory) -> Iterator[
         if cls.__magic__ is None:
             raise NotImplementedError()
         for binary_va in procmem.findv(cls.__magic__):
-            binary_procmem_dmp = cls.from_memory(procmem.slicev(binary_va))
+            binary_procmem_dmp = cls.from_memory_slice(procmem, binary_va)
             if binary_procmem_dmp.is_valid():
                 yield binary_procmem_dmp
             binary_procmem_img = binary_procmem_dmp.image

malduck/procmem/idamem.py

@@ -49,7 +49,7 @@ def __len__(self) -> int:
         return self.idamem.regions[-1].end_offset
 
     def slice(
-        self, from_offset: Optional[int], to_offset: Optional[int]
+        self, from_offset: Optional[int] = None, to_offset: Optional[int] = None
     ) -> "MemoryBuffer":
         # HACK: IDAMemoryBuffer depends on region information from IDAProcessMemory
         # Let's assume that MemoryBuffer is never directly sliced and regions

malduck/procmem/membuf.py

@@ -18,7 +18,7 @@ def __len__(self) -> int:
 
     @abstractmethod
     def slice(
-        self, from_offset: Optional[int], to_offset: Optional[int]
+        self, from_offset: Optional[int] = None, to_offset: Optional[int] = None
     ) -> "MemoryBuffer":
         raise NotImplementedError
 
@@ -34,27 +34,38 @@ def __init__(
     ) -> None:
         if type(buf) is memoryview:
             self.buf = buf
-        elif type(buf) is bytearray:
+        elif type(buf) in (bytearray, bytes):
             self.buf = memoryview(buf)
-        elif type(buf) is bytes:
-            # Copy bytes to bytearray to make it patchable
-            self.buf = memoryview(bytearray(buf))
         else:
-            raise TypeError("Buffer in PlainMemoryBuffer must be bytes or bytearray")
+            raise TypeError(
+                "Buffer in PlainMemoryBuffer must be memoryview, bytes or bytearray"
+            )
 
     def __getitem__(self, item: slice) -> bytes:
         return bytes(self.buf[item])
 
     def __setitem__(self, item: slice, value: bytes) -> None:
+        if self.buf.readonly:
+            # If buffer is read-only, make a copy (on write)
+            patchable_buf = memoryview(bytearray(self.buf))
+            self.buf.release()
+            self.buf = patchable_buf
         self.buf[item] = value
 
     def __len__(self) -> int:
         return len(self.buf)
 
     def slice(
-        self, from_offset: Optional[int], to_offset: Optional[int]
+        self, from_offset: Optional[int] = None, to_offset: Optional[int] = None
     ) -> "MemoryBuffer":
-        return PlainMemoryBuffer(self.buf[from_offset:to_offset])
+        """
+        Creates a derived MemoryBuffer object representing slice of an underlying memory.
+
+        Derived buffer is readonly, so __setitem__ will first make a copy before applying
+        changes. It means that changes on parent buffer may be seen in derived buffers,
+        but not the other way.
+        """
+        return PlainMemoryBuffer(self.buf[from_offset:to_offset].toreadonly())
 
     def release(self) -> None:
         self.buf.release()
@@ -83,7 +94,7 @@ def __init__(
                 super().__init__(memoryview(self.mapped_buf))
             except RuntimeError:
                 # Fallback to file.read()
-                super().__init__(memoryview(bytearray(self.opened_file.read())))
+                super().__init__(memoryview(self.opened_file.read()))
                 self.opened_file.close()
                 self.opened_file = None
 

malduck/procmem/procmem.py

@@ -151,10 +151,47 @@ def from_memory(cls, memory, base=None, **kwargs):
         :rtype: :class:`ProcessMemory`
         """
         copied = cls(
-            memory.m, base=base or memory.imgbase, regions=memory.regions, **kwargs
+            memory.memory.slice(),
+            base=base or memory.imgbase,
+            regions=memory.regions,
+            **kwargs,
         )
         return copied
 
+    @classmethod
+    def from_memory_slice(cls, memory, addr, length=None):
+        """
+        Creates a new ProcessMemory object that represents a slice of ProcessMemory.
+        Slice can't be cross-region. If you don't provide length, it will be limited to
+        the length of the region containing provided virtual address.
+
+        :param addr: Starting virtual address
+        :type addr: int
+        :param length: Length of slice (optional)
+        :type length: Optional[int]
+        """
+        region = memory.addr_region(addr)
+        # Boundary check
+        if region is None or (length is not None and region.end < (addr + length)):
+            raise ValueError(
+                "Sliced range must be contained within single, existing region"
+            )
+        new_regions = [
+            Region(
+                addr,
+                region.size if length is None else length,
+                region.state,
+                region.type_,
+                region.protect,
+                0,
+            )
+        ]
+        slice_from_offset = region.offset + (addr - region.addr)
+        slice_to_offset = region.offset + length if length is not None else None
+        new_memory = memory.memory.slice(slice_from_offset, slice_to_offset)
+        new_imgbase = addr
+        return cls(new_memory, base=new_imgbase, regions=new_regions)
+
     @property
     def length(self):
         """
@@ -409,40 +446,6 @@ def readv_until(self, addr, s):
         idx = chunk.find(s)
         return chunk[:idx] if idx >= 0 else chunk
 
-    def slicev(self, addr, length=None):
-        """
-        Creates a new ProcessMemory object that represents a slice of ProcessMemory.
-
-        Slice can't be cross-region. If you don't provide length, it will be limited to
-        the length of the region containing provided virtual address.
-
-        :param addr: Starting virtual address
-        :type addr: int
-        :param length: Length of slice (optional)
-        :type length: Optional[int]
-        """
-        region = self.addr_region(addr)
-        # Boundary check
-        if region is None or (length is not None and region.end < (addr + length)):
-            raise ValueError(
-                "Sliced range must be contained within single, existing region"
-            )
-        new_regions = [
-            Region(
-                addr,
-                region.size if length is None else length,
-                region.state,
-                region.type_,
-                region.protect,
-                0,
-            )
-        ]
-        slice_from_offset = region.offset + (addr - region.addr)
-        slice_to_offset = region.offset + length if length is not None else None
-        new_memory = self.memory.slice(slice_from_offset, slice_to_offset)
-        new_imgbase = addr
-        return self.__class__(new_memory, base=new_imgbase, regions=new_regions)
-
     def patchp(self, offset, buf):
         """
         Patch bytes under specified offset

malduck/procmem/procmem.pyi

@@ -95,7 +95,6 @@ class ProcessMemory:
     ) -> Iterator[Tuple[int, bytes]]: ...
     def readv(self, addr: int, length: Optional[int] = None) -> bytes: ...
     def readv_until(self, addr: int, s: bytes) -> bytes: ...
-    def slicev(self: T, addr: int, length: Optional[int] = None) -> T: ...
     def patchp(self, offset: int, buf: bytes) -> None: ...
     def patchv(self, addr: int, buf: bytes) -> None: ...
     @overload