"""
IDAPython script to annotate IDB files with information extracted from the
MSDN documentation including functions, arguments and constants.

Authors: Moritz Raabe, William Ballenthin
Copyright 2014 Mandiant, A FireEye Company

Mandiant licenses this file to you under the Apache License, Version
2.0 (the "License"); you may not use this file except in compliance with the
License. You may obtain a copy of the License at:

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied. See the License for the specific language governing
permissions and limitations under the License.
"""

import os.path
import logging
import idc
import idautils
import idaapi
from textwrap import fill
from time import strftime
import IDB_MSDN_Annotator.xml_parser

CREATE_BACKUP = True  # indicate if a backup should be created
# The file should be located in the the MSDN_data directory
MSDN_INFO_FILE = 'msdn_data.xml'
NEW_SEGMENT_SIZE = 0x800  # size of the new segment
NEW_SEGMENT_NAME = '.msdn'  # name of the new segment
COMMENT_MAX_WIDTH = 40  # maximum column width of comments
MAX_ARG_DESCR_LEN = 600  # maximum string length of argument descriptions
PREVIOUS_INSTR_DELTA = 0xF  # range in where to look for previous instruction
ARG_SEARCH_THRESHOLD = 0xFF  # range in where to look for arguments
# name of enum holding NULL values and description
NULL_ENUM_NAME = 'Null_Enum'

g_logger = logging.getLogger(__name__)


class FailedToExpandSegmentException(Exception):

    def __init__(self, message):
        super(FailedToExpandSegmentException, self).__init__(message)
        self.message = message


class ArgumentNotFoundException(Exception):

    def __init__(self, message):
        super(ArgumentNotFoundException, self).__init__(message)
        self.message = message


class FailedToAppendSegmentException(Exception):

    def __init__(self, message):
        super(FailedToAppendSegmentException, self).__init__(message)
        self.message = message


class RenamingException(Exception):

    def __init__(self, message):
        super(RenamingException, self).__init__(message)
        self.message = message


class NoInputFileException(Exception):

    def __init__(self, message):
        super(NoInputFileException, self).__init__(message)
        self.message = message


def make_import_names_callback(library_calls, library_addr):
    """ Return a callback function used by idaapi.enum_import_names(). """
    def callback(ea, name, ordinal):
        """ Callback function to retrieve code references to library calls. """
        library_calls[name] = []
        library_addr[name] = ea
        for ref in idautils.CodeRefsTo(ea, 0):
            library_calls[name].append(ref)
        return True  # True -> Continue enumeration
    return callback


def get_imports(library_calls, library_addr):
    """ Populate dictionaries with import information. """
    import_names_callback = make_import_names_callback(library_calls,
                                                       library_addr)
    for i in range(0, idaapi.get_import_module_qty()):
        idaapi.enum_import_names(i, import_names_callback)


def format_comment(comment_string, width=COMMENT_MAX_WIDTH):
    """ Return UTF encoded string limited to 'width' characters per line. """
    return str(fill(comment_string, width))


def add_fct_descr(ea, function, rep):
    """ Insert a (repeatable) comment describing the function at ea.

    Arguments:
    ea -- effective address where the comment is added
    function -- function object holding data
    rep -- add repeatable comment (True/False)
    """
    descr = format_comment(function.description) + '\n' + \
        format_comment('RETURN VALUE: ' + function.returns)
    # Both functions do not return
    if rep:
        idc.set_cmt(ea, descr, 1)
    else:
        idc.set_cmt(ea, descr, 0)


def get_end_of_last_segment():
    """ Return the last segment's end address. """
    last_ea = 0
    for segment in idautils.Segments():
        if idc.get_segm_end(segment) > last_ea:
            last_ea = idc.get_segm_end(segment)
    return last_ea


