injector/analyzer/main.py

84 lines
2.7 KiB
Python

import base64
from io import BytesIO
import json
import pefile
def main():
subject = pefile.PE("subjects\\main.exe")
def _single_or_none(x):
items = [i for i in x]
if len(items) == 0:
return None
assert len(items) == 1
return items[0]
def _dump(fname, section):
with open(fname, "wb") as f:
if section is not None:
if isinstance(section, bytes):
f.write(section)
else:
f.write(section.get_data())
for i in subject.sections:
print(i)
text_section = _single_or_none(i for i in subject.sections if i.Name == b".text\0\0\0")
data_section = _single_or_none(i for i in subject.sections if i.Name == b".data\0\0\0")
rdata_section = _single_or_none(i for i in subject.sections if i.Name == b".rdata\0\0")
_dump("dumps\\text.dat", text_section)
_dump("dumps\\data.dat", data_section)
_dump("dumps\\rdata.dat", rdata_section)
relevant_sections = [section for section in (text_section, data_section, rdata_section) if section is not None]
if len(relevant_sections) == 0:
raise ValueError("no sections to plot")
min_address = min(i.VirtualAddress for i in relevant_sections)
max_address = max(_round_up_to_page(i.VirtualAddress + i.SizeOfRawData) for i in relevant_sections)
print(min_address, max_address)
buffer = bytearray(max_address - min_address)
for section in relevant_sections:
data = section.get_data() # TODO: De-pad the text section from 0xccs
start = section.VirtualAddress - min_address
buffer[start:start+len(data)] = data
buffer = bytes(buffer)
_dump("dumps\\starting_state.dat", buffer)
binary = {
"startingState": base64.b64encode(buffer).decode("utf8"),
"imports": [],
}
# find imports
# print(subject)
# print(dir(subject))
for entry in subject.DIRECTORY_ENTRY_IMPORT:
# print(entry.dll)
for imp in entry.imports:
# print(dir(imp))
import_address = imp.address - subject.OPTIONAL_HEADER.ImageBase - min_address
print(hex(import_address), imp.name)
binary["imports"].append({
"dll": entry.dll.decode("utf8"),
"symbol": imp.name.decode("utf8"),
"address": import_address,
})
entry_point_rva = subject.OPTIONAL_HEADER.AddressOfEntryPoint
binary["entryPoint"] = entry_point_rva - min_address
with open("binaries/main.json", "wt") as f:
f.write(json.dumps(binary, indent=4))
def _round_up_to_page(x: int):
# TODO: Is this the page size on x64? I think it is
return ((x + 0x1000 - 1) // 0x1000) * 0x1000
if __name__ == "__main__":
main()