import argparse


def alignmentUp(size, alignment):
    if size % alignment != 0:
        size += alignment
    return size


def rva2Offset(sections, alignment, rva):
    for section in sections:
        begin = section['virtualAddress']
        end = begin + alignmentUp(section['virtualSize'], alignment)
        if begin < rva < end:
            return rva - begin + section['pointerToRawData']

    return 0

def va2Offset(sections, imageBase, alignment, va):
    return rva2Offset(sections, alignment, va - imageBase)




def main():
    p = argparse.ArgumentParser(description='PE format parser')
    p.add_argument('filename')

    args = p.parse_args()

    with open(args.filename, 'rb') as f:

        assert f.read(2) == b'MZ'
        f.seek(60)
        lfanew = int.from_bytes(f.read(4), byteorder='little')
        f.seek(lfanew)

        # IMAGE_NT_HEADERS
        assert f.read(4) == b'PE\0\0'

        #skip IMAGE_FILE_HEADER
        f.seek(f.tell() + 20)

        f.seek(f.tell() + 2)
        numberOfSections = int.from_bytes(f.read(2), byteorder='little')
        f.seek(f.tell() - 4)



        #IMAGE_OPTIONAL_HEADER32
        assert f.read(2) == b'\x0b\x01'
        f.seek(f.tell() - 2)
        f.seek(f.tell() + 0x1c)
        imageBase = int.from_bytes(f.read(4), byteorder='little')
        sectionAlignment = int.from_bytes(f.read(4), byteorder='little')
        f.seek(f.tell() - 0x24)
        f.seek(f.tell() + 96)
        exportAddress = int.from_bytes(f.read(4), byteorder='little')
        exportSize = int.from_bytes(f.read(4), byteorder='little')
        print("virtual address {}".format(hex(exportAddress)))
        print("size {}".format(hex(exportSize)))
        print("numberOfSections {}".format(numberOfSections))
        print("imageBase {}".format(hex(imageBase)))
        print("sectionAlignment {}".format(sectionAlignment))
        f.seek(f.tell() - 8)
        f.seek(f.tell() + 8*16)


        sections = []

        for idx in range(0, numberOfSections):
            name = f.read(8).decode()
            virtualSize = int.from_bytes(f.read(4), byteorder='little')
            virtualAddress = int.from_bytes(f.read(4), byteorder='little')
            sizeOfRawData = int.from_bytes(f.read(4), byteorder='little')
            pointerToRawData = int.from_bytes(f.read(4), byteorder='little')
            pointerToRelocations = int.from_bytes(f.read(4), byteorder='little')
            pointerToLinenumbers = int.from_bytes(f.read(4), byteorder='little')
            numberOfRelocations = int.from_bytes(f.read(2), byteorder='little')
            numberOfLinenumbers = int.from_bytes(f.read(2), byteorder='little')
            characteristics = int.from_bytes(f.read(4), byteorder='little')
            sections.append({
                'name': name,
                'virtualSize': virtualSize,
                'virtualAddress': virtualAddress,
                'sizeOfRawData': sizeOfRawData,
                'pointerToRawData': pointerToRawData,
                'pointerToRelocations': pointerToRelocations,
                'pointerToLinenumbers': pointerToLinenumbers,
                'numberOfRelocations': numberOfRelocations,
                'numberOfLinenumbers': numberOfLinenumbers,
                'characteristics': characteristics
            })

        print("sections: {}".format(sections))


        f.seek(rva2Offset(sections, sectionAlignment, exportAddress))

        f.seek(f.tell() + 0x14)
        numberOfFunctions = int.from_bytes(f.read(4), byteorder='little')
        numberOfName = int.from_bytes(f.read(4), byteorder='little')
        addressOfFunctions = int.from_bytes(f.read(4), byteorder='little')
        addressOfNames = int.from_bytes(f.read(4), byteorder='little')
        addressOfNameOrdinals = int.from_bytes(f.read(4), byteorder='little')
        print("numberOfFunctions {}".format(numberOfFunctions))
        print("numberOfName {}".format(numberOfName))
        print("addressOfFunctions {}".format(hex(addressOfFunctions)))
        print("addressOfNames {}".format(hex(addressOfNames)))
        print("addressOfNameOrdinals {}".format(hex(addressOfNameOrdinals)))

        def readString():

            str = ''
            ch = f.read(1)
            while ch != b'\x00':
                str += ch.decode()
                ch = f.read(1)

            return str

        for i in range(0, numberOfFunctions):
            f.seek(rva2Offset(sections, sectionAlignment, addressOfFunctions) + 0x4 * i)
            functionAddress = f.read(4).hex()
            f.seek(rva2Offset(sections, sectionAlignment, addressOfNameOrdinals) + 0x2 * i)
            functionOrdinal = f.read(2).hex()
            f.seek(rva2Offset(sections, sectionAlignment, addressOfNames) + 0x4 * i)
            functionNameAddress = int.from_bytes(f.read(4), byteorder='little')
            f.seek(rva2Offset(sections, sectionAlignment, functionNameAddress))
            functionName = readString()

            print("0x{} 0x{} {}".format(functionAddress, functionOrdinal, functionName))
















if __name__ == '__main__':
    main()