feat(r2): wrapper class R2Qiling and R2Mem

chinggg · chinggg · commit b16e2c7ea07e · 2022-12-30T21:11:43.000+08:00
diff --git a/examples/extensions/r2/deflat_r2.py b/examples/extensions/r2/deflat_r2.py
@@ -9,7 +9,7 @@
 
 from qiling import Qiling
 from qiling.const import QL_VERBOSE
-from qiling.extensions.r2 import R2, R2Deflator
+from qiling.extensions.r2 import R2Qiling
 
 
 
diff --git a/qiling/extensions/r2/__init__.py b/qiling/extensions/r2/__init__.py
@@ -1,2 +1,39 @@
+from qiling import Qiling
 from .r2 import R2
+from .mem import R2Mem
 from .deflat import R2Deflator
+
+from unicorn.unicorn_const import UC_PROT_READ, UC_PROT_WRITE, UC_PROT_EXEC
+
+
+class R2Qiling(Qiling):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._mem = R2Mem(self.mem)
+        self.r2 = R2(self)
+
+
+def uc2perm(ps: int) -> str:
+    perms_d = {
+        UC_PROT_READ  : 'r',
+        UC_PROT_WRITE : 'w',
+        UC_PROT_EXEC  : 'x'
+    }
+
+    return ''.join(val if idx & ps else '-' for idx, val in perms_d.items())
+
+def assert_mem_equal(ql: 'R2Qiling'):
+    map_info = ql.mem.map_info
+    mem_regions = list(ql.uc.mem_regions())
+    assert len(map_info) == len(mem_regions), f'len: map_info={len(map_info)} != mem_regions={len(mem_regions)}'
+    for i, mem_region in enumerate(mem_regions):
+        s, e, p, _, _, data = map_info[i]
+        if (s, e - 1, p) != mem_region:
+            ql.log.error('map_info:')
+            print('\n'.join(ql.mem.get_formatted_mapinfo()))
+            ql.log.error('uc.mem_regions:')
+            print('\n'.join(f'{s:010x} - {e:010x}   {uc2perm(p)}' for (s, e, p) in mem_regions))
+            raise AssertionError(f'(start, end, perm): map_info={(s, e - 1, p)} != mem_region={mem_region}')
+        uc_mem = ql.mem.read(mem_region[0], mem_region[1] - mem_region[0] + 1)
+        assert len(data) == len(uc_mem), f'len of {i} mem: map_info={len(data)} != mem_region={len(uc_mem)}'
+        assert data == uc_mem, f'Memory region {i} {mem_region[0]:#x} - {mem_region[1]:#x} not equal to map_info[{i}]'
diff --git a/qiling/extensions/r2/mem.py b/qiling/extensions/r2/mem.py
@@ -0,0 +1,189 @@
+import ctypes
+
+
+from qiling.os.memory import QlMemoryManager, MapInfoEntry
+from qiling.exception import QlMemoryMappedError
+
+from typing import Any, Callable, Iterator, List, Mapping, MutableSequence, Optional, Pattern, Sequence, Tuple, Union
+
+from unicorn import UC_PROT_NONE, UC_PROT_READ, UC_PROT_WRITE, UC_PROT_EXEC, UC_PROT_ALL
+
+class R2Mem(QlMemoryManager):
+    '''A wrapper for QlMemoryManager that uses map_ptr and store raw memory in map_info
+       NOTE: ql.mem already contains map_infor after loader.run(), so instead of super().__init__(),
+       we accept mem object to simulate inheritance by composition 
+    '''
+
+    def __init__(self, mem: QlMemoryManager):
+        self.__dict__.update(mem.__dict__)
+        self._convert_map()
+
+    def _convert_map(self):
+        '''Clean existing map_info and remap memory'''
+        mapinfo = self.map_info.copy()
+        self.map_info = []
+        self.cmap = {}
+        for s, e, p, label, _mmio in mapinfo:
+          data = self.read(s, e - s)
+          self.ql.uc.mem_unmap(s, e - s)
+          self.map(s, e - s, p, label, data)
+    
+    def map(self, addr: int, size: int, perms: int = UC_PROT_ALL, info: Optional[str] = None, ptr: Optional[bytearray] = None):
+        """Map a new memory range.
+
+        Args:
+            addr: memory range base address
+            size: memory range size (in bytes)
+            perms: requested permissions mask
+            info: range label string
+            ptr: pointer to use (if any)
+
+        Raises:
+            QlMemoryMappedError: in case requested memory range is not fully available
+        """
+
+        assert perms & ~UC_PROT_ALL == 0, f'unexpected permissions mask {perms}'
+
+        if not self.is_available(addr, size):
+            for line in self.get_formatted_mapinfo():
+                print(line)
+            raise QlMemoryMappedError(f'Requested memory {addr:#x} + {size:#x} is unavailable')
+
+        buf = self.map_ptr(addr, size, perms, ptr)
+        self.add_mapinfo(addr, addr + size, perms, info or '[mapped]', is_mmio=False, data=buf)
+    
+    def map_ptr(self, addr: int, size: int, perms: int = UC_PROT_ALL, buf: Optional[bytearray] = None) -> bytearray:
+        """Map a new memory range allocated as Python bytearray, will not affect map_info
+
+        Args:
+            addr: memory range base address
+            size: memory range size (in bytes)
+            perms: requested permissions mask
+            buf: bytearray already allocated (if any)
+
+        Returns:
+            bytearray with size, should be added to map_info by caller
+        """
+        buf = buf or bytearray(size)
+        buf_type = ctypes.c_ubyte * size
+        cdata = buf_type.from_buffer(buf)
+        self.cmap[addr] = cdata
+        self.ql.uc.mem_map_ptr(addr, size, perms, cdata)
+        return buf
+
+    def add_mapinfo(self, mem_s: int, mem_e: int, mem_p: int, mem_info: str, is_mmio: bool = False, data : bytearray = None):
+        """Add a new memory range to map.
+
+        Args:
+            mem_s: memory range start
+            mem_e: memory range end
+            mem_p: permissions mask
+            mem_info: map entry label
+            is_mmio: memory range is mmio
+        """
+        self.map_info.append((mem_s, mem_e, mem_p, mem_info, is_mmio, data))
+        self.map_info.sort(key=lambda tp: tp[0])
+
+    def del_mapinfo(self, mem_s: int, mem_e: int):
+        """Subtract a memory range from map, will destroy data and unmap uc mem in the range.
+
+        Args:
+            mem_s: memory range start
+            mem_e: memory range end
+        """
+
+        tmp_map_info: MutableSequence[MapInfoEntry] = []
+
+        for s, e, p, info, mmio, data in self.map_info:
+            if e <= mem_s:
+                tmp_map_info.append((s, e, p, info, mmio, data))
+                continue
+
+            if s >= mem_e:
+                tmp_map_info.append((s, e, p, info, mmio, data))
+                continue
+
+            del self.cmap[s]  # remove cdata reference starting at s
+            if s < mem_s:
+                self.ql.uc.mem_unmap(s, mem_s - s)
+                self.map_ptr(s, mem_s - s, p, data[:mem_s - s])
+                tmp_map_info.append((s, mem_s, p, info, mmio, data[:mem_s - s]))
+
+            if s == mem_s:
+                pass
+
+            if e > mem_e:
+                self.ql.uc.mem_unmap(mem_e, e - mem_e)
+                self.map_ptr(mem_e, e - mem_e, p, data[mem_e - e:])
+                tmp_map_info.append((mem_e, e, p, info, mmio, data[mem_e - e:]))
+
+            if e == mem_e:
+                pass
+
+            del data[mem_s - s:mem_e - s]
+
+        self.map_info = tmp_map_info
+
+    def change_mapinfo(self, mem_s: int, mem_e: int, mem_p: Optional[int] = None, mem_info: Optional[str] = None, data: Optional[bytearray] = None):
+        tmp_map_info: Optional[MapInfoEntry] = None
+        info_idx: int = None
+
+        for idx, map_info in enumerate(self.map_info):
+            if mem_s >= map_info[0] and mem_e <= map_info[1]:
+                tmp_map_info = map_info
+                info_idx = idx
+                break
+
+        if tmp_map_info is None:
+            self.ql.log.error(f'Cannot change mapinfo at {mem_s:#08x}-{mem_e:#08x}')
+            return
+
+        if mem_p is not None:
+            data = data or self.read(mem_s, mem_e - mem_s).copy()
+            assert(len(data) == mem_e - mem_s)
+            self.unmap(mem_s, mem_e - mem_s)
+            self.map_ptr(mem_s, mem_e - mem_s, mem_p, data)
+            self.add_mapinfo(mem_s, mem_e, mem_p, mem_info or tmp_map_info[3], tmp_map_info[4], data)
+            return
+
+        if mem_info is not None:
+            self.map_info[info_idx] = (tmp_map_info[0], tmp_map_info[1], tmp_map_info[2], mem_info, tmp_map_info[4], tmp_map_info[5])
+    
+    def save(self):
+        """Save entire memory content.
+        """
+
+        mem_dict = {
+            "ram" : [],
+            "mmio" : []
+        }
+
+        for lbound, ubound, perm, label, is_mmio, data in self.map_info:
+            if is_mmio:
+                mem_dict['mmio'].append((lbound, ubound, perm, label, *self.mmio_cbs[(lbound, ubound)]))
+            else:
+                data = self.read(lbound, ubound - lbound)  # read instead of using data from map_info to avoid error
+                mem_dict['ram'].append((lbound, ubound, perm, label, data))
+
+        return mem_dict
+
+    def restore(self, mem_dict):
+        """Restore saved memory content.
+        """
+
+        for lbound, ubound, perms, label, data in mem_dict['ram']:
+            self.ql.log.debug(f'restoring memory range: {lbound:#08x} {ubound:#08x} {label}')
+
+            size = ubound - lbound
+            if self.is_available(lbound, size):
+                self.ql.log.debug(f'mapping {lbound:#08x} {ubound:#08x}, mapsize = {size:#x}')
+                self.map(lbound, size, perms, label, data)
+
+            self.ql.log.debug(f'writing {len(data):#x} bytes at {lbound:#08x}')
+            self.write(lbound, bytes(data))
+
+        for lbound, ubound, perms, label, read_cb, write_cb in mem_dict['mmio']:
+            self.ql.log.debug(f"restoring mmio range: {lbound:#08x} {ubound:#08x} {label}")
+
+            #TODO: Handle overlapped MMIO?
+            self.map_mmio(lbound, ubound - lbound, read_cb, write_cb, info=label)
diff --git a/qiling/extensions/r2/r2.py b/qiling/extensions/r2/r2.py
@@ -17,7 +17,7 @@
 from .deflat import R2Deflator
 
 if TYPE_CHECKING:
-    from qiling.core import Qiling
+    from qiling.extensions.r2 import R2Qiling
 
 def perm2uc(permstr: str) -> int:
     '''convert "-rwx" to unicorn const'''
@@ -208,7 +208,7 @@ def end(self):
 
 
 class R2:
-    def __init__(self, ql: "Qiling", baseaddr=(1 << 64) - 1, loadaddr=0):
+    def __init__(self, ql: 'R2Qiling', baseaddr=(1 << 64) - 1, loadaddr=0):
         super().__init__()
         self.ql = ql
         # r2 -B [baddr]   set base address for PIE binaries
@@ -242,7 +242,7 @@ def _rbuf_map(self, cbuf: ctypes.Array, perm: int = UC_PROT_ALL, addr: int = 0,
         desc = libr.r_io_open_buffer(self._r2i, rbuf, UC_PROT_ALL, 0)  # last arg `mode` is always 0 in r2 code
         libr.r_io.r_io_map_add(self._r2i, desc.contents.fd, desc.contents.perm, delta, addr, len(cbuf))
 
-    def _setup_mem(self, ql: 'Qiling'):
+    def _setup_mem(self, ql: 'R2Qiling'):
         if not hasattr(ql, '_mem'):
             return
         for start, _end, perms, _label, _mmio, _buf in ql.mem.map_info:
@@ -417,11 +417,11 @@ def set_backtrace(self, target: Union[int, str]):
         '''Set backtrace at target address before executing'''
         if isinstance(target, str):
             target = self.where(target)
-        def bt_hook(__ql: "Qiling", *args):
+        def bt_hook(__ql: 'R2Qiling', *args):
             print(self._backtrace_fuzzy())
         self.ql.hook_address(bt_hook, target)
 
-    def disassembler(self, ql: 'Qiling', addr: int, size: int, filt: Pattern[str]=None) -> int:
+    def disassembler(self, ql: 'R2Qiling', addr: int, size: int, filt: Pattern[str]=None) -> int:
         '''A human-friendly monkey patch of QlArchUtils.disassembler powered by r2, can be used for hook_code
             :param ql: Qiling instance
             :param addr: start address for disassembly
