Make the type coercions in compiler more explicit

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)