feat: add syscall pre/post hooks

This also sets up a default hook that prevents
stdio from being closed.

examples/strace.py

@@ -0,0 +1,43 @@
+from pyda import *
+from pwnlib.elf.elf import ELF
+from pwnlib.util.packing import u64
+import string
+import sys
+
+p = process()
+
+e = ELF(p.exe_path)
+e.address = p.maps[p.exe_path].base
+
+plt_map = { e.plt[x]: x for x in e.plt }
+
+def guess_arg(x):
+    printable_chars = bytes(string.printable, 'ascii')
+
+    # Is pointer?
+    if x > 0x100000000:
+        try:
+            data = p.read(x, 0x20)
+            if all([c in printable_chars for c in data[:4]]):
+                if 0 in data:
+                    return str(data[:data.index(0)])
+                else:
+                    return str(data[:20]) + "..."
+
+        except Exception as e:
+            pass
+
+    return hex(x)
+
+def syscall_pre_hook(p, num):
+    print(f"[syscall {num}] (" + ", ".join([
+        f"rdi={guess_arg(p.regs.rdi)}",
+        f"rsi={guess_arg(p.regs.rsi)}",
+        f"rdx={guess_arg(p.regs.rdx)}",
+        f"rcx={guess_arg(p.regs.rcx)}",
+    ]) + ")")
+
+for snum in range(500):
+    p.syscall_pre(snum, syscall_pre_hook)
+
+p.run()

lib/pyda/base.py

@@ -6,11 +6,11 @@
 
 INIT = False
 
-def process():
+def process(**kwargs):
     global INIT
 
     # todo: remove the bogus argument
-    proc = Process(pyda_core.process(""))
+    proc = Process(pyda_core.process(""), **kwargs)
 
     if not INIT:
         # by this point, hacks/ is in pythonpath

lib/pyda/process.py

@@ -3,13 +3,46 @@
 import pyda_core
 
 class Process():
-    def __init__(self, handle):
+    def __init__(self, handle, prevent_close_stdio=True):
         self._p = handle
         self._hooks = {}
+        self._syscall_pre_hooks = {}
+        self._syscall_post_hooks = {}
+        self._registered_syscall_pre_hook = False
+        self._registered_syscall_post_hook = False
+        self._has_run = False
+
+        def prevent_close(p, num):
+            if p.regs.rdi in [0, 1, 2]:
+                p.regs.rax = 0
+                return False # pre-hooks that return False will prevent the syscall from executing
+
+            return None
+
+        if prevent_close_stdio:
+            self.syscall_pre(3, prevent_close)
 
     def _hook_dispatch(self, addr):
         for h in self._hooks[addr]:
             h(self)
+
+    def _syscall_pre_hook_dispatch(self, syscall_num):
+        if syscall_num in self._syscall_pre_hooks:
+            results = []
+            for h in self._syscall_pre_hooks[syscall_num]:
+                results.append(h(self, syscall_num))
+
+            if False in results and True in results:
+                raise RuntimeError("Cannot have mixed return values from syscall pre-hooks")
+            elif False in results:
+                return False
+            elif True in results:
+                return True
+
+    def _syscall_post_hook_dispatch(self, syscall_num):
+        if syscall_num in self._syscall_pre_hooks:
+            for h in self._syscall_pre_hooks[syscall_num]:
+                h(self, syscall_num)
 
     def hook(self, addr, callback):
         if addr not in self._hooks:
@@ -38,6 +71,32 @@ def after_call_hook(p):
 
         self.hook(addr, call_hook)
 
+    def syscall_pre(self, syscall_num, callback):
+        if self._has_run:
+            raise RuntimeError("Cannot add syscall hooks after process has started")
+
+        if not self._registered_syscall_pre_hook:
+            self._p.set_syscall_pre_hook(lambda p, syscall_num: self._syscall_pre_hook_dispatch(syscall_num))
+            self._registered_syscall_pre_hook = True
+
+        if syscall_num not in self._syscall_pre_hooks:
+            self._syscall_pre_hooks[syscall_num] = [callback]
+        else:
+            self._syscall_pre_hooks[syscall_num].append(callback)
+
+    def syscall_post(self, syscall_num, callback):
+        if self._has_run:
+            raise RuntimeError("Cannot add syscall hooks after process has started")
+
+        if not self._registered_syscall_post_hook:
+            self._p.set_syscall_post_hook(lambda p, syscall_num: self._syscall_post_hook_dispatch(syscall_num))
+            self._registered_syscall_post_hook = True
+
+        if syscall_num not in self._syscall_post_hooks:
+            self._syscall_post_hooks[syscall_num] = [callback]
+        else:
+            self._syscall_post_hooks[syscall_num].append(callback)
+
     def set_thread_entry(self, callback):
         self._p.set_thread_init_hook(lambda p: callback(self))
 