def expand_segment(ea):
    """ Expand last segment so it can hold more MSDN argument information.

    Argument:
    ea -- effective address within last segment
    """
    start = idc.get_segm_start(ea)
    end = idc.get_segm_end(ea)
    if end != get_end_of_last_segment():
        raise FailedToExpandSegmentException('Can only expand last segment.')
    if start == idaapi.BADADDR or end == idaapi.BADADDR:
        raise FailedToExpandSegmentException('Invalid start or end address.')
    new_end = end + NEW_SEGMENT_SIZE / 2
    if not idc.set_segment_bounds(ea, start, new_end, idaapi.SEGMOD_KEEP):
        raise FailedToExpandSegmentException('Setting segment bounds failed.')


def get_segment_end_ea(ea):
    """ Return address where next MSDN info can be written to in added
    segment.

    Argument:
    ea -- effective address within added segment where search starts
    """
    addr = ea
    while idc.get_cmt(addr, 0) is not None:
        addr = addr + 1
    if addr > idc.get_segm_end(ea):
        g_logger.debug('Address {} out of segment bounds. Expanding segment.'
                       .format(hex(addr)))
        try:
            expand_segment(ea)
        except FailedToExpandSegmentException as e:
            g_logger.warning(e.message)
            raise e
    else:
        return addr


def append_segment(segment_name):
    """ Add a new segment to the IDB file and return its starting address.
    Information about function arguments will be stored here. Only works if the
    segment name is not used yet. This does not affect the original binary.

    Arguments:
    segment_name -- the name of the segment to be added
    """
    for segment in idautils.Segments():
        if idc.get_segm_name(segment) == segment_name:
            g_logger.warning('Segment ' + segment_name + ' already exists')
            return idc.get_segm_start(segment)

    new_segment_start = get_end_of_last_segment()
    g_logger.debug('Adding new segment at 0x%08x' % new_segment_start)
    if not idc.AddSeg(new_segment_start, (new_segment_start+NEW_SEGMENT_SIZE),
                      0, 1, 0, idaapi.scPub) == 1:
        raise FailedToAppendSegmentException('Could not add segment')
    # set new segment's attributes
    if not idc.set_segm_name(new_segment_start, segment_name):
        raise FailedToAppendSegmentException('Could not rename segment')
    if not idc.set_segm_class(new_segment_start, 'DATA'):
        raise FailedToAppendSegmentException('Could not set segment class')
    if not idc.set_segm_alignment(new_segment_start, idc.saRelPara):
        raise FailedToAppendSegmentException('Could not align segment')
    if not idc.set_segm_addressing(new_segment_start, 1):  # 1 -- 32 bit
        raise FailedToAppendSegmentException(
            'Could not set segment addressing')
    return new_segment_start


def name_exists(name):
    """ Return 'True' if name exists in current IDB file. """
    for _, existing_names in idautils.Names():  # generates (addr, name) tuples
        if name in existing_names:
            return True
    return False


def add_arg_descr(function, segment_ea, arg_description_format):
    """ Name address in added segment annotated with argument descriptions.

    Arguments:
    function -- function object
    segment_ea -- start looking for empty byte to annotate from this ea

    Return:
    next possible free address to add information to
    """
    # No arguments
    if len(function.arguments) == 0:
        return segment_ea
    for argument in function.arguments:
        try:
            free_ea = get_segment_end_ea(segment_ea)
        except FailedToExpandSegmentException as e:
            raise e

        fields = {
            "function_name": function.name,
            "function_dll":  function.dll,
            "argument_name": argument.name,
        }
        name = arg_description_format.format(**fields)
        if not name_exists(name):
            g_logger.debug(' Adding name {} at {}'.format(name, hex(free_ea)))
            idaapi.set_name(free_ea, name)
            description = argument.description[:MAX_ARG_DESCR_LEN]
            idc.set_cmt(free_ea, format_comment(description), 0)
        else:
            g_logger.debug(' Name %s already exists' % name)
    return (free_ea + 1)


