Merge pull request #3 from KVM-VMI/linux_argument_map

Linux argument map

nitro/backends/arguments.py

@@ -1,21 +1,51 @@
-from enum import Enum
+import struct
 
+from enum import Enum
+from nitro.event import SyscallType
 
 class SyscallArgumentType(Enum):
     register = 0
     memory = 1
 
 
 class ArgumentMap:
+
+    ARG_SIZE = {
+        SyscallType.syscall: 'P',   # x64 -> 8 bytes
+        SyscallType.sysenter: 'I'   # x32 -> 4 bytes
+    }
+
     __slots__ = (
         "event",
         "name",
         "process",
-        "modified"
+        "modified",
+        "arg_size_format"
     )
 
     def __init__(self, event, name, process):
         self.event = event
         self.name = name
         self.process = process
         self.modified = {}
+        self.arg_size_format = self.ARG_SIZE[self.event.type]
+
+    def get_argument_value(self, arg_type, opaque):
+        if arg_type == SyscallArgumentType.register:
+            value = self.event.get_register(opaque)
+        else:
+            # memory
+            size = struct.calcsize(self.arg_size_format)
+            addr = self.event.regs.rsp + (opaque * size)
+            value, *rest = struct.unpack(self.arg_size_format, self.process.read_memory(addr, size))
+        return value
+
+    def set_argument_value(self, arg_type, opaque, value):
+        if arg_type == SyscallArgumentType.register:
+            self.event.update_register(opaque, value)
+        else:
+            # memory
+            size = struct.calcsize(self.arg_size_format)
+            addr = self.event.regs.rsp + (opaque * size)
+            buffer = struct.pack(self.arg_size_format, value)
+            self.process.write_memory(addr, buffer)

nitro/backends/linux/arguments.py

@@ -1,8 +1,49 @@
-from nitro.backends.arguments import ArgumentMap
+import struct
+
+from nitro.event import SyscallType
+from nitro.backends.arguments import ArgumentMap, SyscallArgumentType
 
 class LinuxArgumentMap(ArgumentMap):
+
+    CONVENTION = {
+        SyscallType.syscall: [
+            (SyscallArgumentType.register, 'rdi'),
+            (SyscallArgumentType.register, 'rsi'),
+            (SyscallArgumentType.register, 'rdx'),
+            (SyscallArgumentType.register, 'r10'),
+            (SyscallArgumentType.register, 'r9'),
+            (SyscallArgumentType.register, 'r8'),
+        ],
+        SyscallType.sysenter: [
+            (SyscallArgumentType.register, 'rbx'),
+            (SyscallArgumentType.register, 'rcx'),
+            (SyscallArgumentType.register, 'rdx'),
+            (SyscallArgumentType.register, 'rsi'),
+            (SyscallArgumentType.register, 'rdi'),
+            (SyscallArgumentType.register, 'rbp'),
+        ],
+    }
+
     def __getitem__(self, index):
-        raise NotImplementedError()
+        try:
+            arg_type, opaque = self.CONVENTION[self.event.type][index]
+        except KeyError as error:
+            raise RuntimeError('Unknown convention') from error
+        except IndexError:
+            raise RuntimeError('Invalid argument index: Linux syscalls are '
+                               'limited to 6 parameters')
+
+        return self.get_argument_value(arg_type, opaque)
+
 
     def __setitem__(self, index, value):
-        raise NotImplementedError()
+        try:
+            arg_type, opaque = self.CONVENTION[self.event.type][index]
+        except KeyError as error:
+            raise RuntimeError('Unknown convention') from error
+        except IndexError:
+            raise RuntimeError('Invalid argument index: Linux syscalls are '
+                               'limited to 6 parameters')
+
+        self.set_argument_value(arg_type, opaque, value)
+        self.modified[index] = value

nitro/backends/windows/arguments.py

@@ -3,6 +3,7 @@
 from nitro.event import SyscallType
 from nitro.backends.arguments import ArgumentMap, SyscallArgumentType
 
+
 class WindowsArgumentMap(ArgumentMap):
 
     CONVENTION = {
@@ -15,35 +16,19 @@ class WindowsArgumentMap(ArgumentMap):
         ],
     }
 
-    ARG_SIZE = {
-        SyscallType.syscall: 'P',   # x64 -> 8 bytes
-        SyscallType.sysenter: 'I'   # x32 -> 4 bytes
-    }
-
-    __slots__ = (
-        "arg_size_format",
-    )
-
     def __init__(self, event, name, process):
         super().__init__(event, name, process)
-        self.arg_size_format = self.ARG_SIZE[self.event.type]
 
     def __getitem__(self, index):
         try:
             arg_type, opaque = self.CONVENTION[self.event.type][index]
         except KeyError as error:
-            raise RuntimeError('Unknown covention') from error
+            raise RuntimeError('Unknown convention') from error
         except IndexError:
             arg_type, opaque = self.CONVENTION[self.event.type][-1]
             opaque += index - len(self.CONVENTION[self.event.type]) + 1
-        if arg_type == SyscallArgumentType.register:
-            value = self.event.get_register(opaque)
-        else:
-            # memory
-            size = struct.calcsize(self.arg_size_format)
-            addr = self.event.regs.rsp + (opaque * size)
-            value, *rest = struct.unpack(self.arg_size_format, self.process.read_memory(addr, size))
-        return value
+
+        return self.get_argument_value(arg_type, opaque)
 
     def __setitem__(self, index, value):
         try:
@@ -53,13 +38,6 @@ def __setitem__(self, index, value):
         except IndexError:
             arg_type, opaque = self.CONVENTION[self.event.type][-1]
             opaque += index - len(self.CONVENTION[self.event.type]) + 1
-        if arg_type == SyscallArgumentType.register:
-            self.event.update_register(opaque, value)
-        else:
-            # memory
-            size = struct.calcsize(self.arg_size_format)
-            addr = self.event.regs.rsp + (opaque * size)
-            buffer = struct.pack(self.arg_size_format, value)
-            self.process.write_memory(addr, buffer)
-        self.modified[index] = value
 
+        self.set_argument_value(arg_type, opaque, value)
+        self.modified[index] = value