A simple script to check PE files for exploit mitigations (/DYNAMICBASE, /NXCOMPAT, /HIGHENTROPYVA) and anomalies

pe_mitigation_check.py

import argparse
import logging
import pefile
import sys
from prettytable import PrettyTable

########
# Author: David Cannings @edeca
#   Date: September 2018
#
# Checks a PE file for DEP, ASLR, high entropy ASLR and ASLR
# anomalies (/DYNAMICBASE set with relocations stripped).
#
# Requirements:
#   $ pip install pefile PTable
########


def _test_opt_header(pe, flag):
    """
    Test a specific flag in the optional header and return True/False.
    """
    header = pe.OPTIONAL_HEADER.DllCharacteristics
    return (header & pefile.DLL_CHARACTERISTICS[flag]) > 0


def check(fn):
    logging.info("Checking file {}".format(fn))

    try:
        pe = pefile.PE(fn, fast_load=True)
    except pefile.PEFormatError:
        logging.error("Couldn't load this file as a Portable Executable")
        sys.exit(1)
    except OSError:
        logging.error("Couldn't load the specifed file, does it exist?")
        sys.exit(1)

    res = {}
    res['aslr'] = _test_opt_header(pe, 'IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE')
    res['aslr64'] = _test_opt_header(
        pe, 'IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA')
    res['dep'] = _test_opt_header(pe, 'IMAGE_DLLCHARACTERISTICS_NX_COMPAT')
    res['aslr_anomaly'] = False

    if res['aslr'] and pe.FILE_HEADER.Characteristics & pefile.IMAGE_CHARACTERISTICS['IMAGE_FILE_RELOCS_STRIPPED']:
        logging.warning("/DYNAMICBASE is enabled but relocations are stripped")
        res['aslr_anomaly'] = True

    return res


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("file", help="portable executable file to process")
    args = parser.parse_args()
    logging.basicConfig(level=logging.DEBUG)

    res = check(args.file)

    flags = {'aslr': '/DYNAMICBASE', 'dep': '/NXCOMPAT',
             'aslr64': '/HIGHENTROPYVA', 'aslr_anomaly': 'Anomalies?'}
    out = PrettyTable()
    out.field_names = ["Feature", "Enabled"]

    for name, flag in flags.iteritems():
        out.add_row([flag, res[name]])

    print(out)


if __name__ == "__main__":
    main()