# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
# 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: [package root]/LICENSE.txt
# 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 logging
import datetime

import idc
import idaapi
import idautils

import capa
import capa.version
import capa.features.common

logger = logging.getLogger("capa")

SUPPORTED_IDA_VERSIONS = [
    "7.1",
    "7.2",
    "7.3",
    "7.4",
    "7.5",
    "7.6",
]

# file type names as returned by idaapi.get_file_type_name()
SUPPORTED_FILE_TYPES = [
    "Portable executable for 80386 (PE)",
    "Portable executable for AMD64 (PE)",
    "Binary file",  # x86/AMD64 shellcode support
]


def inform_user_ida_ui(message):
    idaapi.info("%s. Please refer to IDA Output window for more information." % message)


def is_supported_ida_version():
    version = idaapi.get_kernel_version()
    if version not in SUPPORTED_IDA_VERSIONS:
        warning_msg = "This plugin does not support your IDA Pro version"
        logger.warning(warning_msg)
        logger.warning(
            "Your IDA Pro version is: %s. Supported versions are: %s." % (version, ", ".join(SUPPORTED_IDA_VERSIONS))
        )
        return False
    return True


def is_supported_file_type():
    file_type = idaapi.get_file_type_name()
    if file_type not in SUPPORTED_FILE_TYPES:
        logger.error("-" * 80)
        logger.error(" Input file does not appear to be a PE file.")
        logger.error(" ")
        logger.error(
            " capa currently only supports analyzing PE files (or binary files containing x86/AMD64 shellcode) with IDA."
        )
        logger.error(" If you don't know the input file type, you can try using the `file` utility to guess it.")
        logger.error("-" * 80)
        return False
    return True


def get_disasm_line(va):
    """ """
    return idc.generate_disasm_line(va, idc.GENDSM_FORCE_CODE)


def is_func_start(ea):
    """check if function stat exists at virtual address"""
    f = idaapi.get_func(ea)
    return f and f.start_ea == ea


def get_func_start_ea(ea):
    """ """
    f = idaapi.get_func(ea)
    return f if f is None else f.start_ea


def get_file_md5():
    """ """
    md5 = idautils.GetInputFileMD5()
    if not isinstance(md5, str):
        md5 = capa.features.common.bytes_to_str(md5)
    return md5


def get_file_sha256():
    """ """
    sha256 = idaapi.retrieve_input_file_sha256()
    if not isinstance(sha256, str):
        sha256 = capa.features.common.bytes_to_str(sha256)
    return sha256


def collect_metadata():
    """ """
    md5 = get_file_md5()
    sha256 = get_file_sha256()

    return {
        "timestamp": datetime.datetime.now().isoformat(),
        # "argv" is not relevant here
        "sample": {
            "md5": md5,
            "sha1": "",  # not easily accessible
            "sha256": sha256,
            "path": idaapi.get_input_file_path(),
        },
        "analysis": {
            "format": idaapi.get_file_type_name(),
            "extractor": "ida",
            "base_address": idaapi.get_imagebase(),
        },
        "version": capa.version.__version__,
    }
