Source code for tritondse.memory

from triton import TritonContext
import bisect
from typing import Optional, Union, Generator, List
from collections import namedtuple
import struct
from contextlib import contextmanager

from tritondse.types import Perm, Addr, ByteSize, Endian

MemMap = namedtuple('Map', "start size perm name")


class MapOverlapException(Exception):
    """
    Exception raised when trying to map a memory area where some of
    the addresses overlap with an already mapped area.
    """
    pass


class MemoryAccessViolation(Exception):
    """
    Exception triggered when accessing memory with
    the wrong permissions.
    """
    def __init__(self, addr: Addr, access: Perm, map_perm: Perm = None, memory_not_mapped: bool = False, perm_error: bool = False):
        """
        :param addr: address where the violation occured
        :param access: type of access performed
        :param map_perm: permission of the memory page of `address`
        :param memory_not_mapped: whether the address was mapped or not
        :param perm_error: whether it is a permission error
        """
        super(MemoryAccessViolation, self).__init__()
        self.address: Addr = addr
        """ address where the violation occurred"""
        self._is_mem_unmapped = memory_not_mapped
        self._is_perm_error = perm_error
        self.access: Perm = access
        """Access type that was performed"""
        self.map_perm: Optional[Perm] = map_perm
        """Permissions of the memory map associated to the address"""

    def is_permission_error(self) -> bool:
        """True if the exception was caused by a permission issue"""
        return self._is_perm_error

    def is_memory_unmapped_error(self) -> bool:
        """
        Return true if the exception was raised due to an access
        to an area not mapped
        """
        return self._is_mem_unmapped

    def __str__(self) -> str:
        if self.is_permission_error():
            return f"(addr:{self.address:#08x}, access:{str(self.access)} on map:{str(self.map_perm)})"
        else:
            return f"({str(self.access)}: {self.address:#08x} unmapped)"

    def __repr__(self):
        return str(self)


STRUCT_MAP = {
    (True, 1): 'B',
    (False, 1): 'b',
    (True, 2): 'H',
    (False, 2): 'h',
    (True, 4): 'I',
    (False, 4): 'i',
    (True, 8): 'Q',
    (False, 8): 'q'
}

ENDIAN_MAP = {
    Endian.LITTLE: "<",
    Endian.BIG: ">"
}


