diff --git a/model/carrier.py b/model/carrier.py index ca1a017..c109fa1 100644 --- a/model/carrier.py +++ b/model/carrier.py @@ -1,7 +1,6 @@ from typing import Dict, List import logging -from model.exehost import DataReuseEntry logger = logging.getLogger("Carrier") @@ -12,6 +11,14 @@ class IatRequest(): self.placeholder: bytes = placeholder # Random bytes as placeholder +class DataReuseEntry(): + def __init__(self, string_ref: str, register: str, randbytes: bytes): + self.string_ref = string_ref + self.register = register + self.randbytes = randbytes + self.data = b'' + self.addr = 0 + class Carrier(): def __init__(self): diff --git a/model/exehost.py b/model/exehost.py index c7065fc..1125d8c 100644 --- a/model/exehost.py +++ b/model/exehost.py @@ -18,19 +18,10 @@ class RelocEntry(): class IatEntry(): - def __init__(self, dll_name, func_name, func_addr): - self.dll_name = dll_name - self.func_name = func_name - self.func_addr = func_addr - - -class DataReuseEntry(): - def __init__(self, string_ref: str, register: str, randbytes: bytes): - self.string_ref = string_ref - self.register = register - self.randbytes = randbytes - self.data = b'' - self.addr = 0 + def __init__(self, dll_name: str, func_name: str, iat_vaddr: int): + self.dll_name: str = dll_name + self.func_name: str = func_name + self.iat_vaddr: int = iat_vaddr class ExeHost(): @@ -42,7 +33,7 @@ class ExeHost(): self.pe: pefile.PE = None self.superpe: SuperPe = None - self.iat = {} # Dict[str, List[Dict[str, str]]] + self.iat: Dict[str, IatEntry] = {} self.base_relocs = [] self.image_base: int = 0 @@ -128,26 +119,13 @@ class ExeHost(): if not dll_name in self.iat: self.iat[dll_name] = [] - self.iat[dll_name].append({ - "dll_name": dll_name, - "func_name": imp_name, - "iat_vaddr": imp_addr - }) + self.iat[dll_name].append(IatEntry(dll_name, imp_name, imp_addr)) def get_vaddr_of_iatentry(self, func_name: str) -> int: for dll_name in self.iat: for entry in self.iat[dll_name]: - if entry["func_name"] == func_name: - return entry["iat_vaddr"] + if entry.func_name == func_name: + return entry.iat_vaddr return None - - ## Other - - def print(self): - logger.info("--( Required IAT Resolves: ") - for _, cap in self.iat_requests.items(): - if cap.addr == 0: - logger.info(" {:28} {}".format(cap.name, "N/A")) - else: - logger.info(" {:28} 0x{:x}".format(cap.name, cap.addr)) + \ No newline at end of file diff --git a/phases/datareuse.py b/phases/datareuse.py index 449f491..e7b5409 100644 --- a/phases/datareuse.py +++ b/phases/datareuse.py @@ -4,7 +4,7 @@ from intervaltree import Interval, IntervalTree from typing import List, Dict import os -from model.exehost import DataReuseEntry +from model.carrier import DataReuseEntry def bytes_to_asm_db(byte_data: bytes) -> bytes: diff --git a/phases/injector.py b/phases/injector.py index 70c7efb..0bfcb33 100644 --- a/phases/injector.py +++ b/phases/injector.py @@ -6,7 +6,7 @@ import time import tempfile import logging -from model.carrier import Carrier +from model.carrier import Carrier, DataReuseEntry from peparser.pehelper import * from model.exehost import * from observer import observer