import logging


def fix_pe_from_memory(pe, imagebase=None):
    """
    Fixes PE file from memory and returns the pe
    :param pe: pefile object, assuming valid
    :param imagebase: If we want to change the image base, set it here as int
    :return:
    """

    # Fix image base according
    if imagebase is not None:
        pe.OPTIONAL_HEADER.ImageBase = int(imagebase, 16)

    for section in pe.sections:
        # Change section address back to raw
        logging.info('==' + section.Name.decode('utf-8',errors='ignore') + '==')
        logging.info('Modifying virtual addresses:')
        logging.info('{} => {}'.format(hex(section.VirtualAddress), hex(section.PointerToRawData)))
        section.VirtualAddress = section.PointerToRawData
    return pe