def find_arg_ea(ea_call, arg_name):
    """ Return ea of argument by looking backwards from library function
    call.

    Arguments:
    ea_call -- effective address of call
    arg_name -- the argument name to look for
    """
    # the search for previous instruction/data will stop at the specified
    # address (inclusive)
    prev_instr = idc.prev_head(ea_call, ea_call - PREVIOUS_INSTR_DELTA)
    while prev_instr > (ea_call - ARG_SEARCH_THRESHOLD) and \
            prev_instr != idaapi.BADADDR:
        # False indicates not to look for repeatable comments
        comment = idc.get_cmt(prev_instr, False)
        if comment == arg_name:
            return prev_instr
        prev_instr = idc.prev_head(
            prev_instr, prev_instr - PREVIOUS_INSTR_DELTA)
    raise ArgumentNotFoundException('  Argument {} not found within threshold'
                                    .format(arg_name))


def add_enums(function):
    """ Add standard enums from parsed MSDN documentation for all imported
    library calls and their arguments.

    Arguments:
    function -- function object
    """
    enum_count = 0
    for argument in function.arguments:
        # Add standard enums
        if not argument.enums:
            g_logger.debug(' No standard constants available for %s' %
                           argument.name)
        else:
            for enum in argument.enums:
                g_logger.debug('  Importing enum %s for argument %s' %
                               (enum, argument.name))
                if idc.import_type(-1, enum) != idaapi.BADADDR:
                    g_logger.debug('  ' + enum + ' ' +
                                   hex(idc.get_enum(enum)) +
                                   ' added successfully')
                    enum_count = enum_count + 1
                else:
                    g_logger.debug('  Could not add ' + enum)

        if not argument.constants:
            # No constants for this argument
            continue

        argument.name = argument.name
        function.name = function.name

        # Add constant descriptions
        for constant in argument.constants:
            if constant.name == 'NULL':
                # Create unique name, so we can add descriptive comment to it
                constant.name = 'NULL_{}_{}'.format(argument.name,
                                                    function.name)
                # Add custom enum for NULL values if it does not exist yet
                enumid = idc.get_enum(NULL_ENUM_NAME)
                if enumid == idaapi.BADADDR:
                    enumid = idc.add_enum(-1, NULL_ENUM_NAME,
                                          idaapi.hex_flag())
                idc.add_enum_member(enumid, constant.name, 0, -1)
                constid = idc.get_enum_member_by_name(constant.name)
                idc.set_enum_member_cmt(constid, format_comment(constant.description),
                                        False)
            else:
                constid = idc.get_enum_member_by_name(constant.name)
                if constid:
                    if idc.set_enum_member_cmt(constid,
                                               format_comment(
                                                   constant.description),
                                               False):
                        g_logger.debug('    Description added for %s' %
                                       constant.name)
                    else:
                        g_logger.debug('    No description added for %s' %
                                       constant.name)
    return enum_count


def get_bitmasks(enumid):
    """ Return list of bitmasks used in enum. """
    bmasks = []
    bid = idc.get_first_bmask(enumid)
    while bid != idaapi.BADADDR:
        bmasks.append(bid)
        bid = idc.get_next_bmask(enumid, bid)
    return bmasks


def get_constant_id(enumid, value):
    """ Return id of constant for specific value in enum. """
    constid = idc.get_enum_member(enumid, value, 0, -1)
    if constid != idaapi.BADADDR and not idc.is_bf(enumid):
        return constid

    for bm in get_bitmasks(enumid):
        constid = idc.get_enum_member(enumid, value, 0, bm)
        if constid != idaapi.BADADDR:
            return constid
    return idaapi.BADADDR


