"""
Start IDA Pro in autonomous mode to dump JSON file of function names { fva: fname }.
Processes a single file or a directory.
Only runs on files with supported file extensions.

Example usage:
  start_ida_dump_fnames.py <candidate files dir>
  start_ida_dump_fnames.py samples\benign
"""

import os
import sys
import json
import hashlib
import logging
import subprocess

import argparse

from scripts.testbed import FNAMES_EXTENSION

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
DUMP_SCRIPT_PATH = os.path.abspath("_dump_fnames.py")

SUPPORTED_EXTENSIONS = [
    ".exe_",
    ".dll_",
    ".sys_",
    ".idb",
    ".i64",
]


logger = logging.getLogger(__name__)


def call_ida_dump_script(sample_path, reprocess):
    """ call IDA in autonomous mode and return True if success, False on failure """
    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

    if sample_path.endswith(".idb") or sample_path.endswith(".i64"):
        sample_path = sample_path[:-4]

    fnames = "%s%s" % (sample_path, FNAMES_EXTENSION)
    if os.path.exists(fnames) and not reprocess:
        logger.info(
            "%s already exists and contains %d function names, provide -r argument to reprocess",
            fnames,
            len(get_function_names(fnames)),
        )
        return True

    out_path = os.path.split(fnames)[-1]  # relative to IDA database file
    args = [IDA_PATH, "-A", '-S%s "%s"' % (DUMP_SCRIPT_PATH, out_path), sample_path]
    logger.debug('calling "%s"' % " ".join(args))
    subprocess.call(args)

    if not os.path.exists(fnames):
        logger.warning("%s was not created", fnames)
        return False

    logger.debug("extracted %d function names to %s", len(get_function_names(fnames)), fnames)
    return True


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


def get_function_names(fnames_file):
    if not os.path.exists(fnames_file):
        return None
    with open(fnames_file, "r") as f:
        return json.load(f)


def main():
    parser = argparse.ArgumentParser(
        description="Launch IDA Pro in autonomous mode to dump function names of a file or of files in a directory"
    )
    parser.add_argument("file_path", type=str, help="File or directory path to analyze")
    parser.add_argument("-r", "--reprocess", action="store_true", default=False, help="Overwrite existing analysis")
    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.exists(args.file_path):
        logger.warning("%s does not exist", args.file_path)
        return -1

    if os.path.isfile(args.file_path):
        call_ida_dump_script(args.file_path, args.reprocess)
        return 0

    errors = 0

    logger.info("processing files in %s with file extension %s", args.file_path, "|".join(SUPPORTED_EXTENSIONS))
    for root, dirs, files in os.walk(args.file_path):
        for file in files:
            if not os.path.splitext(file)[1] in SUPPORTED_EXTENSIONS:
                logger.debug("%s does not have supported file extension", file)
                continue
            path = os.path.join(root, file)
            if not call_ida_dump_script(path, args.reprocess):
                errors += 1

    if errors:
        logger.warning("encountered %d errors", errors)

    return 0


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