"""
Start IDA Pro in autonomous mode to export images of function graphs.

Example usage:
  start_ida_export_fimages.py <target file> <output dir> -f <function list>
  start_ida_export_fimages.py test.exe imgs -f 0x401000,0x402F90
"""

import os
import imp
import sys
import hashlib
import logging
import subprocess

import argparse

try:
    imp.find_module("graphviz")
    from graphviz import Source

    graphviz_found = True
except ImportError:
    graphviz_found = False


IDA32_PATH = "C:\\Program Files\\IDA Pro 7.3\\ida.exe"
IDA64_PATH = "C:\\Program Files\\IDA Pro 7.3\\ida64.exe"

# expected in same directory as this file
EXPORT_SCRIPT_PATH = os.path.abspath("_export_fimages.py")


logger = logging.getLogger(__name__)


def export_fimages(file_path, out_dir, functions, manual=False):
    """
    Export images of function graphs.
    :param file_path: file to analyze
    :param out_dir: output directory
    :param functions: list of strings of hex formatted fvas
    :param manual: non-autonomous mode
    :return: True on success, False otherwise
    """
    if not graphviz_found:
        logger.warning("please install graphviz to export images")
        return False

    if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    script_args = [os.path.abspath(out_dir)] + functions
    call_ida_script(EXPORT_SCRIPT_PATH, script_args, file_path, manual)

    img_count = 0
    for root, dirs, files in os.walk(out_dir):
        for file in files:
            if not file.endswith(".dot"):
                continue
            try:
                s = Source.from_file(file, directory=out_dir)
                s.render(file, directory=out_dir, format="png", cleanup=True)
                img_count += 1
            except BaseException:
                logger.warning("graphviz error rendering file")
    if img_count > 0:
        logger.info('exported %d function graph images to "%s"', img_count, os.path.abspath(out_dir))
        return True
    else:
        logger.warning("failed to export function graph images")
        return False


def call_ida_script(script_path, script_args, sample_path, manual):
    logger.info("processing %s (MD5: %s)", sample_path, get_md5_hexdigest(sample_path))

    # TODO detect 64-bit binaries
    if os.path.splitext(sample_path)[-1] == ".i64":
        IDA_PATH = IDA64_PATH
    else:
        IDA_PATH = IDA32_PATH

    args = [IDA_PATH, "-A", "-S%s %s" % (script_path, " ".join(script_args)), sample_path]

    if manual:
        args.remove("-A")

    logger.debug('calling "%s"' % " ".join(args))
    if subprocess.call(args) == 0:
        return True
    else:
        return False


def get_md5_hexdigest(sample_path):
    m = hashlib.md5()
    with open(sample_path, "rb") as f:
        m.update(f.read())
    return m.hexdigest()


def main():
    parser = argparse.ArgumentParser(
        description="Launch IDA Pro in autonomous mode to export images of function graphs"
    )
    parser.add_argument("file_path", type=str, help="File to export from")
    parser.add_argument("out_dir", type=str, help="Export target directory")
    parser.add_argument("-f", "--functions", action="store", help="Comma separated list of functions to export")
    parser.add_argument("-m", "--manual", action="store_true", help="Manual mode: show IDA dialog boxes")
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")
    args = parser.parse_args(args=sys.argv[1:])

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
        logging.getLogger().setLevel(logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)
        logging.getLogger().setLevel(logging.INFO)

    if not os.path.isfile(args.file_path):
        logger.warning("%s is not a file", args.file_path)
        return -1

    functions = args.functions.split(",")
    export_fimages(args.file_path, args.out_dir, functions, args.manual)

    return 0


if __name__ == "__main__":
    sys.exit(main())