def rename_constant(arg_ea, fct_name, arg_name, arg_enums):
    """ Rename constants to values from standard enumerations. """
    instruction = idc.print_insn_mnem(arg_ea)
    if instruction == 'push':
        op_num = 0
    elif instruction == 'mov':
        op_num = 1
    else:
        raise RenamingException('Constant: unhandled instruction ' +
                                instruction)

    op_val = idc.get_operand_value(arg_ea, op_num)
    # NULL
    if op_val == 0:
        targetid = idc.get_enum_member_by_name(
            'NULL_{}_{}'.format(arg_name, fct_name))
        serial = 0
        enumid = idc.get_enum(NULL_ENUM_NAME)
        constid = idc.get_enum_member(enumid, 0, serial, -1)
        while constid != idaapi.BADADDR:
            if constid == targetid:
                idc.op_enum(arg_ea, op_num, enumid, serial)
                return
            serial = serial + 1
            constid = idc.get_enum_member(enumid, 0, serial, -1)

    # All other constants
    op_type = idc.get_operand_type(arg_ea, op_num)
    if op_type == idaapi.o_imm:
        # only one choice
        if len(arg_enums) == 1:
            enumid = idc.get_enum(arg_enums[0])
            idc.op_enum(arg_ea, op_num, enumid, 0)
            return

        for enum in arg_enums:
            enumid = idc.get_enum(enum)
            constid = get_constant_id(enumid, op_val)
            if constid == idaapi.BADADDR:
                # Not in this enum
                continue
            else:
                # Found the right enum
                idc.op_enum(arg_ea, op_num, enumid, 0)
                return


def rename_argument(ea, function, argument, arg_description_format):
    """ Rename function's argument comment at ea based on config string. """
    fields = {
        "function_name": function.name,
        "function_dll":  function.dll,
        "argument_name": argument.name,
    }
    new_arg = str(arg_description_format.format(**fields).encode('utf-8'))
    idc.set_cmt(ea, new_arg, 0)


def rename_args_and_consts(ref, function, conf_constants_import,
                           conf_arguments_annotate,
                           conf_arg_description_format):
    """ Rename arguments and constants for a function called at 'ref'. """
    for argument in function.arguments:
        try:
            arg_ea = find_arg_ea(ref, argument.name)
        except ArgumentNotFoundException as e:
            g_logger.debug(e.message)
            continue
        if conf_constants_import and argument.enums != []:
            g_logger.debug('  renaming constant {} ({})'.format(argument.name,
                                                                hex(arg_ea)))
            try:
                rename_constant(arg_ea, function.name, argument.name,
                                argument.enums)
            except RenamingException as e:
                g_logger.warning(e)

        if conf_arguments_annotate:
            g_logger.debug('  renaming argument {} ({})'.format(argument.name,
                                                                hex(arg_ea)))
            rename_argument(arg_ea, function, argument,
                            conf_arg_description_format)


def backup_database():
    """ Backup the database to a file similar to IDA's snapshot function. """
    time_string = strftime('%Y%m%d%H%M%S')
    file = idc.get_root_filename()
    if not file:
        raise NoInputFileException('No input file provided')
    input_file = file.rsplit('.', 1)[0]
    backup_file = '%s_%s.idb' % (input_file, time_string)
    g_logger.info('Backing up database to file ' + backup_file)
    idc.save_database(backup_file, idaapi.DBFL_BAK)


def get_data_files(dir):
    """ Return alphabetical sorted list of all found XML data files in
    directory, excluding the main database file (MSDN_INFO_FILE).

    Argument:
    dir -- path where XML data files reside
    """
    data_files = os.listdir(dir)
    if MSDN_INFO_FILE in data_files:
        data_files.remove(MSDN_INFO_FILE)
        return data_files
    else:
        raise IOError('Main database file ' + MSDN_INFO_FILE + ' not found' +
                      ' in ' + dir)


def parse_xml_data_files(msdn_data_dir):
    """ Return dictionary holding function information.

    Arguments:
    msdn_data_dir -- path to the directory storing the XML data files
    """
    functions_map = {}

    # Parse main database file first
    xml_file = os.path.join(msdn_data_dir, MSDN_INFO_FILE)
    functions = xml_parser.parse(xml_file)
    for function in functions:
        functions_map[function.name] = function

    # Parse additional files
    data_files = get_data_files(msdn_data_dir)
    for file in data_files:
        xml_file = os.path.join(msdn_data_dir, file)
        additional_functions = xml_parser.parse(xml_file)

        # Merge functions or add new function
        for a_function in additional_functions.iter():
            if a_function.name in functions_map:
                functions_map[a_function.name].merge(a_function)
            else:
                functions_map[a_function.name] = a_function
    return functions_map


