# Copyright (C) 2020 Zeropoint Dynamics

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see
# <http://www.gnu.org/licenses/>.
# ======================================================================

import ida_kernwin
import idaapi
import idc


class ApplyZelosOverlay(idaapi.action_handler_t):  # pragma: no cover
    """
    Applies the overlay from the user-selected memory snapshot.
    """

    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        import json

        filepath = ida_kernwin.ask_file(
            False, "*.zmu;*.overlay;*", "Load Zelos Overlay..."
        )
        if filepath is None:
            return
        f = open(filepath, "r")
        zelos_data = f.read()
        f.close()

        zelos_data = zelos_data[len("DISAS\n") :]
        zelos_dump = json.loads(zelos_data)

        # Apply the overlay data
        for comment in zelos_dump["comments"]:
            ea = comment["address"]
            try:
                comment_text = str(comment["text"])
            except UnicodeEncodeError:
                comment_text = ""
            color = comment.get("color", 0x73F0DF)

            # Set color of instruction line
            idaapi.set_item_color(ea, color)
            idaapi.set_cmt(ea, comment_text, False)

            # Set function name if not already changed
            idc.get_func_attr(ea, idc.FUNCATTR_START)
            name = idc.get_func_name(ea)
            if len(name) > 0 and name.startswith("zmu_") is False:
                idc.set_name(ea, "zmu_" + name)

        return 1

    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


class zelosoverlay_t(idaapi.plugin_t):  # pragma: no cover
    """
    Adds a Zelos {View} menu option for loading an overlay.
    """

    flags = 0
    comment = "Load an overlay file generated by Zelos."
    help = comment
    wanted_name = "ZelosOverlay"
    wanted_hotkey = ""
    menu_name = "View/"
    menu_context = []

    def init(self):
        zelos_overlay_action = idaapi.action_desc_t(
            "zelosoverlay:action",
            "Load Zelos Overlay...",
            ApplyZelosOverlay(),
            "",
            "Load Zelos Overlay...",
            199,
        )

        idaapi.register_action(zelos_overlay_action)
        idaapi.attach_action_to_menu(
            "View/", "zelosoverlay:action", idaapi.SETMENU_APP
        )

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        pass


def PLUGIN_ENTRY():  # pragma: no cover
    return zelosoverlay_t()
