mirror of
https://github.com/dobin/SuperMega
synced 2026-06-02 17:27:10 +00:00
Dobin Rutishauser
96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
from typing import Dict, List
|
|
import logging
|
|
import pefile
|
|
|
|
from model.defs import *
|
|
from pe.superpe import SuperPe, PeSection
|
|
|
|
|
|
logger = logging.getLogger("Injectable")
|
|
|
|
|
|
class IatRequest():
|
|
def __init__(self, name: str, placeholder: bytes):
|
|
self.name: str = name # Function Name, like "VirtualAlloc"
|
|
self.references: List[bytes] = []
|
|
self.add_reference(placeholder)
|
|
|
|
def add_reference(self, placeholder):
|
|
self.references.append(placeholder)
|
|
|
|
|
|
class DataReuseReference():
|
|
def __init__(self, placeholder: bytes, register: str):
|
|
self.placeholder: bytes = placeholder
|
|
self.register: str = register
|
|
|
|
|
|
class DataReuseEntry():
|
|
def __init__(self, string_ref: str, in_code: bool = False):
|
|
self.string_ref: str = string_ref # "$SG72513"
|
|
self.data: bytes = b'' # the content/data
|
|
self.addr: int = 0 # where content/data is stored
|
|
self.in_code: bool = in_code # is the data in code section
|
|
|
|
self.references: List[DataReuseReference] = []
|
|
|
|
|
|
def add_reference(self, placeholder, register):
|
|
self.references.append(DataReuseReference(placeholder, register))
|
|
|
|
|
|
class Injectable():
|
|
def __init__(self, exe_file: str):
|
|
self.iat_requests: List[IatRequest] = []
|
|
self.reusedata_fixups: List[DataReuseEntry] = []
|
|
self.exe_filepath: str = exe_file
|
|
self.superpe: SuperPe = None
|
|
|
|
def init(self) -> bool:
|
|
# check if file exists
|
|
if not os.path.exists(self.exe_filepath):
|
|
logger.error("Injectable file does not exist: {}".format(self.exe_filepath))
|
|
return False
|
|
self.superpe = SuperPe(self.exe_filepath)
|
|
return True
|
|
|
|
|
|
# IAT
|
|
|
|
def add_iat_request(self, func_name: str, placeholder: bytes):
|
|
# existing?
|
|
for iat in self.iat_requests:
|
|
if iat.name == func_name:
|
|
iat.add_reference(placeholder)
|
|
return
|
|
|
|
# new
|
|
self.iat_requests.append(IatRequest(func_name, placeholder))
|
|
|
|
def get_all_iat_requests(self) -> List[IatRequest]:
|
|
return self.iat_requests
|
|
|
|
def get_unresolved_iat(self):
|
|
"""Returns a list of IAT entries not available in the PE file"""
|
|
functions = []
|
|
for iat in self.iat_requests:
|
|
if self.superpe.get_vaddr_of_iatentry(iat.name) == None:
|
|
functions.append(iat.name)
|
|
return functions
|
|
|
|
|
|
# Data Reuse
|
|
|
|
def add_datareuse_fixup(self, fixup: DataReuseEntry):
|
|
logger.debug("---( Add datareuse: {}".format(fixup.string_ref))
|
|
self.reusedata_fixups.append(fixup)
|
|
|
|
def get_all_reusedata_fixups(self) -> List[DataReuseEntry]:
|
|
return self.reusedata_fixups
|
|
|
|
def get_reusedata_fixup(self, string_ref) -> DataReuseEntry:
|
|
for entry in self.reusedata_fixups:
|
|
if entry.string_ref == string_ref:
|
|
return entry
|
|
return None
|
|
|