Merge pull request #2 from KVM-VMI/refactor_windows_process

Refactor windows process
Soft · Jun 29, 2017 · f09465c · f09465c
2 parents 6146b5d + ebdbddd
commit f09465c
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 20 deletions.
diff --git a/nitro/backends/linux/backend.py b/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.process import Process
 from nitro.backends.backend import Backend
 from nitro.backends.linux.arguments import LinuxArgumentMap
 

diff --git a/nitro/process.py → nitro/backends/process.py b/nitro/process.py → nitro/backends/process.py
@@ -1,12 +1,15 @@
-
 class Process:
 
-    def __init__(self, cr3, descriptor, name, pid, libvmi):
+    __slots__ = (
+        "libvmi",
+        "cr3",
+        "descriptor",
+    )
+
+    def __init__(self, libvmi, cr3, descriptor):
+        self.libvmi = libvmi
         self.cr3 = cr3
         self.descriptor = descriptor
-        self.name = name
-        self.pid = pid
-        self.libvmi = libvmi
 
     def as_dict(self):
         info = {

diff --git a/nitro/backends/windows/backend.py b/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')

diff --git a/nitro/backends/windows/process.py b/nitro/backends/windows/process.py
@@ -0,0 +1,18 @@
+from nitro.backends.process import Process
+
+class WindowsProcess(Process):
+
+    __slots__ = (
+        "name",
+        "pid",
+    )
+
+    # descriptor is the start address of the EPROCESS struct
+    def __init__(self, libvmi, cr3, descriptor):
+        super().__init__(libvmi, cr3, descriptor)
+        # get name
+        image_file_name_off = self.descriptor + self.libvmi.get_offset('win_pname')
+        self.name = self.libvmi.read_str_va(image_file_name_off, 0)
+        # get pid
+        unique_processid_off = self.descriptor + self.libvmi.get_offset('win_pid')
+        self.pid = self.libvmi.read_addr_va(unique_processid_off, 0)