'''
IDAPython script to export DOT files of function graphs.

Example usage (via IDA autonomous mode):
  ida.exe -A -S_export_fimages.py "<output dir>" <fva1> [<fva2> ...] <sample_path>
'''

import os

import idc
import idaapi
import ida_gdl


def main():
    if len(idc.ARGV) < 3:
        # requires output directory and function VAs argument(s)
        idc.qexit(-1)

    # wait for auto-analysis to finish
    idc.auto_wait()

    out_dir = idc.ARGV[1]
    fvas = [int(fva, 0x10) for fva in idc.ARGV[2:]]
    idb_name = os.path.split(idc.get_idb_path())[-1]

    for fva in fvas:
        fstart = idc.get_func_attr(fva, idc.FUNCATTR_START)
        name = '%s_0x%x' % (idb_name.replace('.', '_'), fstart)
        out_path = os.path.join(out_dir, name)
        fname = idc.get_name(fstart)

        if not ida_gdl.gen_flow_graph(out_path, '%s (0x%x)' % (fname, fstart), idaapi.get_func(fstart), 0, 0,
                                      ida_gdl.CHART_GEN_DOT | ida_gdl.CHART_PRINT_NAMES):
            print 'IDA error generating flow graph'
        # TODO add label to DOT file, see https://stackoverflow.com/a/6452088/10548020
        # TODO highlight where rule matched

    # exit IDA
    idc.qexit(0)


if __name__ == '__main__':
    main()