# Copyright 2020 Google LLC
#
# Licensed 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.

"""
IDA Pro script that imports a capa report,
produced via `capa --json /path/to/sample`,
into the current database.

It will mark up functions with their capa matches, like:

    ; capa: print debug messages (host-interaction/log/debug/write-event)
    ; capa: delete service (host-interaction/service/delete)
    ; Attributes: bp-based frame

    public UninstallService
    UninstallService proc near
    ...

To use, invoke from the IDA Pro scripting dialog,
such as via Alt-F9,
and then select the existing capa report from the file system.

This script will verify that the report matches the workspace.
Check the output window for any errors, and/or the summary of changes.
"""

import logging
from pathlib import Path

import ida_nalt
import ida_funcs
import ida_kernwin

import capa.rules
import capa.features.freeze
import capa.render.result_document

logger = logging.getLogger("capa")


def append_func_cmt(va, cmt, repeatable=False):
    """
    add the given comment to the given function,
    if it doesn't already exist.
    """
    func = ida_funcs.get_func(va)
    if not func:
        raise ValueError("not a function")

    existing = ida_funcs.get_func_cmt(func, repeatable) or ""
    if cmt in existing:
        return

    if len(existing) > 0:
        new = existing + "\n" + cmt
    else:
        new = cmt

    ida_funcs.set_func_cmt(func, new, repeatable)


def main():
    path = ida_kernwin.ask_file(False, "*", "capa report")
    if not path:
        return 0

    result_doc = capa.render.result_document.ResultDocument.from_file(Path(path))
    meta, capabilities = result_doc.to_capa()

    # in IDA 7.4, the MD5 hash may be truncated, for example:
    # wanted: 84882c9d43e23d63b82004fae74ebb61
    # found: b'84882C9D43E23D63B82004FAE74EBB6\x00'
    #
    # see: https://github.com/idapython/bin/issues/11
    a = meta.sample.md5.lower()
    b = bytes.hex(ida_nalt.retrieve_input_file_md5()).lower()
    if not a.startswith(b):
        logger.error("sample mismatch")
        return -2

    rows = []
    for name in capabilities.matches.keys():
        rule = result_doc.rules[name]
        if rule.meta.lib:
            continue
        if rule.meta.is_subscope_rule:
            continue
        if rule.meta.scopes.static == capa.rules.Scope.FUNCTION:
            continue

        ns = rule.meta.namespace

        for address, _ in rule.matches:
            if address.type != capa.features.freeze.AddressType.ABSOLUTE:
                continue

            va = address.value
            rows.append((ns, name, va))

    # order by (namespace, name) so that like things show up together
    rows = sorted(rows)
    for ns, name, va in rows:
        if ns:
            cmt = name + f"({ns})"
        else:
            cmt = name

        logger.info("0x%x: %s", va, cmt)
        try:
            # message will look something like:
            #
            #     capa: delete service (host-interaction/service/delete)
            append_func_cmt(va, "capa: " + cmt, repeatable=False)
        except ValueError:
            continue

    logger.info("ok")


main()