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()