1 files changed, 64 insertions, 2 deletions
diff --git a/tests/test_process.py b/tests/test_process.py
index 276f6f5..57e5b3d 100644
--- a/tests/test_process.py
+++ b/tests/test_process.py
@@ -1,5 +1,67 @@
+import pytest
+
 from hypodermic.process import Process
+from hypodermic.shellcode import assemble
+
+# The program to be tested. Not necessarily important, the binary just
+# has to exist.
+PROCNAME = "/bin/ls"
+
+
+@pytest.fixture
+def process():
+    return Process(path=PROCNAME)
+
+
+# Asserts that instructions can be written at the instruction pointer.
+def test_process_memio(process):
+    if process.arch == "x64":
+        ip = process.get_register("rip")
+        backup = process.read_bytes(ip, 0x10)
+        assert len(backup) == 0x10
+
+        assert process.write_bytes(ip, b'\x90' * 0x10) == 0x10
+        assert process.read_bytes(ip, 0x10) == b'\x90' * 0x10
+
+        assert process.write_bytes(ip, backup) == 0x10
+    else:
+        ip = process.get_register("eip")
+        backup = process.read_bytes(ip, 0x10)
+        assert len(backup) == 0x10
+
+        assert process.write_bytes(ip, b'\x90' * 0x10) == 0x10
+        assert process.read_bytes(ip, 0x10) == b'\x90' * 0x10
+
+        assert process.write_bytes(ip, backup) == 0x10
+
+
+# Asserts that writes to the extended %ax register are persistent.
+def test_process_registers(process):
+    if process.arch == "x64":
+        rax = process.get_register("rax")
+        process.set_register("rax", 0xdeadbeef)
+        assert process.get_register("rax") == 0xdeadbeef
+        process.set_register("rax", rax)
+        assert process.get_register("rax") == rax
+    else:
+        eax = process.get_register("eax")
+        process.set_register("eax", 0xdeadbeef)
+        assert process.get_register("eax") == 0xdeadbeef
+        process.set_register("eax", eax)
+        assert process.get_register("eax") == eax
 
 
-def test_create_process():
-    p = Process(path="/bin/ls")
+# Asserts that code can be run and that its effects are persistent.
+def test_run_code(process):
+    if process.arch == "x64":
+        code = assemble("movq $0xdeadbeef, %rax;")
+        rax = process.get_register("rax")
+        process.run_code(code, preserve=["rax"])
+        assert process.get_register("rax") == 0xdeadbeef
+        process.set_register("rax", rax)
+    else:
+        code = assemble("movl $0xdeadbeef, %eax;")
+        eax = process.get_register("eax")
+        process.run_code(code, preserve=["eax"])
+        assert process.get_register("eax") == 0xdeadbeef
+        process.set_register("eax", eax)