Initial unit testing

4 files changed, 89 insertions, 4 deletions

diff --git a/hypodermic/main.py b/hypodermic/main.py
index 3cbb047..7a671f6 100644
--- a/hypodermic/main.py
+++ b/hypodermic/main.py
@@ -111,5 +111,4 @@ def main():
     else:
         alert("Attaching to process with pid {}...".format(args.attach))
         p = Process(pid=args.attach)
-    print(p.dlopen("/usr/lib/libyggdrasil.so"))
     p.continue_until_haulted()
diff --git a/hypodermic/memory.py b/hypodermic/memory.py
index 0d86855..56a3fc3 100644
--- a/hypodermic/memory.py
+++ b/hypodermic/memory.py
@@ -48,7 +48,7 @@ def parse_device(line: str) -> Device:
         The parsed Device object.
     """
     major, minor = line.split(':')
-    return Device(major, minor)
+    return Device(int(major), int(minor))
 
 
 def parse_perms(line: str) -> Perms:
diff --git a/tests/test_memory.py b/tests/test_memory.py
new file mode 100644
index 0000000..46a22b7
--- /dev/null
+++ b/tests/test_memory.py
@@ -0,0 +1,24 @@
+from hypodermic.memory import Device, Perms, Region, parse_region
+
+TEST_LINE = "00400000-00408000 r-xp 00000000 08:12 35923730    /usr/bin/cat"
+
+
+def test_parse_line():
+    res = parse_region(TEST_LINE)
+    assert isinstance(res, Region)
+    assert isinstance(res.dev, Device)
+    assert isinstance(res.perms, Perms)
+
+    assert res.start == 0x00400000
+    assert res.end == 0x00408000
+
+    assert res.perms.r
+    assert not res.perms.w
+    assert res.perms.x
+    assert not res.perms.s
+
+    assert res.off == 0
+    assert res.dev.major == 8
+    assert res.dev.minor == 12
+    assert res.inode == 35923730
+    assert res.path == "/usr/bin/cat"
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)