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
from model.defs import *
from model.carrier import Carrier, DataReuseEntry
from model.injectable import Injectable, DataReuseEntry
from observer import observer
from helper import *
from phases.asmtextparser import parse_asm_text_file
@@ -25,11 +25,11 @@ class AsmTest(unittest.TestCase):
def test_asm_fixup(self):
asm_in: FilePath = "tests/data/peb_walk_pre_fixup.asm"
asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
injectable = Injectable("fake.exe")
injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings()
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], 272 ; 00ab4130H
@@ -47,19 +47,19 @@ class AsmTest(unittest.TestCase):
def test_asm_iat_request(self):
asm_in: FilePath = "tests/data/iat_reuse_pre_fixup.asm"
asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
injectable = Injectable("fake.exe")
injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings()
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.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.assertTrue(len(req2.references[0]), 6) # 6 random bytes
@@ -79,14 +79,14 @@ class AsmTest(unittest.TestCase):
def test_data_reuse_entries(self):
asm_in = "tests/data/data_reuse_pre_fixup.asm"
asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
injectable = Injectable("fake.exe")
injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings()
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
data_reuse_entries = carrier.get_all_reusedata_fixups()
data_reuse_entries = injectable.get_all_reusedata_fixups()
self.assertEqual(2+1, len(data_reuse_entries))
entry = data_reuse_entries[0+1]
@@ -104,11 +104,11 @@ class AsmTest(unittest.TestCase):
asm_in = "tests/data/data_reuse_pre_fixup.asm"
asm_text = file_readall_text(asm_in)
carrier = Carrier("fake.exe")
carrier.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
injectable = Injectable("fake.exe")
injectable.add_datareuse_fixup(DataReuseEntry("supermega_payload"))
settings: Settings = Settings()
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??
self.assertTrue("\tDB " in asm_text_lines[108-1-1])