def main(config=None):
    g_logger.info('Starting script execution')

    if CREATE_BACKUP:
        # Backup database before making any changes
        try:
            backup_database()
        except NoInputFileException as e:
            g_logger.warn('Quitting execution: ' + e.message)
            return

    # Default config in case none is provided
    config['arg_description_format'] = '{argument_name}_{function_name}'
    if not config:
        config = {}
        config['functions_annotate'] = True
        config['functions_repeatable_comment'] = False
        config['arguments_annotate'] = True
        config['constants_import'] = True
        config['msdn_data_dir'] = os.path.abspath(
            os.path.join(idaapi.get_user_idadir(), 'MSDN_data'))

    # Parse XML files and populate dictionary
    msdn_data_dir = config['msdn_data_dir']
    if not os.path.exists(msdn_data_dir):
        g_logger.error(
            'Configured msdn data directory does not exist: %s', msdn_data_dir)
        return
    functions_map = parse_xml_data_files(msdn_data_dir)

    # Retrieve all imported functions, store data in dictionaries
    g_logger.debug('Retrieving imported functions')
    library_calls = {}  # maps function_name -> CodeRefTo
    library_addr = {}  # maps function_name -> ea in import table
    get_imports(library_calls, library_addr)

    # Append segment where function argument information will be stored
    try:
        g_logger.debug('Appending new segment %s' % NEW_SEGMENT_NAME)
        free_ea = append_segment(NEW_SEGMENT_NAME)
    except FailedToAppendSegmentException(Exception) as e:
        g_logger.debug(e.message)
        return

    g_logger.debug('Starting annotations')
    functions_not_found = []
    for fct_name, eas in library_calls.items():
        if fct_name not in functions_map:
            # sometimes function info is available, but the import is more
            # specific, e.g., lstrcat vs. lstrcatA/W
            if fct_name[:-1] in functions_map:
                library_addr[fct_name[:-1]] = library_addr[fct_name]
                fct_name = fct_name[:-1]
            elif fct_name[6:] in functions_map:
                # handle mangled names (starting with __imp_)
                library_addr[fct_name[len('__imp_'):]] = library_addr[fct_name]
                fct_name = fct_name[len('__imp_'):]
            else:
                functions_not_found.append(fct_name)
                continue

        g_logger.debug('Working on function %s' % fct_name)
        if config['functions_annotate']:
            # Add function description to import table
            add_fct_descr(library_addr[fct_name],
                          functions_map[fct_name],
                          config['functions_repeatable_comment'])

        if config['constants_import']:
            # Add enums for extracted constant data
            num_added_enums = add_enums(functions_map[fct_name])
            if num_added_enums:
                g_logger.debug(' Added {} ENUMs for {}'.format(num_added_enums,
                                                               fct_name))

        # Add argument description in newly created segment
        free_ea = add_arg_descr(functions_map[fct_name],
                                free_ea, config['arg_description_format'])

        # Rename arguments and constants so they link to set names
        for ref in eas:
            g_logger.debug(' Enhancing argument and constant info for {} ({})'
                           .format(fct_name, hex(ref)))
            rename_args_and_consts(ref, functions_map[fct_name],
                                   config['constants_import'],
                                   config['arguments_annotate'],
                                   config['arg_description_format'])

    # Report
    logging.info('\n======================')
    logging.info('MSDN Annotator SUMMARY')
    logging.info('======================')
    logging.info(' Functions not found')
    logging.info(' -------------------')
    i = 1
    for f in functions_not_found:
        logging.info('  {}\t{}'.format(i, f))
        i += 1
    logging.info('')


if __name__ == '__main__':
    main()