[docs] class Memory(object): """ Memory representation of the current :py:class:`ProcessState` object. It wraps all interaction with Triton's memory context to provide high-level function. It adds a segmentation and memory permission model at the top of Triton. It also overrides __getitem__ and the slice mechanism to be able read and write concrete memory values in a Pythonic manner. """ def __init__(self, ctx: TritonContext, endianness: Endian = Endian.LITTLE): """ :param ctx: TritonContext to interface with """ self.ctx: TritonContext = ctx """Underlying Triton context""" self._linear_map_addr = [] # List of [map_start, map_end, map_start, map_end ...] self._linear_map_map = [] # List of [MemMap, None, MemMap, None ...] self._segment_enabled = True self._endian = endianness self._endian_key = ENDIAN_MAP[self._endian] self._mem_cbs_enabled = True # self._maps = {} # Addr: -> Map
[docs] def set_endianess(self, en: Endian) -> None: """ Set the endianness of memory accesses. By default, endianess is little. :param en: Endian: Endianess to use. :return: None """ self._endian = en self._endian_key = ENDIAN_MAP[self._endian]
@property def _ptr_size(self) -> int: return self.ctx.getGprSize() @property def segmentation_enabled(self) -> bool: """ returns whether segmentation enforcing is enabled :return: True if segmentation is enabled """ return self._segment_enabled
[docs] def disable_segmentation(self) -> None: """ Turn-off segmentation enforcing. """ self._segment_enabled = False
[docs] def enable_segmentation(self) -> None: """ Turn-off segmentation enforcing. """ self._segment_enabled = True
[docs] def set_segmentation(self, enabled: bool) -> None: """ Set the segmentation enforcing with the given boolean. """ self._segment_enabled = enabled
[docs] @contextmanager def without_segmentation(self, disable_callbacks=False) -> Generator['Memory', None, None]: """ Context manager enabling manipulating temporarily the memory without considering the memory permissions. E.g: It enables writing data in a memory mapped in RX :param disable_callbacks: Whether to disable memory callbacks that could have been set :return: """ previous = self._segment_enabled self.disable_segmentation() cbs = self._mem_cbs_enabled self._mem_cbs_enabled = not disable_callbacks yield self self._mem_cbs_enabled = cbs self.set_segmentation(previous)
[docs] def callbacks_enabled(self) -> bool: """ Return whether memory callbacks are enabled. :return: True if callbacks are enabled """ return self._mem_cbs_enabled
[docs] def get_maps(self) -> Generator[MemMap, None, None]: """ Iterate all the memory maps defined, including all memory areas allocated on the heap. :return: generator of all :py:class:`MemMap` objects """ yield from (x for x in self._linear_map_map if x)
[docs] def map(self, start, size, perm: Perm = Perm.R | Perm.W | Perm.X, name="") -> MemMap: """ Map the given address and size in memory with the given permission. :raise MapOverlapException: In the case the map overlap an existing mapping :param start: address to map :param size: size to map :param perm: permission :param name: name to given to the memory region :return: MemMap freshly mapped """ def _map_idx(idx): self._linear_map_addr.insert(idx, start + size - 1) # end address is included self._linear_map_addr.insert(idx, start) self._linear_map_map.insert(idx, None) memmap = MemMap(start, size, perm, name) self._linear_map_map.insert(idx, memmap) return memmap if not self._linear_map_addr: # Nothing mapped yet return _map_idx(0) idx = bisect.bisect_left(self._linear_map_addr, start) if idx == len(self._linear_map_addr): # It should be mapped at the end return _map_idx(idx) addr = self._linear_map_addr[idx] if (idx % 2) == 0: # We are on a start address if start < addr and start+size <= addr: # Can fit before return _map_idx(idx) else: # there is an overlap raise MapOverlapException(f"0x{start:08x}:{size} overlap with map: 0x{addr:08x} (even)") else: # We are on an end address prev = self._linear_map_addr[idx-1] raise MapOverlapException(f"0x{start:08x}:{size} overlap with map: 0x{prev:08x} (odd)")
[docs] def unmap(self, addr: Addr) -> None: """ Unmap the :py:class:`MemMap` object mapped at the address. The address can be within the map and not requires pointing at the head. :param addr: address to unmap :return: None """ def _unmap_idx(idx): self._linear_map_addr.pop(idx) # Pop the start self._linear_map_addr.pop(idx) # Pop the end self._linear_map_map.pop(idx) # Pop the object self._linear_map_map.pop(idx) # Pop the None padding idx = bisect.bisect_left(self._linear_map_addr, addr) try: mapaddr = self._linear_map_addr[idx] if (idx % 2) == 0: # We are on a start address (meaning we should be exactly on map start other unmapped) if addr == mapaddr: # We are exactly on the map address _unmap_idx(idx) else: raise MemoryAccessViolation(addr, Perm(0), memory_not_mapped=True) else: # We are on an end address _unmap_idx(idx-1) except IndexError: raise MemoryAccessViolation(addr, Perm(0), memory_not_mapped=True)
[docs] def mprotect(self, addr: Addr, perm: Perm) -> None: """ Update the map at the given address with permissions provided in argument. :param addr: address of the map of which to change permission :param perm: permission to assign :return: None """ idx = bisect.bisect_left(self._linear_map_addr, addr) try: if (idx % 2) == 0: # We are on a start address (meaning we should be exactly on map start other unmapped) map = self._linear_map_map[idx] self._linear_map_map[idx] = MemMap(map.start, map.size, perm, map.name) # replace map with new perms else: # We are on an end address map = self._linear_map_map[idx-1] self._linear_map_map[idx-1] = MemMap(map.start, map.size, perm, map.name) # replace map with new perms except IndexError: raise MemoryAccessViolation(addr, Perm(0), memory_not_mapped=True)
def __setitem__(self, key: Addr, value: bytes) -> None: """ Assign the given value at the address given by the key. The value must be bytes but can be multiple bytes. Warning: You cannot use the slice API on this function. :param key: address to write to :param value: content to write :raise MemoryAccessViolation: in case of invalid access """ if isinstance(key, slice): raise TypeError("slice unsupported for __setitem__") else: self.write(key, value) def __getitem__(self, item: Union[Addr, slice]) -> bytes: """ Read the memory at the given address. If the key is an integer reads a single byte. If the key is a slice: read addr+size bytes in memory. :param item: address, or address:size to read :return: memory content :raise MemoryAccessViolation: if the access is invalid """ if isinstance(item, slice): return self.read(item.start, item.stop) elif isinstance(item, int): return self.read(item, 1)
[docs] def write(self, addr: Addr, data: bytes) -> None: """ Write the given `data` bytes at `addr` address. :param addr: address where to write :param data: data to write :return: None """ if self._segment_enabled: map = self._get_map(addr, len(data)) if map is None: raise MemoryAccessViolation(addr, Perm.W, memory_not_mapped=True) if Perm.W not in map.perm: raise MemoryAccessViolation(addr, Perm.W, map_perm=map.perm, perm_error=True) return self.ctx.setConcreteMemoryAreaValue(addr, data)
[docs] def read(self, addr: Addr, size: ByteSize) -> bytes: """ Read `size` bytes at `addr` address. :param addr: address to read :param size: size of content to read :return: bytes read """ if self._segment_enabled: map = self._get_map(addr, size) if map is None: raise MemoryAccessViolation(addr, Perm.R, memory_not_mapped=True) if Perm.R not in map.perm: raise MemoryAccessViolation(addr, Perm.R, map_perm=map.perm, perm_error=True) return self.ctx.getConcreteMemoryAreaValue(addr, size)
def _get_map(self, ptr: Addr, size: ByteSize) -> Optional[MemMap]: """ Internal function returning the MemMap object associated with any address. It returns None if part of the memory range falls out of a memory mapping. Complexity is O(log(n)) :param ptr: address in memory :param size: size of the memory :return: True if mapped """ idx = bisect.bisect_left(self._linear_map_addr, ptr) try: addr = self._linear_map_addr[idx] if (idx % 2) == 0: # We are on a start address (meaning we should be exactly on map start other unmapped) end = self._linear_map_addr[idx+1] return self._linear_map_map[idx] if (ptr == addr and ptr+size <= end+1) else None else: # We are on an end address start = self._linear_map_addr[idx-1] return self._linear_map_map[idx-1] if (start <= addr and ptr+size <= addr+1) else None # fit into the map except IndexError: return None # Either raised when linear_map is empty or the address is beyond everything that is mapped
[docs] def get_map(self, addr: Addr, size: ByteSize = 1) -> Optional[MemMap]: """ Find the MemMap associated with the given address and returns it if any. :param addr: Address of the map (or any map inside) :param size: size of bytes for which we want the map :return: MemMap if found """ return self._get_map(addr, size)
[docs] def find_map(self, name: str) -> Optional[List[MemMap]]: """ Find a map given its name. :param name: Map name :return: MemMap if found """ l = [] for map in (x for x in self._linear_map_map if x): if map.name == name: l.append(map) return l
[docs] def map_from_name(self, name: str) -> MemMap: """ Return a map from its name. This function assumes the map is present. :raise AssertionError: If the map is not found :param name: Map name :return: MemMap """ for map in (x for x in self._linear_map_map if x): if map.name == name: return map assert False
[docs] def is_mapped(self, ptr: Addr, size: ByteSize = 1) -> bool: """ The function checks whether the memory is mapped or not. The implementation return False if the memory chunk overlap on two memory regions. Complexity is O(log(n)) :param ptr: address in memory :param size: size of the memory :return: True if mapped """ return self._get_map(ptr, size) is not None
[docs] def has_ever_been_written(self, ptr: Addr, size: ByteSize) -> bool: """ Returns whether the given range of addresses has previously been written or not. (Do not take in account the memory mapping). :param ptr: The pointer to check :type ptr: :py:obj:`tritondse.types.Addr` :param size: Size of the memory range to check :return: True if all addresses have been defined """ return self.ctx.isConcreteMemoryValueDefined(ptr, size)
[docs] def read_uint(self, addr: Addr, size: ByteSize = 4): """ Read in the process memory a **little-endian** integer of the ``size`` at ``addr``. :param addr: Address at which to read data :type addr: :py:obj:`tritondse.types.Addr` :param size: Number of bytes to read :type size: Union[str, :py:obj:`tritondse.types.ByteSize`] :return: Integer value read :raise struct.error: If value can't fit in `size` """ data = self.read(addr, size) return struct.unpack(self._endian_key+STRUCT_MAP[(True, size)], data)[0]
[docs] def read_sint(self, addr: Addr, size: ByteSize = 4): """ Read in the process memory a **little-endian** integer of the ``size`` at ``addr``. :param addr: Address at which to read data :type addr: :py:obj:`tritondse.types.Addr` :param size: Number of bytes to read :type size: Union[str, :py:obj:`tritondse.types.ByteSize`] :return: Integer value read :raise struct.error: If value can't fit in `size` """ data = self.read(addr, size) return struct.unpack(self._endian_key+STRUCT_MAP[(False, size)], data)[0]
[docs] def read_ptr(self, addr: Addr) -> int: """ Read in the process memory a little-endian integer of size :py:attr:`tritondse.ProcessState.ptr_size` :param addr: Address at which to read data :type addr: :py:obj:`tritondse.types.Addr` :return: Integer value read """ return self.read_uint(addr, self._ptr_size)
[docs] def read_char(self, addr: Addr) -> int: """ Read a char in memory (1-byte) following endianess. :param addr: address to read :return: char value as int """ return self.read_sint(addr, 1)
[docs] def read_uchar(self, addr: Addr) -> int: """ Read an unsigned char in memory (1-byte) following endianness. :param addr: address to read :return: unsigned char value as int """ return self.read_uint(addr, 1)
[docs] def read_int(self, addr: Addr) -> int: """ Read a signed integer in memory (4-byte) following endianness. :param addr: address to read :return: signed integer value as int """ return self.read_sint(addr, 4)
[docs] def read_word(self, addr: Addr) -> int: """ Read signed word in memory (2-byte) following endianness. :param addr: address to read :return: signed word value as int """ return self.read_uint(addr, 2)
[docs] def read_dword(self, addr: Addr) -> int: """ Read signed double word in memory (4-byte) following endianness. :param addr: address to read :return: dword value as int """ return self.read_uint(addr, 4)
[docs] def read_qword(self, addr: Addr) -> int: """ Read signed qword in memory (8-byte) following endianness. :param addr: address to read :return: qword value as int """ return self.read_uint(addr, 8)
[docs] def read_long(self, addr: Addr) -> int: """ Read 'C style' long in memory (4-byte) following endianness. :param addr: address to read :return: value as int """ return self.read_sint(addr, 4)
[docs] def read_ulong(self, addr: Addr) -> int: """ Read unsigned long in memory (4-byte) following endianness. :param addr: address to read :return: unsigned long value as int """ return self.read_uint(addr, 4)
[docs] def read_long_long(self, addr: Addr) -> int: """ Read long long in memory (8-byte) following endianness. :param addr: address to read :return: long long value as int """ return self.read_sint(addr, 8)
[docs] def read_ulong_long(self, addr: Addr) -> int: """ Read unsigned long long in memory (8-byte) following endianness. :param addr: address to read :return: unsigned long long value as int """ return self.read_uint(addr, 8)
[docs] def read_string(self, addr: Addr) -> str: """ Read a string in process memory at the given address .. warning:: The memory read is unbounded. Thus the memory is iterated up until finding a 0x0. :returns: the string read in memory :rtype: str """ s = "" index = 0 while True: val = self.read_uint(addr+index, 1) if not val: return s s += chr(val) index += 1
[docs] def write_int(self, addr: Addr, value: int, size: ByteSize = 4): """ Write in the process memory the given integer value of the given size at a specific address. :param addr: Address at which to read data :param value: data to write represented as an integer :param size: Number of bytes to read :raise struct.error: If integer value cannot fit in `size` """ self.write(addr, struct.pack(self._endian_key+STRUCT_MAP[(value >= 0, size)], value))
[docs] def write_ptr(self, addr: Addr, value: int) -> None: """ Similar to :py:meth:`write_int` but the size is automatically adjusted to be ``ptr_size``. :param addr: address where to write data :type addr: :py:obj:`tritondse.types.Addr` :param value: pointer value to write :type value: int :raise struct.error: If integer value cannot fit in a pointer size """ self.write_int(addr, value, self._ptr_size)
[docs] def write_char(self, addr: Addr, value: int) -> None: """ Write the integer value as a single byte in memory. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a byte (>255) """ self.write_int(addr, value, 1)
[docs] def write_word(self, addr: Addr, value: int) -> None: """ Write the word (2-byte) in memory following endianess. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a word """ self.write_int(addr, value, 2)
[docs] def write_dword(self, addr: Addr, value: int) -> None: """ Write the word (4-byte) in memory following endianess. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a dword """ self.write_int(addr, value, 4)
[docs] def write_qword(self, addr: Addr, value: int) -> None: """ Write the qword (8-byte) in memory following endianess. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a qword """ self.write_int(addr, value, 8)
[docs] def write_long(self, addr: Addr, value: int) -> None: """ Write a "C style" long (4-byte) in memory following endianess. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a long """ return self.write_int(addr, value, 4)
[docs] def write_long_long(self, addr: Addr, value: int) -> None: """ Write the "C style" long long (8-byte) in memory following endianess. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a long long """ return self.write_int(addr, value, 8)