fix: unittest

This commit is contained in:
Dobin Rutishauser
2024-06-23 17:10:06 +02:00
parent fc64068139
commit 720d85c994
+17 -17
View File
@@ -3,7 +3,7 @@ import unittest
import logging import logging
from model.defs import * from model.defs import *
from model.carrier import Carrier, DataReuseEntry from model.injectable import Injectable, DataReuseEntry
from observer import observer from observer import observer
from helper import * from helper import *
from phases.asmtextparser import parse_asm_text_file from phases.asmtextparser import parse_asm_text_file
@@ -25,11 +25,11 @@ class AsmTest(unittest.TestCase):
def test_asm_fixup(self): def test_asm_fixup(self):
asm_in: FilePath = "tests/data/peb_walk_pre_fixup.asm" asm_in: FilePath = "tests/data/peb_walk_pre_fixup.asm"
asm_text = file_readall_text(asm_in) asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe") injectable = Injectable("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload")) injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings() settings: Settings = Settings()
settings.payload_location = PayloadLocation.DATA settings.payload_location = PayloadLocation.DATA
asm_text_lines = parse_asm_text_file(carrier, asm_text, settings) asm_text_lines = parse_asm_text_file(injectable, asm_text, settings)
# cmp DWORD PTR n$1[rsp], 11223344 ; 00ab4130H # cmp DWORD PTR n$1[rsp], 11223344 ; 00ab4130H
# cmp DWORD PTR n$1[rsp], 272 ; 00ab4130H # cmp DWORD PTR n$1[rsp], 272 ; 00ab4130H
@@ -47,19 +47,19 @@ class AsmTest(unittest.TestCase):
def test_asm_iat_request(self): def test_asm_iat_request(self):
asm_in: FilePath = "tests/data/iat_reuse_pre_fixup.asm" asm_in: FilePath = "tests/data/iat_reuse_pre_fixup.asm"
asm_text = file_readall_text(asm_in) asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe") injectable = Injectable("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload")) injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings() settings: Settings = Settings()
settings.payload_location = PayloadLocation.DATA settings.payload_location = PayloadLocation.DATA
asm_text_lines = parse_asm_text_file(carrier, asm_text, settings) asm_text_lines = parse_asm_text_file(injectable, asm_text, settings)
self.assertEqual(len(carrier.iat_requests), 2) self.assertEqual(len(injectable.iat_requests), 2)
req1 = carrier.iat_requests[0] req1 = injectable.iat_requests[0]
self.assertEqual(req1.name, "GetEnvironmentVariableW") self.assertEqual(req1.name, "GetEnvironmentVariableW")
self.assertTrue(len(req1.references[0]), 6) # 6 random bytes self.assertTrue(len(req1.references[0]), 6) # 6 random bytes
req2 = carrier.iat_requests[1] req2 = injectable.iat_requests[1]
self.assertEqual(req2.name, "VirtualAlloc") self.assertEqual(req2.name, "VirtualAlloc")
self.assertTrue(len(req2.references[0]), 6) # 6 random bytes self.assertTrue(len(req2.references[0]), 6) # 6 random bytes
@@ -79,14 +79,14 @@ class AsmTest(unittest.TestCase):
def test_data_reuse_entries(self): def test_data_reuse_entries(self):
asm_in = "tests/data/data_reuse_pre_fixup.asm" asm_in = "tests/data/data_reuse_pre_fixup.asm"
asm_text = file_readall_text(asm_in) asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe") injectable = Injectable("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload")) injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings() settings: Settings = Settings()
settings.payload_location = PayloadLocation.DATA settings.payload_location = PayloadLocation.DATA
asm_text_lines = parse_asm_text_file(carrier, asm_text, settings) asm_text_lines = parse_asm_text_file(injectable, asm_text, settings)
asm_text = masm_shc(asm_text_lines) # optional here asm_text = masm_shc(asm_text_lines) # optional here
data_reuse_entries = carrier.get_all_reusedata_fixups() data_reuse_entries = injectable.get_all_reusedata_fixups()
self.assertEqual(2+1, len(data_reuse_entries)) self.assertEqual(2+1, len(data_reuse_entries))
entry = data_reuse_entries[0+1] entry = data_reuse_entries[0+1]
@@ -104,11 +104,11 @@ class AsmTest(unittest.TestCase):
asm_in = "tests/data/data_reuse_pre_fixup.asm" asm_in = "tests/data/data_reuse_pre_fixup.asm"
asm_text = file_readall_text(asm_in) asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe") injectable = Injectable("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload")) injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings() settings: Settings = Settings()
settings.payload_location = PayloadLocation.DATA settings.payload_location = PayloadLocation.DATA
asm_text_lines = parse_asm_text_file(carrier, asm_text, settings) asm_text_lines = parse_asm_text_file(injectable, asm_text, settings)
# why -1 -1?? # why -1 -1??
self.assertTrue("\tDB " in asm_text_lines[108-1-1]) self.assertTrue("\tDB " in asm_text_lines[108-1-1])