diff --git a/compiler/main.py b/compiler/main.py index 86a6ca2..6a55c4d 100644 --- a/compiler/main.py +++ b/compiler/main.py @@ -1,9 +1,7 @@ -import base64 from dataclasses import dataclass from io import BytesIO -import json from typing import Generator -import pefile +from pefile import PE, Structure @dataclass class Import: @@ -36,38 +34,40 @@ def _single[T](ts: Generator[T]) -> T: raise ValueError(f"expected 1, got {len(items)}") -def _create_binary(subject: pefile.PE) -> Binary: +def _create_binary(subject: PE) -> Binary: optional_header = subject.OPTIONAL_HEADER - assert isinstance(optional_header, pefile.Structure) - text_section = _single(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") + assert isinstance(optional_header, Structure) + text_section: Structure = _single(i for i in subject.sections if i.Name == b".text\0\0\0") + data_section: Structure | None = _single_or_none(i for i in subject.sections if i.Name == b".data\0\0\0") + rdata_section: Structure | None = _single_or_none(i for i in subject.sections if i.Name == b".rdata\0\0") - relevant_sections = [section for section in (text_section, data_section, rdata_section) if section is not None] + relevant_sections: list[Structure] = [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(i.VirtualAddress + i.SizeOfRawData for i in relevant_sections) + min_address = min(getattr(i, "VirtualAddress") for i in relevant_sections) + max_address = max(getattr(i, "VirtualAddress") + getattr(i, "SizeOfRawData") for i in relevant_sections) buffer = bytearray(max_address - min_address) for section in relevant_sections: - data = section.get_data() - start = section.VirtualAddress - min_address + data = getattr(section, "get_data")() + start = getattr(section, "VirtualAddress") - min_address buffer[start:start+len(data)] = data starting_state = bytes(buffer) entry_point_rva = getattr(optional_header, "AddressOfEntryPoint") - entry_point = (entry_point_rva - min_address) + entry_point = entry_point_rva - min_address imports: list[Import] = [] for entry in getattr(subject, "DIRECTORY_ENTRY_IMPORT"): - library: bytes = entry.dll + entry: Structure + library: bytes = getattr(entry, "dll") procedures: list[tuple[bytes, int]] = [] - for imp in entry.imports: - import_address_rva = imp.address - getattr(optional_header, "ImageBase") + for imp in getattr(entry, "imports"): + imp: Structure + import_address_rva = getattr(imp, "address") - getattr(optional_header, "ImageBase") import_address = import_address_rva - min_address - procedures.append((imp.name, import_address)) + procedures.append((getattr(imp, "name"), import_address)) imports.append(Import(library, procedures)) @@ -151,7 +151,7 @@ def _encode_binary(binary: Binary) -> bytes: return out.getbuffer() def main(): - subject = pefile.PE("subjects\\main.exe") + subject = PE("subjects\\main.exe") binary = _create_binary(subject) code = _encode_binary(binary)