Hopefully fixed the Linux backend after Process restructuring

nitro/backends/linux/backend.py

@@ -5,7 +5,7 @@
 
 from nitro.syscall import Syscall
 from nitro.event import SyscallDirection
-from nitro.process import Process
+from nitro.backends.linux.process import LinuxProcess
 from nitro.backends.backend import Backend
 from nitro.backends.linux.arguments import LinuxArgumentMap
 
@@ -20,10 +20,8 @@ class LinuxBackend(Backend):
         "nb_vcpu",
         "syscall_stack",
         "tasks_offset",
-        "pid_offset",
         "mm_offset",
         "pgd_offset",
-        "name_offset"
     )
 
     def __init__(self, domain, libvmi):
@@ -37,10 +35,8 @@ def __init__(self, domain, libvmi):
         self.syscall_stack = tuple([] for _ in range(self.nb_vcpu))
 
         self.tasks_offset = self.libvmi.get_offset("linux_tasks")
-        self.pid_offset = self.libvmi.get_offset("linux_pid")
         self.mm_offset = self.libvmi.get_offset("linux_mm")
         self.pgd_offset = self.libvmi.get_offset("linux_pgd")
-        self.name_offset = self.libvmi.get_offset("linux_name")
 
     def process_event(self, event):
         process = self.associate_process(event.sregs.cr3)
@@ -71,8 +67,6 @@ def associate_process(self, cr3):
         head = self.libvmi.translate_ksym2v("init_task") # get the address of swapper's task_struct
         next_ = head
         while True: # Maybe this should have a sanity check stopping it
-            pid = self.libvmi.read_32(next_ + self.pid_offset, 0)
-
             mm = self.libvmi.read_addr_va(next_ + self.mm_offset, 0)
             if not mm:
                 mm = self.libvmi.read_addr_va(next_ + self.mm_offset + VOID_P_SIZE, 0)
@@ -81,11 +75,9 @@ def associate_process(self, cr3):
                 pgd_phys_addr = self.libvmi.translate_kv2p(pgd)
                 if cr3 == pgd_phys_addr:
                     # Eventually, I would like to look for the executable name from mm->exe_file->f_path
-                    name = self.libvmi.read_str_va(next_ + self.name_offset, 0)
-                    process = Process(cr3, next_, name, pid, self.libvmi)
-                    return process
+                    return LinuxProcess(self.libvmi, cr3, next_)
             else:
-                logging.debug("no mm found for pid %s", pid)
+                logging.debug("missing mm")
             next_ = self.libvmi.read_addr_va(next_ + self.tasks_offset, 0) - self.tasks_offset
             if next_ == head:
                 break

nitro/backends/linux/process.py

@@ -0,0 +1,20 @@
+from nitro.backends.process import Process
+
+class LinuxProcess(Process):
+    __slots__ = (
+        "task_struct",
+        "name",
+        "pid"
+    )
+
+    def __init__(self, libvmi, cr3, task_struct):
+        super().__init__(libvmi, cr3)
+        pid_offset = self.libvmi.get_offset("linux_pid")
+        name_offset = self.libvmi.get_offset("linux_name")
+
+        self.task_struct = task_struct
+        self.pid = self.libvmi.read_32(self.task_struct + pid_offset, 0)
+        self.name = self.libvmi.read_str_va(self.task_struct + name_offset, 0)
+
+
+

nitro/process.py → nitro/backends/process.py

@@ -1,19 +1,28 @@
 
-class Process:
+class Process():
 
-    def __init__(self, cr3, descriptor, name, pid, libvmi):
-        self.cr3 = cr3
-        self.descriptor = descriptor
-        self.name = name
-        self.pid = pid
+    __slots__ = (
+        "libvmi",
+        "cr3",
+    )
+
+    def __init__(self, libvmi, cr3):
         self.libvmi = libvmi
+        self.cr3 = cr3
+
+    @property
+    def pid(self):
+        raise NotImplementedError("pid must be overridden by a subclass")
+
+    @property
+    def name(self):
+        raise NotImplementedError("name must be overridden by a subclass")
 
     def as_dict(self):
-        info = {
+        return {
             'name': self.name,
             'pid': self.pid
         }
-        return info
 
     def read_memory(self, addr, count):
         return self.libvmi.read_va(addr, self.pid, count)

nitro/backends/windows/backend.py

@@ -11,7 +11,7 @@
 
 from nitro.event import SyscallDirection, SyscallType
 from nitro.syscall import Syscall
-from nitro.process import Process
+from nitro.backends.windows.process import WindowsProcess
 from nitro.backends.backend import Backend
 from nitro.backends.windows.arguments import WindowsArgumentMap
 
@@ -22,7 +22,9 @@ class WindowsBackend(Backend):
         "nb_vcpu",
         "syscall_stack",
         "sdt",
-        "processes"
+        "tasks_offset",
+        "pdbase_offset",
+        "processes",
     )
 
     def __init__(self, domain, libvmi):
@@ -35,7 +37,10 @@ def __init__(self, domain, libvmi):
         self.sdt = None
         self.load_symbols()
 
-        # run libvmi helper subprocess
+        # get offsets
+        self.tasks_offset = self.libvmi.get_offset("win_tasks")
+        self.pdbase_offset = self.libvmi.get_offset("win_pdbase")
+
         self.processes = {}
 
     def process_event(self, event):
@@ -126,22 +131,14 @@ def find_eprocess(self, cr3):
 
         while flink != ps_head:
             # get start of EProcess
-            start_eproc = flink - self.libvmi.get_offset('win_tasks')
+            start_eproc = flink - self.tasks_offset
             # move to start of DirectoryTableBase
-            directory_table_base_off = start_eproc + self.libvmi.get_offset('win_pdbase')
+            directory_table_base_off = start_eproc + self.pdbase_offset
             # read directory_table_base
             directory_table_base = self.libvmi.read_addr_va(directory_table_base_off, 0)
             # compare to our cr3
             if cr3 == directory_table_base:
-                # get name
-                image_file_name_off = start_eproc + self.libvmi.get_offset('win_pname')
-                image_file_name = self.libvmi.read_str_va(image_file_name_off, 0)
-                # get pid
-                unique_processid_off = start_eproc + self.libvmi.get_offset('win_pid')
-                pid = self.libvmi.read_addr_va(unique_processid_off, 0)
-                eprocess = Process(cr3, start_eproc, image_file_name, pid, self.libvmi)
-                return eprocess
-
+                return WindowsProcess(self.libvmi, cr3, start_eproc)
             # read new flink
             flink = self.libvmi.read_addr_va(flink, 0)
         raise RuntimeError('Process not found')

nitro/backends/windows/process.py

@@ -0,0 +1,19 @@
+from nitro.backends.process import Process
+
+class WindowsProcess(Process):
+
+    __slots__ = (
+        "eproc",
+        "name",
+        "pid"
+    )
+
+    def __init__(self, libvmi, cr3, eproc):
+        super().__init__(libvmi, cr3)
+        self.eproc = eproc
+        # get name
+        image_file_name_off = self.eproc + self.libvmi.get_offset('win_pname')
+        self.name = self.libvmi.read_str_va(image_file_name_off, 0)
+        # get pid
+        unique_processid_off = self.eproc + self.libvmi.get_offset('win_pid')
+        self.pid = self.libvmi.read_addr_va(unique_processid_off, 0)