@@ -61,6 +120,7 @@ def __getattr__(self, name):
         raise AttributeError(f"Invalid attribute '{name}'. Did you mean 'regs.{name}'?")
 
     def run(self):
+        self._has_run = True
         self._p.run()
 
     @property

pyda_core/pyda_core.c

@@ -23,6 +23,8 @@ pyda_process* pyda_mk_process() {
     proc->main_thread = pyda_mk_thread(proc);
     proc->callbacks = NULL;
     proc->thread_init_hook = NULL;
+    proc->syscall_pre_hook = NULL;
+    proc->syscall_post_hook = NULL;
     proc->py_obj = NULL;
 
     pthread_condattr_t condattr;
@@ -257,6 +259,26 @@ void pyda_set_thread_init_hook(pyda_process *p, PyObject *callback) {
     Py_INCREF(callback);
 }
 
+void pyda_set_syscall_pre_hook(pyda_process *p, PyObject *callback) {
+    // NOTE: GIL is held
+
+    if (p->syscall_pre_hook)
+        Py_DECREF(p->syscall_pre_hook);
+
+    p->syscall_pre_hook = callback;
+    Py_INCREF(callback);
+}
+
+void pyda_set_syscall_post_hook(pyda_process *p, PyObject *callback) {
+    // NOTE: GIL is held
+
+    if (p->syscall_post_hook)
+        Py_DECREF(p->syscall_post_hook);
+
+    p->syscall_post_hook = callback;
+    Py_INCREF(callback);
+}
+
 int pyda_flush_hooks() {
     pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
     pyda_process *p = t->proc;
@@ -284,28 +306,74 @@ pyda_hook* pyda_get_callback(pyda_process *p, void* addr) {
     return NULL;
 }
 
-void pyda_hook_cleancall(pyda_hook *cb) {
-    pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
+static void thread_prepare_for_python_entry(PyGILState_STATE *gstate, pyda_thread *t, void* pc) {
     if (t->skip_next_hook) {
         t->skip_next_hook = 0;
         return;
     }
 
     if (t->errored) return;
 
-    PyGILState_STATE gstate;
-    gstate = PyGILState_Ensure();
+    *gstate = PyGILState_Ensure();
 
     void *drcontext = dr_get_current_drcontext();
     t->cur_context.size = sizeof(dr_mcontext_t);
     t->cur_context.flags = DR_MC_ALL; // dr_redirect_execution requires it
     dr_get_mcontext(drcontext, &t->cur_context);
-    t->cur_context.pc = (app_pc)cb->addr;
+
+    if (pc)
+        t->cur_context.pc = (app_pc)pc;
+
     t->rip_updated_in_cleancall = 0;
+}
+
+static void thread_prepare_for_python_return(PyGILState_STATE *gstate, pyda_thread *t, void* hook_addr) {
+    void *drcontext = dr_get_current_drcontext();
+
+    // Syscall hooks are not allowed to modify PC
+    if (!hook_addr) {
+        if (t->rip_updated_in_cleancall) {
+            dr_fprintf(STDERR, "\n[Pyda] ERROR: Syscall hooks are not allowed to modify PC. Skipping future hooks.\n");
+            dr_flush_file(STDERR);
+            t->errored = 1;
+        }
+        pyda_flush_hooks(); // There is no risk of invalidating the current block here, since we are about to do a syscall
+        dr_set_mcontext(drcontext, &t->cur_context);
+        PyGILState_Release(*gstate);
+        return;
+    }
+
+    if (t->cur_context.pc == (app_pc)hook_addr && t->rip_updated_in_cleancall) {
+        if (t->rip_updated_in_cleancall) {
+            dr_fprintf(STDERR, "\n[Pyda] ERROR: Hook updated RIP to the same address. This is UB. Skipping future hooks.\n");
+            dr_flush_file(STDERR);
+            t->errored = 1;
+        }
+    }
+
+    if (pyda_flush_hooks() || t->rip_updated_in_cleancall) {
+        if (t->cur_context.pc == hook_addr) {
+            t->skip_next_hook = 1;
+        }
+        // we need to call dr_redirect_execution
+        PyGILState_Release(*gstate);
+        dr_redirect_execution(&t->cur_context);
+    } else {
+        dr_set_mcontext(drcontext, &t->cur_context);
+        PyGILState_Release(*gstate);
+    }
+}
+
+void pyda_hook_cleancall(pyda_hook *cb) {
+    PyGILState_STATE gstate;
+    pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
+
+    thread_prepare_for_python_entry(&gstate, t, cb->addr);
 
     DEBUG_PRINTF("cleancall %p %p %p\n", cb, cb->py_func, t);
 
     PyObject *result = PyObject_CallFunctionObjArgs(cb->py_func, t->proc->py_obj, NULL);
+
     if (result == NULL) {
         dr_fprintf(STDERR, "\n[Pyda] ERROR: Hook call failed. Skipping future hooks on thread %d\n", t->tid);
         dr_flush_file(STDERR);
@@ -318,27 +386,46 @@ void pyda_hook_cleancall(pyda_hook *cb) {
     }
 
     DEBUG_PRINTF("cleancall ret %p %p %p\n", cb, cb->py_func, t);
+    thread_prepare_for_python_return(&gstate, t, cb->addr);
+}
 
-    if (t->cur_context.pc == (app_pc)cb->addr && t->rip_updated_in_cleancall) {
-        if (t->rip_updated_in_cleancall) {
-            dr_fprintf(STDERR, "\n[Pyda] Hook updated RIP to the same address. This is UB. Skipping future hooks.\n");
-            dr_flush_file(STDERR);
-            t->errored = 1;
-            // dr_abort();
-        }
-    }
+int pyda_hook_syscall(int syscall_num, int is_pre) {
+    PyGILState_STATE gstate;
+    pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
 
-    if (pyda_flush_hooks() || t->rip_updated_in_cleancall) {
-        if (t->cur_context.pc == cb->addr) {
-            t->skip_next_hook = 1;
-        }
-        // we need to call dr_redirect_execution
-        PyGILState_Release(gstate);
-        dr_redirect_execution(&t->cur_context);
+    PyObject *hook = (is_pre ? t->proc->syscall_pre_hook : t->proc->syscall_post_hook);
+    if (!hook) return 1;
+
+    thread_prepare_for_python_entry(&gstate, t, NULL);
+
+    DEBUG_PRINTF("syscall %d pre %d\n", syscall_num, is_pre);
+
+    int should_run = 1;
+
+    PyObject *syscall_num_obj = PyLong_FromLong(syscall_num);
+    PyObject *result = PyObject_CallFunctionObjArgs(hook, t->proc->py_obj, syscall_num_obj, NULL);
+
+    Py_DECREF(syscall_num_obj);
+
+    if (result == NULL) {
+        dr_fprintf(STDERR, "\n[Pyda] ERROR: Syscall hook call failed. Skipping future hooks on thread %d\n", t->tid);
+        dr_flush_file(STDERR);
+        t->errored = 1;
+        PyErr_Print();
+        dr_fprintf(STDERR, "\n");
+    } else if (is_pre && PyBool_Check(result)) { 
+        // Should run
+        should_run = PyObject_IsTrue(result);
+        DEBUG_PRINTF("syscall pre_hook returned %d\n", should_run);
     } else {
-        dr_set_mcontext(drcontext, &t->cur_context);
-        PyGILState_Release(gstate);
+        Py_DECREF(result);
+        DEBUG_PRINTF("syscall hook returned non-bool\n");
     }
+
+    DEBUG_PRINTF("syscall ret %d pre %d\n", syscall_num, is_pre);
+    thread_prepare_for_python_return(&gstate, t, NULL);
+
+    return should_run;
 }
 
 #endif

pyda_core/pyda_core.h

@@ -36,6 +36,9 @@ struct pyda_process_s {
 
     pyda_thread *main_thread;
     PyObject *thread_init_hook;
+    PyObject *syscall_pre_hook;
+    PyObject *syscall_post_hook;
+
     PyObject *py_obj;
 
     pthread_cond_t thread_exit_cond;
@@ -89,12 +92,14 @@ void pyda_initial_break(pyda_thread *t);
 void pyda_add_hook(pyda_process *p, uint64_t addr, PyObject *callback);
 void pyda_remove_hook(pyda_process *p, uint64_t addr);
 void pyda_set_thread_init_hook(pyda_process *p, PyObject *callback);
+void pyda_set_syscall_pre_hook(pyda_process *p, PyObject *callback);
+void pyda_set_syscall_post_hook(pyda_process *p, PyObject *callback);
 pyda_hook* pyda_get_callback(pyda_process *p, void* addr);
 
 // These can only be called from application threads
 int pyda_flush_hooks();
 void pyda_hook_cleancall(pyda_hook *cb);
-
+int pyda_hook_syscall(int syscall_num, int is_pre);
 
 #ifndef PYDA_DYNAMORIO_CLIENT