372 lines
12 KiB
Python
372 lines
12 KiB
Python
|
# Lazymio (mio@lazym.io)
|
||
|
|
|
||
|
|
from typing import Dict, List
|
||
|
|
from unicorn import *
|
||
|
|
from unicorn.x86_const import UC_X86_REG_EIP, UC_X86_REG_RIP
|
||
|
|
from unicorn.arm64_const import UC_ARM64_REG_PC
|
||
|
|
from unicorn.arm_const import UC_ARM_REG_PC, UC_ARM_REG_CPSR
|
||
|
|
from unicorn.mips_const import UC_MIPS_REG_PC
|
||
|
|
from unicorn.m68k_const import UC_M68K_REG_PC
|
||
|
|
from unicorn.riscv_const import UC_RISCV_REG_PC
|
||
|
|
from unicorn.ppc_const import UC_PPC_REG_PC
|
||
|
|
from unicorn.sparc_const import UC_SPARC_REG_PC
|
||
|
|
|
||
|
|
import gevent
|
||
|
|
import gevent.threadpool
|
||
|
|
import gevent.lock
|
||
|
|
import threading
|
||
|
|
|
||
|
|
# This class is named UnicornTask be design since it's not a
|
||
|
|
# real thread. The expected usage is to inherit this class
|
||
|
|
# and overwrite specific methods.
|
||
|
|
#
|
||
|
|
# This class is a friend class of MultiTaskUnicorn
|
||
|
|
class UnicornTask:
|
||
|
|
|
||
|
|
def __init__(self, uc: Uc, begin: int, end: int, task_id = None):
|
||
|
|
self._uc = uc
|
||
|
|
self._begin = begin
|
||
|
|
self._end = end
|
||
|
|
self._stop_request = False
|
||
|
|
self._ctx = None
|
||
|
|
self._task_id = None
|
||
|
|
self._arch = self._uc._arch
|
||
|
|
self._mode = self._uc._mode
|
||
|
|
|
||
|
|
@property
|
||
|
|
def pc(self):
|
||
|
|
""" Get current PC of the thread. This property should only be accessed when
|
||
|
|
the task is running.
|
||
|
|
"""
|
||
|
|
raw_pc = self._raw_pc()
|
||
|
|
if (self._uc.reg_read(UC_ARM_REG_CPSR) & (1 << 5)):
|
||
|
|
return raw_pc | 1
|
||
|
|
else:
|
||
|
|
return raw_pc
|
||
|
|
|
||
|
|
def _raw_pc(self):
|
||
|
|
# This extension is designed to be independent of Qiling, so let's
|
||
|
|
# do this manually...
|
||
|
|
if self._arch == UC_ARCH_X86:
|
||
|
|
if (self._mode & UC_MODE_32) != 0:
|
||
|
|
return self._uc.reg_read(UC_X86_REG_EIP)
|
||
|
|
elif (self._mode & UC_MODE_64) != 0:
|
||
|
|
return self._uc.reg_read(UC_X86_REG_RIP)
|
||
|
|
elif self._arch == UC_ARCH_MIPS:
|
||
|
|
return self._uc.reg_read(UC_MIPS_REG_PC)
|
||
|
|
elif self._arch == UC_ARCH_ARM:
|
||
|
|
return self._uc.reg_read(UC_ARM_REG_PC)
|
||
|
|
|
||
|
|
elif self._arch == UC_ARCH_ARM64:
|
||
|
|
return self._uc.reg_read(UC_ARM64_REG_PC)
|
||
|
|
elif self._arch == UC_ARCH_PPC:
|
||
|
|
return self._uc.reg_read(UC_PPC_REG_PC)
|
||
|
|
elif self._arch == UC_ARCH_M68K:
|
||
|
|
return self._uc.reg_read(UC_M68K_REG_PC)
|
||
|
|
elif self._arch == UC_ARCH_SPARC:
|
||
|
|
return self._uc.reg_read(UC_SPARC_REG_PC)
|
||
|
|
elif self._arch == UC_ARCH_RISCV:
|
||
|
|
return self._uc.reg_read(UC_RISCV_REG_PC)
|
||
|
|
|
||
|
|
# Really?
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def _reach_end(self):
|
||
|
|
# We may stop due to the scheduler asks us to, so check it manually.
|
||
|
|
#print(f"{hex(self._raw_pc())} {hex(self._end)}")
|
||
|
|
return self._raw_pc() == self._end
|
||
|
|
|
||
|
|
def save(self):
|
||
|
|
""" This method is used to save the task context.
|
||
|
|
Overwrite this method to implement specifc logic.
|
||
|
|
"""
|
||
|
|
return self._uc.context_save()
|
||
|
|
|
||
|
|
def restore(self, context):
|
||
|
|
""" This method is used to restore the task context.
|
||
|
|
Overwrite this method to implement specific logic.
|
||
|
|
"""
|
||
|
|
self._uc.context_restore(context)
|
||
|
|
self._begin = self.pc
|
||
|
|
|
||
|
|
def on_start(self):
|
||
|
|
""" This callback is triggered when a task gets scheduled.
|
||
|
|
"""
|
||
|
|
if self._ctx:
|
||
|
|
self.restore(self._ctx)
|
||
|
|
|
||
|
|
def on_interrupted(self, ucerr: int):
|
||
|
|
""" This callback is triggered when a task gets interrupted, which
|
||
|
|
is useful to emulate a clock interrupt.
|
||
|
|
"""
|
||
|
|
self._ctx = self.save()
|
||
|
|
|
||
|
|
def on_exit(self):
|
||
|
|
""" This callback is triggered when a task is about to exit.
|
||
|
|
"""
|
||
|
|
pass
|
||
|
|
|
||
|
|
# This manages nested uc_emu_start calls and is designed as a friend
|
||
|
|
# class of MultiTaskUnicorn.
|
||
|
|
class NestedCounter:
|
||
|
|
|
||
|
|
def __init__(self, mtuc: "MultiTaskUnicorn"):
|
||
|
|
self._mtuc = mtuc
|
||
|
|
|
||
|
|
def __enter__(self, *args, **kwargs):
|
||
|
|
self._mtuc._nested_started += 1
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, *args, **kwargs):
|
||
|
|
self._mtuc._nested_started -= 1
|
||
|
|
|
||
|
|
# This mimic a Unicorn object by maintaining the same interface.
|
||
|
|
# If no task is registered, the behavior is exactly the same as
|
||
|
|
# a normal unicorn.
|
||
|
|
#
|
||
|
|
# Note: To implement a non-block syscall:
|
||
|
|
# 1. Record the syscall in the hook, but **do nothing**
|
||
|
|
# 2. Stop emulation.
|
||
|
|
# 3. Handle the syscall in the on_interruped callback with
|
||
|
|
# proper **gevent functions** like gevent.sleep instead
|
||
|
|
# of time.sleep.
|
||
|
|
# 4. In this case, gevent would schedule another task to
|
||
|
|
# take emulation if the task is blocked.
|
||
|
|
#
|
||
|
|
# Bear in mind that only one task can be picked to emulate at
|
||
|
|
# the same time.
|
||
|
|
class MultiTaskUnicorn(Uc):
|
||
|
|
|
||
|
|
def __init__(self, arch, mode, interval: int = 100):
|
||
|
|
""" Create a MultiTaskUnicorn object.
|
||
|
|
Interval: Sceduling interval in **ms**. The longger interval, the better
|
||
|
|
performance but less interrupts.
|
||
|
|
"""
|
||
|
|
super().__init__(arch, mode)
|
||
|
|
self._interval = interval
|
||
|
|
self._tasks = {} # type: Dict[int, UnicornTask]
|
||
|
|
self._task_id_counter = 2000
|
||
|
|
self._to_stop = False
|
||
|
|
self._cur_utk_id = None
|
||
|
|
self._running = False
|
||
|
|
self._run_lock = threading.RLock()
|
||
|
|
self._multitask_enabled = False
|
||
|
|
self._count = 0
|
||
|
|
self._nested_started = 0
|
||
|
|
|
||
|
|
@property
|
||
|
|
def current_thread(self):
|
||
|
|
return self._tasks[self._cur_utk_id]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def running(self):
|
||
|
|
return self._running
|
||
|
|
|
||
|
|
def _next_task_id(self):
|
||
|
|
while self._task_id_counter in self._tasks:
|
||
|
|
self._task_id_counter += 1
|
||
|
|
|
||
|
|
return self._task_id_counter
|
||
|
|
|
||
|
|
def _emu_start_locked(self, begin: int, end: int, timeout: int, count: int):
|
||
|
|
with self._run_lock:
|
||
|
|
try:
|
||
|
|
self._running = True
|
||
|
|
super().emu_start(begin, end, timeout, count)
|
||
|
|
self._running = False
|
||
|
|
except UcError as err:
|
||
|
|
return err.errno
|
||
|
|
|
||
|
|
return UC_ERR_OK
|
||
|
|
|
||
|
|
def _emu_start_utk_locked(self, utk: UnicornTask):
|
||
|
|
return self._emu_start_locked(utk._begin, utk._end, 0, 0)
|
||
|
|
|
||
|
|
def _timeout_main(self, timeout: int):
|
||
|
|
|
||
|
|
gevent.sleep(timeout / 1000)
|
||
|
|
|
||
|
|
self.tasks_stop()
|
||
|
|
|
||
|
|
def _task_main(self, utk_id: int):
|
||
|
|
|
||
|
|
utk = self._tasks[utk_id]
|
||
|
|
use_count = (self._count > 0)
|
||
|
|
|
||
|
|
while True:
|
||
|
|
# utk may be stopped before running once, check it.
|
||
|
|
if utk._stop_request:
|
||
|
|
return # If we have to stop due to a tasks_stop, we preserve all threads so that we may resume.
|
||
|
|
|
||
|
|
self._cur_utk_id = utk_id
|
||
|
|
|
||
|
|
utk.on_start()
|
||
|
|
|
||
|
|
with gevent.Timeout(self._interval / 1000, False):
|
||
|
|
try:
|
||
|
|
pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
|
||
|
|
task = pool.spawn(self._emu_start_utk_locked, utk) # Run unicorn in a separate thread.
|
||
|
|
task.wait()
|
||
|
|
finally:
|
||
|
|
if not task.done():
|
||
|
|
# Interrupted by timeout, in this case we call uc_emu_stop.
|
||
|
|
super().emu_stop()
|
||
|
|
|
||
|
|
# Wait until we get the result.
|
||
|
|
ucerr = task.get()
|
||
|
|
|
||
|
|
if utk._reach_end():
|
||
|
|
utk._stop_request = True
|
||
|
|
|
||
|
|
if use_count:
|
||
|
|
self._count -= 1
|
||
|
|
|
||
|
|
if self._count <= 0:
|
||
|
|
self.tasks_stop()
|
||
|
|
return
|
||
|
|
|
||
|
|
if utk._stop_request:
|
||
|
|
utk.on_exit()
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
utk.on_interrupted(ucerr)
|
||
|
|
|
||
|
|
if self._to_stop:
|
||
|
|
return
|
||
|
|
|
||
|
|
# on_interrupted callback may have asked us to stop.
|
||
|
|
if utk._stop_request:
|
||
|
|
utk.on_exit()
|
||
|
|
break
|
||
|
|
|
||
|
|
# Give up control at once.
|
||
|
|
gevent.sleep(0)
|
||
|
|
|
||
|
|
del self._tasks[utk_id]
|
||
|
|
|
||
|
|
def tasks_save(self):
|
||
|
|
""" Save all tasks' contexts.
|
||
|
|
"""
|
||
|
|
return { k: v.save() for k, v in self._tasks.items() }
|
||
|
|
|
||
|
|
def tasks_restore(self, tasks_context: dict):
|
||
|
|
""" Restore the contexts of all tasks.
|
||
|
|
"""
|
||
|
|
for task_id, context in tasks_context:
|
||
|
|
if task_id in self._tasks:
|
||
|
|
self._tasks[task_id].restore(context)
|
||
|
|
|
||
|
|
def task_create(self, utk: UnicornTask):
|
||
|
|
""" Create a unicorn task. utk should be a initialized UnicornTask object.
|
||
|
|
If the task_id is not set, we generate one.
|
||
|
|
utk: The task to add.
|
||
|
|
"""
|
||
|
|
if not isinstance(utk, UnicornTask):
|
||
|
|
raise TypeError("Expect a UnicornTask or derived class")
|
||
|
|
self._multitask_enabled = True
|
||
|
|
if utk._task_id is None:
|
||
|
|
utk._task_id = self._next_task_id()
|
||
|
|
self._tasks[utk._task_id] = utk
|
||
|
|
return utk._task_id
|
||
|
|
|
||
|
|
def task_exit(self, utk_id):
|
||
|
|
""" Stop a task.
|
||
|
|
|
||
|
|
utk_id: The id returned from task_create.
|
||
|
|
"""
|
||
|
|
if utk_id not in self._tasks:
|
||
|
|
return
|
||
|
|
|
||
|
|
if utk_id == self._cur_utk_id and self._running:
|
||
|
|
self.emu_stop()
|
||
|
|
|
||
|
|
self._tasks[utk_id]._stop_request = True
|
||
|
|
|
||
|
|
def emu_start(self, begin: int, end: int, timeout: int, count: int):
|
||
|
|
""" Emulate an area of code just once. This overwrites the original emu_start interface and
|
||
|
|
provides extra cares when multitask is enabled. If no task is registerd, this call bahaves
|
||
|
|
like the original emu_start.
|
||
|
|
NOTE: Calling this method may cause current greenlet to be switched out.
|
||
|
|
begin, end, timeout, count: refer to Uc.emu_start
|
||
|
|
"""
|
||
|
|
if self._multitask_enabled:
|
||
|
|
if self._nested_started > 0:
|
||
|
|
with NestedCounter(self):
|
||
|
|
pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
|
||
|
|
task = pool.spawn(self._emu_start_locked, begin, end, timeout, count)
|
||
|
|
ucerr = task.get()
|
||
|
|
|
||
|
|
if ucerr != UC_ERR_OK:
|
||
|
|
raise UcError(ucerr)
|
||
|
|
|
||
|
|
return ucerr
|
||
|
|
else:
|
||
|
|
|
||
|
|
# Assume users resume on the last thread (and that should be the case)
|
||
|
|
if self._cur_utk_id in self._tasks:
|
||
|
|
self._tasks[self._cur_utk_id]._begin = begin
|
||
|
|
self._tasks[self._cur_utk_id]._end = end
|
||
|
|
else:
|
||
|
|
print(f"Warning: Can't found last thread we scheduled")
|
||
|
|
|
||
|
|
# This translation is not accurate, though.
|
||
|
|
self.tasks_start(count, timeout)
|
||
|
|
else:
|
||
|
|
return super().emu_start(begin, end, timeout, count)
|
||
|
|
|
||
|
|
def emu_stop(self):
|
||
|
|
""" Stop the emulation. If no task is registerd, this call bahaves like the original emu_stop.
|
||
|
|
"""
|
||
|
|
if self._multitask_enabled:
|
||
|
|
if self._running:
|
||
|
|
super().emu_stop()
|
||
|
|
# Stop the world as original uc_emu_stop does
|
||
|
|
if self._nested_started == 1:
|
||
|
|
self.tasks_stop()
|
||
|
|
else:
|
||
|
|
super().emu_stop()
|
||
|
|
|
||
|
|
def tasks_stop(self):
|
||
|
|
""" This will stop all running tasks. If no task is registered, this call does nothing.
|
||
|
|
"""
|
||
|
|
if self._multitask_enabled:
|
||
|
|
self._to_stop = True
|
||
|
|
if self._running:
|
||
|
|
super().emu_stop()
|
||
|
|
|
||
|
|
def tasks_start(self, count: int = 0, timeout: int = 0):
|
||
|
|
""" This will start emulation until all tasks get done.
|
||
|
|
count: Stop after sceduling *count* times. <=0 disables this check.
|
||
|
|
timeout: Stop after *timeout* ms. <=0 disables this check.
|
||
|
|
"""
|
||
|
|
workset = [] # type: List[gevent.Greenlet]
|
||
|
|
self._to_stop = False
|
||
|
|
self._count = count
|
||
|
|
|
||
|
|
if self._nested_started != 0:
|
||
|
|
print("Warning: tasks_start is called inside an uc_emu_start!")
|
||
|
|
return
|
||
|
|
|
||
|
|
with NestedCounter(self):
|
||
|
|
|
||
|
|
if self._count <= 0:
|
||
|
|
self._count = 0
|
||
|
|
|
||
|
|
if self._multitask_enabled:
|
||
|
|
|
||
|
|
if timeout > 0:
|
||
|
|
workset.append(gevent.spawn(self._timeout_main, timeout))
|
||
|
|
|
||
|
|
while len(self._tasks) != 0 and not self._to_stop:
|
||
|
|
|
||
|
|
new_workset = [ v for v in workset if not v.dead]
|
||
|
|
|
||
|
|
for utk_id in self._tasks:
|
||
|
|
new_workset.append(gevent.spawn(self._task_main, utk_id))
|
||
|
|
|
||
|
|
workset = new_workset
|
||
|
|
|
||
|
|
gevent.joinall(workset, raise_error=True)
|
||
|
|
|
||
|
|
if len(self._tasks) == 0:
|
||
|
|
self._multitask_enabled = False
|