test: add unittests

.vscode/settings.json

@@ -0,0 +1,11 @@
+{
+    "python.testing.unittestArgs": [
+        "-v",
+        "-s",
+        "./tests",
+        "-p",
+        "test_*.py"
+    ],
+    "python.testing.pytestEnabled": false,
+    "python.testing.unittestEnabled": true
+}

requirements.txt

@@ -1,3 +1,4 @@
+pytest
 pefile
 capstone
 keystone-engine

tests/test_asm.py

@@ -0,0 +1,67 @@
+import shutil
+from typing import List
+import unittest
+import logging
+
+from phases.compiler import fixup_asm_file, fixup_iat_reuse
+from model import ExeInfo
+from defs import *
+
+
+class AsmTest(unittest.TestCase):
+    def test_asm_fixup(self):
+        path_in: FilePath = "tests/data/peb_walk_pre_fixup.asm"
+        path_working: FilePath = "tests/data/peb_walk_pre_fixup.asm.test"
+
+        shutil.copy(path_in, path_working)
+        fixup_asm_file(path_working, 272)
+        with open(path_working, "r") as f:
+            lines = f.readlines()
+
+        # cmp     DWORD PTR n$1[rsp], 11223344            ; 00ab4130H
+        # cmp     DWORD PTR n$1[rsp], 272         ; 00ab4130H
+        self.assertTrue(", 272" in lines[192-1])
+        self.assertTrue("11223344" not in lines[192-1])
+        
+        # mov     r8, QWORD PTR supermega_payload
+        # lea	r8, [shcstart]
+        self.assertTrue("lea	r8, [shcstart]" in lines[198-1])
+        self.assertTrue("supermega_payload" not in lines[198-1])
+
+        # shcstart:
+        self.assertTrue("shcstart:" in lines[213-1])
+
+        os.remove(path_working)
+
+
+    def test_asm_iat_fixup(self):
+        path_in: FilePath = "tests/data/iat_reuse_pre_fixup.asm"
+        path_working: FilePath = "tests/data/iat_reuse_pre_fixup.asm.test"
+        shutil.copy(path_in, path_working)
+
+        exe_info = ExeInfo()
+
+        fixup_iat_reuse(path_working, exe_info)
+        self.assertTrue(len(exe_info.iat_resolves), 2)
+
+        self.assertTrue("GetEnvironmentVariableW" in exe_info.iat_resolves)
+        self.assertEqual(exe_info.iat_resolves["GetEnvironmentVariableW"].name, "GetEnvironmentVariableW")
+        self.assertEqual(exe_info.iat_resolves["GetEnvironmentVariableW"].addr, 0)
+        self.assertTrue(len(exe_info.iat_resolves["GetEnvironmentVariableW"].id), 6) # 6 random bytes
+
+        with open(path_working, "r") as f:
+            lines = f.readlines()
+        
+        # added ; at the beginning
+        self.assertTrue(lines[13-1].startswith("; EXTRN	__imp_GetEnvironmentVariableW:PROC"))
+        self.assertTrue(lines[14-1].startswith("; EXTRN	__imp_VirtualAlloc:PROC"))
+
+        # 	call	QWORD PTR __imp_GetEnvironmentVariableW
+        # 	DB 044H, 0aeH, 06cH, 0b6H, 072H, 07cH
+        self.assertTrue(lines[158-1].startswith("	DB "))
+
+        # 	call	QWORD PTR __imp_VirtualAlloc
+        # 	DB 0c7H, 0b6H, 0feH, 0dcH, 0b2H, 0c6H
+        self.assertTrue(lines[183-1].startswith("	DB "))
+
+        os.remove(path_working)