# plugins

repository: smuniz/point-source
file: debugger_plugins/pointsource_plugin.py
description: Python-based C Decompiler based on LLVM
wanted_name: Point source decompiler
comment: Point source C decompiler
created_at: 2014-10-03 13:17:13+00:00
pushed_at: 2017-03-15 14:10:43+00:00
url:
https://github.com/smuniz/point-source/blob/a35a25cdb810b7e86457cb74d0641df23bdfdac2/debugger_plugins/pointsource_plugin.py

content:
#!/usr/bin/env python
# 
# Copyright (c) 2017 Sebastian Muniz
# 
# This code is part of point source decompiler
#

from sys import path as sys_path, exit, stdout
from inspect import getfile, currentframe
from os import path
from traceback import format_exc

#
# Add path to system paths.
#
#current_filename = getfile(currentframe())
#print "-----> %s" % current_filename
#path_to_current_module = os.path.realpath(os.path.dirname(current_filename))
path_to_current_module = r"z:\home\oem\projects\decompiler\point-source.git"

if path_to_current_module not in sys_path:
    print "[+] Fixing Python modules path..."
    sys_path.append(path_to_current_module)

#print "-----> %s" % path_to_current_module

import idaapi

from pointsource import PointSource, PointSourceException

__description__ = "Point Source Decompiler plugin"
__version__ = "0.1"


class PointSourcePlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Point source C decompiler"

    help = "Point source C decompiler help"
    wanted_name = "Point source decompiler"
    wanted_hotkey = ""

    DECOMPILE_HOTKEY = "Alt-F5"

    def init(self):
        """Initialize the plugin."""
        # Set debugging flag during development
        self.debug = True

        self.print_banner()
        try:
            #
            # Initializing decompiler in order to know if the current
            # architecture is supported or not.
            #
            self.decompiler = PointSource()

            #
            # Registering shortcuts.
            #
            new_menu = idaapi.add_menu_item("Edit/Plugins", "-", None, 0, self.do_nothing, ())
            if not new_menu:
                #print "[-] Unable to add menu separator."
                del new_menu
            new_menu = idaapi.add_menu_item("Edit/Plugins", "PointSource : Decompile function", self.DECOMPILE_HOTKEY, 0, self.decompile_function, ())
            if not new_menu:
                #print "[-] Unable to add menu item."
                del new_menu
            
            new_hotkey = idaapi.add_hotkey(self.DECOMPILE_HOTKEY, self.decompile_function)
            if not new_hotkey:
                #print "[-] Unable to add hotkey."
                del new_hotkey

            # Remain in memory because this is a supported architecture and
            # we'll get some work to do.
            return idaapi.PLUGIN_KEEP

        except PointSourceException, err:
            print "[-] Unable to initialize decompiler : %s" % err

        # This architecture is unsupported so we unload ourselves and leave.
        return idaapi.PLUGIN_UNL

    def run(self, argument):
        """Execute the script when invoked."""
        try:
            # Remove the modal dialogue 
            old = idaapi.disable_script_timeout()

            idaapi.show_wait_box("Decompiling function...")

            self.decompiler.set_screen_address_to_decompile()
            self.decompiler.decompile()

            # Re-enable the original timeout.
            idaapi.set_script_timeout(old)

        except PointSourceException, err:
            print "[-] Unable to run decompiler : %s" % err

        except Exception, err:
            if self.debug:
                print format_exc()
            print "[-] Exception : %s" % err

    def term(self):
        """Terminate the plugin."""
        print "Terminating plugin : %s" % __description__

    def print_banner(self):
        banner = [
          "%s has been loaded (v%s)" % (__description__, __version__),
          "  The hotkeys are : %s -> decompile" % self.DECOMPILE_HOTKEY,
          "  Additional options available in the Edit/Plugins menu.",
          "  contact Point Source Team on Slack at: point-source.slack.com",
        ]
        sepline = '-' * (max([len(s) for s in banner])+1)

        print(sepline)
        print("\n".join(banner))
        print(sepline)

    def do_nothing(self):
        """Empty stub for blank menu line."""
        pass

    def decompile_function(self):
        """Decompile the function under the cursor."""
        #print "Decompile function at 0x%X" % idaapi.get_screen_ea()
        self.run(0)

def PLUGIN_ENTRY():
    return PointSourcePlugin()


repository: baidubce/bce-qianfan-sdk
file: python/qianfan/common/cli/plugin.py
created_at: 2023-10-19 04:07:27+00:00
pushed_at: 2025-04-18 09:51:11+00:00
description:
Provide best practices for LMOps, as well as elegant and convenient access to the features of the Qianfan MaaS Platform. (提供大模型工具链最佳实践，以及优雅且便捷地访问千帆大模型平台）

url:
https://github.com/baidubce/bce-qianfan-sdk/blob/3eb1bcec265f88b40b22a2517c470707e82ab652/python/qianfan/common/cli/plugin.py

content:
# Copyright (c) 2023 Baidu, 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
#
#     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 json
import os
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import unquote, urlparse

import typer
from prompt_toolkit import prompt
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.document import Document
from prompt_toolkit.validation import ValidationError
from rich import print as rprint
from rich.console import Group, RenderableType
from rich.live import Live
from rich.markdown import Markdown
from rich.spinner import Spinner
from rich.text import Text

import qianfan
from qianfan import QfRole
from qianfan.common.cli.utils import (
    BosPathValidator,
    InputEmptyValidator,
    credential_required,
    render_response_debug_info,
    replace_logger_handler,
)
from qianfan.resources.typing import QfMessages
from qianfan.utils.bos_uploader import BosHelper, parse_bos_path


class PluginInputValidator(InputEmptyValidator):
    """
    Validator for input in plugin
    """

    def validate(self, document: Document) -> None:
        """
        validate input:
        - input must not be empty
        - if input is /image, file path must be provided
        """
        super().validate(document)
        text = document.text.strip()
        if text.startswith(PluginClient.IMAGE_PROMPT):
            path = text[len(PluginClient.IMAGE_PROMPT) :].strip()
            if len(path) == 0:
                raise ValidationError(
                    message="Image file path must be provided (e.g. /image car.jpg)."
                )


class PluginClient(object):
    """
    Client object for the chat command
    """

    END_PROMPT = "/exit"
    RESET_PROMPT = "/reset"
    IMAGE_PROMPT = "/image"
    FILE_PROMPT = "/file"
    HELP_PROMPT = "/help"

    HELP_MESSAGES = {
        END_PROMPT: "End the conversation",
        RESET_PROMPT: "Reset the conversation",
        IMAGE_PROMPT: "Attach a local image to the conversation (e.g. /image car.jpg)",
        FILE_PROMPT: "Attach a local file to the conversation (e.g. /file xx.docx)",
        HELP_PROMPT: "Print help message",
    }

    def __init__(
        self,
        model: str,
        endpoint: Optional[str],
        multi_line: bool,
        debug: bool,
        plugins: List[str],
        bos_path: Optional[str],
        **kwargs: Any,
    ) -> None:
        """
        Init the chat client
        """
        if endpoint is None:
            self.client = qianfan.Plugin(model=model)
        else:
            self.client = qianfan.Plugin(endpoint=endpoint)
        self.msg_history = QfMessages()
        self.multi_line = multi_line
        self.console = replace_logger_handler()
        self.inference_args = kwargs
        self.bos_path = bos_path
        self.plugins = plugins
        self.debug = debug
        if self._is_qianfan_plugin():
            self.HELP_MESSAGES.pop(self.FILE_PROMPT)
        self.input_completer = WordCompleter(
            list(self.HELP_MESSAGES.keys()), sentence=True, meta_dict=self.HELP_MESSAGES
        )

    def print_hint_msg(self) -> None:
        """
        Print hint message when startup
        """
        if self.multi_line:
            rprint(
                "[bold]Hint[/bold]: [green bold]Press Esc before Enter[/] to submit"
                f" your message, or use '{self.HELP_PROMPT}' to acquire more commands."
            )
        else:
            rprint(
                "[bold]Hint[/bold]: Press enter to submit your message, or use"
                f" '{self.HELP_PROMPT}' to acquire more commands.."
            )
            rprint(
                "[bold]Hint[/bold]: If you want to submit multiple lines, use the"
                " '--multi-line' option."
            )
        if self._is_qianfan_plugin():
            rprint(f"[dim]Using qianfan plugin with {self.plugins}...[/]")
        else:
            rprint("[dim]Using ERNIE Bot plugin...[/]")

    def get_bos_path(self) -> str:
        """
        Get bos path. If bos_path is not provided, prompt user to input
        """
        if self.bos_path is None:
            rprint("Please input bos bucket path [dim](bos:/<bucket>/<path>)[/]: ")
            self.bos_path = prompt("bos:/", validator=BosPathValidator())
            self.bos_path = "bos:/" + self.bos_path
        if not self.bos_path.endswith("/"):
            self.bos_path = self.bos_path + "/"
        return self.bos_path

    def upload_file_to_bos(self, filepath: str) -> Tuple[str, str]:
        """
        Upload file to bos and get bos_url and http_url
        """
        bos_helper = BosHelper()
        bucket, bos_path = parse_bos_path(self.get_bos_path())

        bos_path = bos_path + os.path.basename(filepath)
        with self.console.status("Uploading file to bos..."):
            bos_helper.upload_file_to_bos(filepath, bos_path, bucket)
        url = bos_helper.get_bos_file_shared_url(bos_path, bucket)
        bos_url = f"bos:/{bucket}{bos_path}"
        return bos_url, url

    def print_help_message(self) -> None:
        """
        Print command introduction
        """
        for k, v in self.HELP_MESSAGES.items():
            rprint(f"[bold green]{k}[/]: {v}")

    def _is_qianfan_plugin(self) -> bool:
        """
        Check if the plugin is qianfan plugin
        """
        return self.client._endpoint is not None

    def chat_in_terminal(self) -> None:
        """
        Chat in terminal
        """

        self.print_hint_msg()
        # loop the conversation
        extra_field: Dict[str, Any] = {}
        while True:
            rprint("\n[yellow bold]Enter your message[/yellow bold]:")
            while True:
                message = prompt(
                    multiline=self.multi_line,
                    validator=PluginInputValidator(),
                    completer=self.input_completer,
                ).strip()
                if message.startswith(self.IMAGE_PROMPT):
                    path = message[len(self.IMAGE_PROMPT) :].strip()
                    if not path.startswith("http"):
                        bos_path, http_path = self.upload_file_to_bos(path)
                        filename = os.path.basename(path)
                        rprint(f"Image has been uploaded to: {bos_path}")
                        rprint(f"Image share url: {http_path}\n")
                        path = http_path
                    else:
                        filename = unquote(os.path.basename(urlparse(path).path))
                    if self._is_qianfan_plugin():
                        rprint(
                            "[yellow bold]Please continue to input your prompt[/yellow"
                            " bold]:"
                        )
                        extra_field["fileurl"] = path
                        continue
                    else:
                        message = f"<img>{filename}</img><url>{path}</url>"
                if not self._is_qianfan_plugin() and message.startswith(
                    self.FILE_PROMPT
                ):
                    path = message[len(self.FILE_PROMPT) :].strip()
                    if not path.startswith("http"):
                        bos_path, http_path = self.upload_file_to_bos(path)
                        filename = os.path.basename(path)
                        rprint(f"File has been uploaded to: {bos_path}")
                        rprint(f"File share url: {http_path}\n")
                        path = http_path
                    else:
                        filename = unquote(os.path.basename(urlparse(path).path))
                    message = f"<file>{filename}</file><url>{path}</url>"

                break

            rprint("\n[blue][bold]Model response[/bold][/blue]:")

            if message == self.END_PROMPT:
                rprint("Bye!")
                raise typer.Exit()
            elif message == self.RESET_PROMPT:
                self.msg_history = QfMessages()
                extra_field = {}
                rprint("Chat history has been cleared.")
                continue
            elif message == self.HELP_PROMPT:
                self.print_help_message()
                continue
            else:
                self.msg_history.append(message)

            with Live(
                Spinner("dots", text="Thinking...", style="status.spinner"),
                auto_refresh=True,
                refresh_per_second=24,
                console=self.console,
            ) as live:
                if self._is_qianfan_plugin():
                    response = self.client.do(
                        message,
                        plugins=self.plugins,
                        llm=self.inference_args,
                        stream=True,
                        history=self.msg_history._to_list()[:-1],
                        **self.inference_args,
                        **extra_field,
                    )
                else:
                    response = self.client.do(
                        self.msg_history,
                        plugins=self.plugins,
                        stream=True,
                        **self.inference_args,
                    )

                m = ""
                plugin_metainfo = {}
                plugin_render_content = None
                for r in response:
                    render_list: List[RenderableType] = []
                    if self._is_qianfan_plugin():
                        m += r["result"]
                        render_list.append(Markdown(m))
                    else:
                        event = r["_event"]
                        if event == "pluginMeta":
                            for meta in r["plugin_metas"]:
                                plugin_metainfo[meta["pluginId"]] = meta
                        elif event == "plugin":
                            plugin_render_list: List[RenderableType] = []
                            for plugin in r["plugin_info"]:
                                plugin_name = plugin["plugin_id"]
                                if plugin_name in plugin_metainfo:
                                    plugin_name = (
                                        f"{plugin_metainfo[plugin_name]['pluginNameForHuman']}"
                                        f"({plugin_metainfo[plugin_name]['pluginNameForModel']})"
                                    )

                                plugin_resp = plugin["plugin_resp"].strip().split("\n")
                                status = None
                                for resp in reversed(plugin_resp):
                                    if len(resp) == 0:
                                        continue
                                    plugin_response = json.loads(resp)
                                    if "actionName" not in plugin_response:
                                        continue
                                    status = plugin_response["actionContent"]
                                    break
                                if status is not None:
                                    if plugin["status"] == "1":
                                        plugin_info_obj = Text.from_markup(
                                            text=f"• {plugin_name}\t{status}",
                                            style="dim",
                                        )
                                    else:
                                        plugin_info_obj = Text(
                                            f"• {plugin_name}", style="dim"
                                        )
                                    plugin_render_list.append(plugin_info_obj)
                            plugin_render_list.append(Text(""))
                            plugin_render_content = Group(*plugin_render_list)

                        elif event == "chat":
                            m += r["result"]
                        if plugin_render_content is not None:
                            render_list.append(plugin_render_content)
                        render_list.append(Markdown(m))

                    if "is_end" not in r or not r["is_end"]:
                        render_list.append(
                            Spinner(
                                "dots", text="Generating...", style="status.spinner"
                            )
                        )
                    stat = r.statistic
                    render_list.append(
                        Text.from_markup(
                            "\n[dim]First token latentcy:"
                            f" {stat['first_token_latency']:.2f}s, Total latency:"
                            f" {stat['total_latency']:.2f}s.[/]"
                        )
                    )
                    if "is_end" in r and r["is_end"]:
                        if "usage" in r:
                            token_usage = r["usage"]
                            render_list.append(
                                Text.from_markup(
                                    f"[dim]Input token: {token_usage['prompt_tokens']},"
                                    " Output token:"
                                    f" {token_usage['completion_tokens']}, Total token:"
                                    f" {token_usage['total_tokens']}.[/]"
                                )
                            )
                        if self.debug:
                            render_list.append(render_response_debug_info(response=r))

                    live.update(Group(*render_list))

            self.msg_history.append(m, role=QfRole.Assistant)


MODEL_ARGUMENTS_PANEL = (
    "Model Arguments (Some arguments are not supported by every model)"
)


@credential_required
def plugin_entry(
    endpoint: Optional[str] = typer.Option(
        None,
        help=(
            "Endpoint of the plugin. ERNIE-4.0-8K plugin will be used if not specified."
        ),
    ),
    # tui: bool = typer.Option(False, help="Using Terminal UI"),
    multi_line: bool = typer.Option(
        False,
        "--multi-line",
        help="Multi-line mode which needs to press Esc before enter to submit message.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        help="Enable debug mode. The request infomation will be printed.",
    ),
    plugins: str = typer.Option(
        ...,
        help=(
            "Plugins enabled. Use comma(,) to split. (e.g."
            " uuid-zhishiku,uuid-chatocr,uuid-weatherforecast)"
        ),
    ),
    bos_path: Optional[str] = typer.Option(None, help="Bos path used for upload file."),
    temperature: Optional[float] = typer.Option(
        None,
        help=(
            "Controls the randomness of the generated text. A higher temperature makes"
            " the model more creative and produces more diverse, but potentially less"
            " coherent."
        ),
        rich_help_panel=MODEL_ARGUMENTS_PANEL,
    ),
    top_p: Optional[float] = typer.Option(
        None,
        help=(
            "Lower top_p value allows the model to focus on a narrowed set of likely"
            " next tokens, making the response more conherent but less random."
        ),
        rich_help_panel=MODEL_ARGUMENTS_PANEL,
    ),
    penalty_score: Optional[float] = typer.Option(
        None,
        help="Penalty scores can be applied to discourage repetition.",
        rich_help_panel=MODEL_ARGUMENTS_PANEL,
    ),
    system: Optional[str] = typer.Option(
        None,
        help="Persona setting for the model.",
        rich_help_panel=MODEL_ARGUMENTS_PANEL,
    ),
    stop: Optional[str] = typer.Option(
        None,
        help="Stop words. Use comma to split multiple stop words.",
        rich_help_panel=MODEL_ARGUMENTS_PANEL,
    ),
) -> None:
    """
    Chat with the LLM with plugins in the terminal.
    """
    model = "EBPluginV2"

    extra_args = {}

    def add_if_not_none(key: str, value: Any) -> None:
        if value is not None:
            extra_args[key] = value

    add_if_not_none("temperature", temperature)
    add_if_not_none("top_p", top_p)
    add_if_not_none("penalty_score", penalty_score)
    add_if_not_none("system", system)

    if stop is not None:
        extra_args["stop"] = stop.split(",")

    client = PluginClient(
        model,
        endpoint,
        multi_line,
        debug=debug,
        plugins=plugins.split(","),
        bos_path=bos_path,
        **extra_args,
    )
    client.chat_in_terminal()


repository: lucasg/idarop
file: plugins/idarop_plugin_t.py
description: ROP database plugin for IDA
wanted_name: IDA ROP
created_at: 2017-08-08 20:09:38+00:00
pushed_at: 2018-06-05 19:52:47+00:00
url:
https://github.com/lucasg/idarop/blob/9b090f98c3bf61c66c59c5b75c57ad115b6a2366/plugins/idarop_plugin_t.py

content:
# IDA libraries
import idaapi
from idaapi import plugin_t

from idarop import IDAROP_VERSION, IDAROP_DESCRIPTION
from idarop.ui import IdaRopManager

class idarop_t(plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = IDAROP_DESCRIPTION
    help = IDAROP_DESCRIPTION
    wanted_name = "IDA ROP"
    wanted_hotkey = ""

    def init(self):
        """ On script initalisation : load previous rop results and init menu items """

        # Only Intel x86/x86-64 are supported
        if idaapi.ph_get_id() == idaapi.PLFM_386:

            global idarop_manager

            # Check if already initialized
            if not 'idarop_manager' in globals():

                idarop_manager = IdaRopManager()
                if idarop_manager.add_menu_items():
                    print("[IdaRop] Failed to initialize IDA Sploiter.")
                    idarop_manager.del_menu_items()
                    del idarop_manager
                    return idaapi.PLUGIN_SKIP
                else:
                    try:
                        idarop_manager.load_internal_db()
                    except Exception as e:
                        pass

                    print("[IdaRop] IDA ROP View  v%s initialized " % IDAROP_VERSION)

            return idaapi.PLUGIN_KEEP
        else:
            return idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass

    def term(self):
        """ On IDA's close event, export the Rop gadget list in a default csv file"""
        idaapi.show_wait_box("Saving gadgets ...")
        try:
            idarop_manager.save_internal_db()
        except Exception as e:
            pass
        idaapi.hide_wait_box()

def PLUGIN_ENTRY():
    return idarop_t()

###############################################################################
# Script / Testing
###############################################################################

def idarop_main():
    global idarop_manager

    if 'idarop_manager' in globals():
        idarop_manager.del_menu_items()
        del idarop_manager

    idarop_manager = IdaRopManager()
    idarop_manager.add_menu_items()

    idarop = idarop_manager.idarop

if __name__ == '__main__':
    #idarop_main()
    pass


repository: Asuri-Team/NUAACTF2022-Challenges
file: pwn/ezmmio/attachments/qiling-1.4.4/qiling/extensions/idaplugin/qilingida.py
description: NUAACTF 2022 challenges source code
wanted_name: Qiling Emulator
created_at: 2022-12-05 02:59:01+00:00
pushed_at: 2023-01-31 10:11:33+00:00
url:
https://github.com/Asuri-Team/NUAACTF2022-Challenges/blob/969b4028ef5f04498b2a49bd53ea52b820f3eb71/pwn/ezmmio/attachments/qiling-1.4.4/qiling/extensions/idaplugin/qilingida.py

content:
# Cross Platform and Multi Architecture Advanced Binary Emulation Framework Plugin For IDA
# Learn how to use? Please visit https://docs.qiling.io/en/latest/ida/

UseAsScript = True

import sys
import collections
import time
import struct
import re
import logging
import shlex
from enum import Enum
from elftools.elf.elffile import ELFFile
from json import load

# IDA Python SDK
from idaapi import *
from idc import *
from idautils import *
import idc
import ida_ua
import ida_idaapi
import ida_funcs
import ida_nalt
import idautils
import ida_xref
import ida_kernwin
import ida_ida
import ida_bytes
import ida_segment
import ida_name
import ida_gdl
import ida_frame
import ida_idp
import ida_auto
import ida_netnode
import ida_hexrays
import ida_range
# PyQt
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtWidgets import (QPushButton, QHBoxLayout)

# Qiling
from qiling import Qiling
from qiling.const import *
from qiling.arch.x86_const import reg_map_32 as x86_reg_map_32
from qiling.arch.x86_const import reg_map_64 as x86_reg_map_64
from qiling.arch.x86_const import reg_map_misc as x86_reg_map_misc
from qiling.arch.x86_const import reg_map_st as x86_reg_map_st
from qiling.arch.arm_const import reg_map as arm_reg_map
from qiling.arch.arm64_const import reg_map as arm64_reg_map
from qiling.arch.mips_const import reg_map as mips_reg_map
from qiling import __version__ as QLVERSION
from qiling.os.filestruct import ql_file
from keystone import *


QilingHomePage = 'https://www.qiling.io'
QilingStableVersionURL = 'https://raw.githubusercontent.com/qilingframework/qiling/master/qiling/__version__.py'
logging.basicConfig(level=logging.INFO, format='[%(levelname)s][%(module)s:%(lineno)d] %(message)s')

class Colors(Enum):
    Blue = 0xE8864A
    Pink = 0xC0C0FB
    White = 0xFFFFFF
    Black = 0x000000
    Green = 0xd3ead9
    Gray = 0xd9d9d9
    Beige = 0xCCF2FF

class IDA:
    def __init__(self):
        pass

    @staticmethod
    def get_function(addr):
        return ida_funcs.get_func(addr)

    @staticmethod
    def get_function_start(addr):
        return IDA.get_function(addr).start_ea

    @staticmethod
    def get_function_end(addr):
        return IDA.get_function(addr).end_ea

    @staticmethod
    def get_function_framesize(addr):
        return IDA.get_function(addr).frsize

    @staticmethod
    def get_function_name(addr):
        return ida_funcs.get_func_name(addr)

    @staticmethod
    def get_functions():
        return [IDA.get_function(func) for func in idautils.Functions()]

    @staticmethod
    def set_color(addr, what, color):
        return idc.set_color(addr, what, color)

    @staticmethod
    def color_block(bb, color):
        for i in range(bb.start_ea, bb.end_ea):
            IDA.set_color(i, idc.CIC_ITEM, color)

    # note:
    # corresponds to IDA graph view
    # a good example to iterate the graph
    # https://github.com/idapython/src/blob/bc9b51b1c70083815a574a57b7a783698de3698d/examples/core/dump_flowchart.py
    # arg can be a function or a (start, end) tuple or an address in the function
    @staticmethod
    def get_flowchart(arg):
        if type(arg) is int:
            func = IDA.get_function(arg)
            if func is None:
                return None
            return ida_gdl.FlowChart(func)
        return ida_gdl.FlowChart(arg)

    @staticmethod
    def get_block(addr):
        flowchart = IDA.get_flowchart(addr)
        for bb in flowchart:
            if bb.start_ea <= addr and addr < bb.end_ea:
                return bb
        return None

    @staticmethod
    def block_is_terminating(bb):
        # fcb_ret: has a retn instruction in the end
        # fcb_noret: in most cases, exit() is called
        # fcb_indjump: jmp $eax
        if (bb.type == ida_gdl.fcb_ret or bb.type == ida_gdl.fcb_noret or
                (bb.type == ida_gdl.fcb_indjump and len(list(bb.succs())) == 0)):
            return True
        for b in bb.succs():
            if b.type == ida_gdl.fcb_extern:
                return True
        return False

    @staticmethod
    def get_starting_block(addr):
        flowchart = IDA.get_flowchart(addr)
        if flowchart is None:
            return None
        func = IDA.get_function(addr)
        for bb in flowchart:
            if bb.start_ea == func.start_ea:
                return bb
        return None

    @staticmethod
    def get_terminating_blocks(addr):
        flowchart = IDA.get_flowchart(addr)
        return [bb for bb in flowchart if IDA.block_is_terminating(bb)]

    @staticmethod
    def get_prev_head(addr, minea=0):
        return ida_bytes.prev_head(addr, minea)

    @staticmethod
    def get_segments():
        r = []
        seg = ida_segment.get_first_seg()
        while seg is not None:
            r.append(seg)
            seg = ida_segment.get_next_seg(seg.start_ea)
        return r

    @staticmethod
    def get_segment_name(s):
        return ida_segment.get_segm_name(s)

    @staticmethod
    def get_segment_by_name(name):
        return ida_segment.get_segm_by_name(name)

    @staticmethod
    def __addr_in_seg(addr):
        segs = IDA.get_segments()
        for seg in segs:
            if addr < seg.end_ea and addr >= seg.start_ea:
                return seg
        return None

    # note: accept name and address in the segment
    @staticmethod
    def get_segment(arg):
        if type(arg) is int:
            return IDA.__addr_in_seg(arg)
        else: # str
            return IDA.get_segment_by_name(arg)

    @staticmethod
    def get_segment_start(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.start_ea
        return None

    @staticmethod
    def get_segment_end(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.end_ea
        return None

    @staticmethod
    def get_segment_perm(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.perm # RWX e.g. 0b101 = R + X
        return None

    @staticmethod
    def get_segment_type(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.type # 0x1 SEG_DATA 0x2 SEG_CODE See doc for details
        return None

    @staticmethod
    def get_instruction(addr):
        r = ida_ua.print_insn_mnem(addr)
        if r == "":
            return None
        return r

    # immidiate value
    @staticmethod
    def get_operand(addr, n):
        return (idc.get_operand_type(addr, n), idc.get_operand_value(addr, n))

    # eax, ecx, etc
    @staticmethod
    def print_operand(addr, n):
        return idc.print_operand(addr, n)

    @staticmethod
    def get_instruction_size(addr):
        return ida_bytes.get_item_size(addr)

    @staticmethod
    def get_instructions_count(begin, end):
        p = begin
        cnt = 0
        while p < end:
            sz = IDA.get_instruction_size(p)
            cnt += 1
            p += sz
        return cnt

    @staticmethod
    def get_name(addr):
        return ida_name.get_name(addr)

    @staticmethod
    def get_name_address(name, addr=0):
        return ida_name.get_name_ea(addr, name)

    @staticmethod
    def get_bytes(addr, l):
        return ida_bytes.get_bytes(addr, l)

    @staticmethod
    def get_byte(addr):
        return ida_bytes.get_byte(addr)

    @staticmethod
    def get_word(addr):
        return ida_bytes.get_word(addr)

    @staticmethod
    def get_dword(addr):
        return ida_bytes.get_dword(addr)

    @staticmethod
    def get_qword(addr):
        return ida_bytes.get_qword(addr)

    @staticmethod
    def get_xrefsto(addr, flags=ida_xref.XREF_ALL):
        return [ref.frm for ref in idautils.XrefsTo(addr, flags)]

    @staticmethod
    def get_xrefsfrom(addr, flags=ida_xref.XREF_ALL):
        return [ref.frm for ref in idautils.XrefsFrom(addr, flags)]

    @staticmethod
    def get_input_file_path():
        return ida_nalt.get_input_file_path()

    @staticmethod
    def get_info_structure():
        return ida_idaapi.get_inf_structure()

    @staticmethod
    def get_main_address():
        return IDA.get_info_structure().main

    @staticmethod
    def get_max_address():
        return IDA.get_info_structure().max_ea

    @staticmethod
    def get_min_address():
        return IDA.get_info_structure().min_ea

    @staticmethod
    def is_big_endian():
        return IDA.get_info_structure().is_be()

    @staticmethod
    def is_little_endian():
        return not IDA.is_big_endian()

    @staticmethod
    def get_filetype():
        info = IDA.get_info_structure()
        ftype = info.filetype
        if ftype == ida_ida.f_MACHO:
            return "macho"
        elif ftype == ida_ida.f_PE or ftype == ida_ida.f_EXE or ftype == ida_ida.f_EXE_old: # is this correct?
            return "pe"
        elif ftype == ida_ida.f_ELF:
            return "elf"
        else:
            return None

    @staticmethod
    def get_ql_arch_string():
        info = IDA.get_info_structure()
        proc = info.procname.lower()
        result = None
        if proc == "metapc":
            result = "x86"
            if info.is_64bit():
                result = "x8664"
        elif "mips" in proc:
            result = "mips"
        elif "arm" in proc:
            result = "arm32"
            if info.is_64bit():
                result = "arm64"
        # That's all we support :(
        return result

    @staticmethod
    def get_current_address():
        return ida_kernwin.get_screen_ea()

    # return (?, start, end)
    @staticmethod
    def get_last_selection():
        return ida_kernwin.read_range_selection(None)

    # Use with skipcalls
    # note that the address is the end of target instruction
    # e.g.:
    # 0x1 push eax
    # 0x4 mov eax, 0
    # call get_frame_sp_delta(0x4) and get -4.
    @staticmethod
    def get_frame_sp_delta(addr):
        return ida_frame.get_sp_delta(IDA.get_function(addr), addr)

    @staticmethod
    def patch_bytes(addr, bs):
        return ida_bytes.patch_bytes(addr, bs)

    @staticmethod
    def fill_bytes(start, end, bs = b'\x90'):
        return ida_bytes.patch_bytes(start, bs*(end-start))

    @staticmethod
    def nop_selection():
        _, start, end = IDA.get_last_selection()
        return IDA.fill_bytes(start, end)

    @staticmethod
    def fill_block(bb, bs=b'\x90'):
        return IDA.fill_bytes(bb.start_ea, bb.end_ea, bs)

    @staticmethod
    def assemble(ea, cs, ip, use32, line):
        return ida_idp.assemble(ea, cs, ip, use32, line)

    @staticmethod
    def create_data(ea, dataflag, size, tid=ida_netnode.BADNODE):
        return ida_bytes.create_data(ea, dataflag, size, tid)

    @staticmethod
    def create_bytes_array(start, end):
        return IDA.create_data(start, ida_bytes.byte_flag(), end-start)

    @staticmethod
    def create_byte(ea, length, force=False):
        return ida_bytes.create_byte(ea, length, force)

    @staticmethod
    def perform_analysis(start, end, final_pass=True):
        return ida_auto.plan_and_wait(start, end)

    @staticmethod
    def get_item_head(ea):
        return ida_bytes.get_item_head(ea)

    @staticmethod
    def get_item_end(ea):
        return ida_bytes.get_item_end(ea)

    @staticmethod
    def get_item_size(ea):
        return ida_bytes.get_item_size(ea)

    @staticmethod
    def get_item(ea):
        return (IDA.get_item_head(ea), IDA.get_item_end(ea))

    @staticmethod
    def is_colored_item(ea):
        return ida_nalt.is_colored_item(ea)

    # NOTE: The [start, end) range should include all control flows except long calls.
    @staticmethod
    def get_micro_code_mba(start, end, decomp_flags=ida_hexrays.DECOMP_WARNINGS, maturity=7):
        mbrgs = ida_hexrays.mba_ranges_t()
        rg = ida_range.range_t(start, end)
        mbrgs.ranges.push_back(rg)
        fl = ida_hexrays.hexrays_failure_t()
        mba = ida_hexrays.gen_microcode(mbrgs, fl, None, decomp_flags, maturity)
        if mba is None:
            logging.error(f"Fail to get mba because: {fl}")
        return mba

    @staticmethod
    def micro_code_from_mbb(mbb):
        cur = mbb.head
        while cur is not None:
            yield cur
            cur = cur.next
        return

### View Class

class QlEmuRegView(simplecustviewer_t):
    def __init__(self, ql_emu_plugin):
        super(QlEmuRegView, self).__init__()
        self.hooks = None
        self.ql_emu_plugin = ql_emu_plugin

    def Create(self):
        title = "QL Register View"
        if not simplecustviewer_t.Create(self, title):
            return False

        self.menu_update = 1

        class Hooks(UI_Hooks):
            class PopupActionHandler(action_handler_t):
                def __init__(self, subview, menu_id):
                    action_handler_t.__init__(self)
                    self.subview = subview
                    self.menu_id = menu_id

                def activate(self, ctx):
                    self.subview.OnPopupMenu(self.menu_id)

                def update(self, ctx):
                    return AST_ENABLE_ALWAYS

            def __init__(self, form):
                UI_Hooks.__init__(self)
                self.form = form

            def finish_populating_widget_popup(self, widget, popup):
                if self.form.title == get_widget_title(widget):
                    attach_dynamic_action_to_popup(widget, popup, action_desc_t(None, "Edit Register", self.PopupActionHandler(self.form, self.form.menu_update),  None, None, -1))

        if self.hooks is None:
            self.hooks = Hooks(self)
            self.hooks.hook()

        return True

    def SetReg(self, addr, ql:Qiling):
        #clear
        self.ClearLines()

        view_title = COLSTR("Reg value at { ", SCOLOR_AUTOCMT)
        view_title += COLSTR("IDA Address:0x%X | QL Address:0x%X" % (addr - self.ql_emu_plugin.qlemu.baseaddr + get_imagebase(), addr), SCOLOR_DREF)
        # TODO: Add disass should be better
        view_title += COLSTR(" }", SCOLOR_AUTOCMT)
        self.AddLine(view_title)
        self.AddLine("")

        reglist = QlEmuMisc.get_reg_map(ql)
        line = ""
        cols = 3
        reglist = [reglist[i:i+cols] for i in range(0,len(reglist),cols)]
        for regs in reglist:
            for reg in regs:
                line += COLSTR(" %4s: " % str(reg), SCOLOR_REG)
                regvalue = ql.arch.regs.read(reg)
                if arch in [QL_ARCH.X8664, QL_ARCH.ARM64]:
                    value_format = "0x%.16X"
                else:
                    value_format = "0x%.8X"
                line += COLSTR(str(value_format % regvalue), SCOLOR_NUMBER)
                # TODO: ljust will looks better
            self.AddLine(line)
            line = ''
        self.AddLine(line)
        self.Refresh()

    def OnPopupMenu(self, menu_id):
        if menu_id == self.menu_update:
            self.ql_emu_plugin.ql_chang_reg()

    def OnClose(self):
        if self.hooks:
            self.hooks.unhook()
            self.hooks = None
        self.ql_emu_plugin.ql_close_reg_view()

class QlEmuStackView(simplecustviewer_t):
    def __init__(self, ql_emu_plugin):
        super(QlEmuStackView, self).__init__()
        self.ql_emu_plugin = ql_emu_plugin

    def Create(self):
        title = "QL Stack View"
        if not simplecustviewer_t.Create(self, title):
            return False
        return True

    def SetStack(self, ql:Qiling):
        self.ClearLines()
        if ql is None:
            return

        sp = ql.arch.regs.arch_sp
        self.AddLine('')
        self.AddLine(COLSTR('  Stack at 0x%X' % sp, SCOLOR_AUTOCMT))
        self.AddLine('')

        reg_bit_size = ql.arch.bits
        reg_byte_size = reg_bit_size // 8
        value_format = '% .16X' if reg_bit_size == 64 else '% .8X'

        for i in range(-30, 30):
            clr = SCOLOR_DREF if i < 0 else SCOLOR_INSN
            cur_addr = (sp + i * reg_byte_size)
            line = ('  ' + value_format + ': ') % cur_addr
            try:
                value = ql.mem.read(cur_addr, reg_byte_size)
                value, = struct.unpack('Q' if reg_bit_size == 64 else 'I', value)
                line += value_format % value
            except Exception:
                line += '?' * reg_byte_size * 2

            self.AddLine(COLSTR(line, clr))

    def OnClose(self):
        self.ql_emu_plugin.ql_close_stack_view()

class QlEmuMemView(simplecustviewer_t):
    def __init__(self, ql_emu_plugin, addr, size):
        super(QlEmuMemView, self).__init__()
        self.ql_emu_plugin = ql_emu_plugin
        self.viewid = addr
        self.addr = addr
        self.size = size
        self.lastContent = []

    def Create(self, title):
        if not simplecustviewer_t.Create(self, title):
            return False
        return True

    def SetMem(self, ql:Qiling):
        self.ClearLines()

        if ql is None:
            return

        try:
            memory = ql.mem.read(self.addr, self.size)
        except:
            return

        size = len(memory)

        view_title = COLSTR("  Memory at [ ", SCOLOR_AUTOCMT)
        view_title += COLSTR("0x%X: %d byte(s)" % (self.addr, size), SCOLOR_DREF)
        view_title += COLSTR(" ]", SCOLOR_AUTOCMT)
        self.AddLine(str(view_title))
        self.AddLine("")
        self.AddLine(COLSTR("                0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F", SCOLOR_AUTOCMT))

        startAddress = self.addr
        line = ""
        chars = ""
        get_char = lambda byte: chr(byte) if 0x20 <= byte <= 0x7E else '.'

        if size != 0:
            for x in range(size):
                if x%16==0:
                    line += COLSTR(" %.12X: " % startAddress, SCOLOR_AUTOCMT)
                if len(self.lastContent) == len(memory):
                    if memory[x] != self.lastContent[x]:
                        line += COLSTR(str("%.2X " % memory[x]), SCOLOR_VOIDOP)
                        chars += COLSTR(get_char(memory[x]), SCOLOR_VOIDOP)
                    else:
                        line += COLSTR(str("%.2X " % memory[x]), SCOLOR_NUMBER)
                        chars += COLSTR(get_char(memory[x]), SCOLOR_NUMBER)
                else:
                    line += COLSTR(str("%.2X " % memory[x]), SCOLOR_NUMBER)
                    chars += COLSTR(get_char(memory[x]), SCOLOR_NUMBER)

                if (x+1)%16==0:
                    line += "  " + chars
                    self.AddLine(line)
                    startAddress += 16
                    line = ""
                    chars = ""

            # add padding
            tail = 16 - size%16
            if tail != 0:
                for x in range(tail): line += "   "
                line += "  " + chars
                self.AddLine(line)

        self.Refresh()
        self.lastContent = memory

    def OnClose(self):
        self.ql_emu_plugin.ql_close_mem_view(self.viewid)

### Dialog Class
class QlEmuMemDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:mem_addr}
BUTTON YES* Add
BUTTON CANCEL Cancel
Show Memory Range
Specify start address and size of new memory range.
<##Address\::{mem_addr}> <##Size\::{mem_size}>
<##Comment\::{mem_cmnt}>
""", {
        'mem_addr': Form.NumericInput(swidth=20, tp=Form.FT_HEX),
        'mem_size': Form.NumericInput(swidth=10, tp=Form.FT_DEC),
        'mem_cmnt': Form.StringInput(swidth=41)
    })

class QlEmuSetupDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:path_name}
BUTTON YES* Start
BUTTON CANCEL Cancel
Setup Qiling
<#Select Rootfs to open#Rootfs path\:        :{path_name}>
<#Custom script path   #Custom script path\: :{script_name}>
<#Custom parameter   #Custom parameter\: :{parameters}>
<#Custom env   #Custom env\: :{env_var}>
""", {
            'path_name': Form.DirInput(swidth=50),
            'script_name': Form.FileInput(swidth=50, open=True),
            'parameters': Form.StringInput(swidth=50),
            'env_var': Form.FileInput(swidth=70, open=True),
        })

class QlEmuSaveDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:path_name}
BUTTON YES* Save
BUTTON CANCEL Cancel
Save Path
<#Save to#Path\::{path_name}>
""", {
        'path_name': Form.FileInput(swidth=50, save=True),
    })

class QlEmuLoadDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:file_name}
BUTTON YES* Load
BUTTON CANCEL Cancel
Load File
<#Load From#File\::{file_name}>
""", {
        'file_name': Form.FileInput(swidth=50, open=True)
    })

class QlEmuAboutDialog(Form):
    def __init__(self, version):
        super(QlEmuAboutDialog, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Qiling Website
Qiling:: About
            {FormChangeCb}
            Qiling IDA plugin v%s with Qiling Framework v%s.
            Author:
            Ziqiao Kong,
            Chenxu Wu,
            Qiling Team.
            Qiling is released under the GPL v2.
            Find more info at https://www.qiling.io
            """ %(version, QLVERSION), {
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(QilingHomePage, new = 2)

        return 1

class QlEmuUpdateDialog(Form):
    def __init__(self, version, message):
        super(QlEmuUpdateDialog, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Qiling Website
Qiling:: Check for update
            {FormChangeCb}
            Your Qiling is v%s
            %s
            """ %(version, message), {
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })
        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(QilingHomePage, new = 2)

        return 1

class QlEmuRegEditDialog(Form):
    def __init__(self, regName):
        Form.__init__(self, r"""STARTITEM {id:reg_val}
BUTTON YES* Save
BUTTON CANCEL Cancel
Register Value
{reg_label}
<##:{reg_val}>
""", {
        'reg_label': Form.StringLabel("Edit [ " + regName + " ] value"),
        'reg_val': Form.NumericInput(tp=Form.FT_HEX, swidth=20)
        })

class QlEmuRegDialog(Choose):
    def __init__(self, reglist, flags=0, width=None, height=None, embedded=False):
        Choose.__init__(
            self, "QL Register Edit",
            [ ["Register", 10 | Choose.CHCOL_PLAIN],
              ["Value", 30] ])
        self.popup_names = ["", "", "Edit Value", ""]
        self.items = reglist

    def show(self):
        return self.Show(True) >= 0

    def OnEditLine(self, n):
        edit_dlg = QlEmuRegEditDialog(self.items[n][0])
        edit_dlg.Compile()
        edit_dlg.reg_val.value = self.items[n][1]
        ok = edit_dlg.Execute()
        if ok == 1:
            newvalue = edit_dlg.reg_val.value
            self.items[n][1] = int("%X" % newvalue, 16)
        self.Refresh()

    def OnGetLine(self, n):
        if self.items[n][2] == 32:
            return [ self.items[n][0], "0x%08X" % self.items[n][1] ]
        if self.items[n][2] == 64:
            return [ self.items[n][0], "0x%16X" % self.items[n][1] ]

    def OnGetSize(self):
        return len(self.items)

    def OnClose(self):
        pass

### Misc
class QlEmuMisc:
    MenuItem = collections.namedtuple("MenuItem", ["action", "handler", "title", "tooltip", "shortcut", "popup"])
    class menu_action_handler(action_handler_t):
        def __init__(self, handler, action):
            action_handler_t.__init__(self)
            self.action_handler = handler
            self.action_type = action

        def activate(self, ctx):
            if ctx.widget_type == BWN_DISASM:
                self.action_handler.ql_handle_menu_action(self.action_type)
            return 1

        # This action is always available.
        def update(self, ctx):
            return AST_ENABLE_ALWAYS

    @staticmethod
    def get_reg_map(ql:Qiling):
        tables = {
            QL_ARCH.X86     : list({**x86_reg_map_32, **x86_reg_map_misc, **x86_reg_map_st}.keys()),
            QL_ARCH.X8664   : list({**x86_reg_map_64, **x86_reg_map_misc, **x86_reg_map_st}.keys()),
            QL_ARCH.ARM     : list({**arm_reg_map}.keys()),
            QL_ARCH.ARM64   : list({**arm64_reg_map}.keys()),
            QL_ARCH.MIPS    : list({**mips_reg_map}.keys()),
        }

        return tables.get(ql.arch.type, [])

    @staticmethod
    def url_download(url):
        try:
            from urllib2 import Request, urlopen, URLError, HTTPError
        except:
            from urllib.request import Request, urlopen
            from urllib.error import URLError, HTTPError

        # create the url and the request
        req = Request(url)

        # Open the url
        try:
            # download this URL
            f = urlopen(req)
            content = f.read()
            return (0, content)

        # handle errors
        except HTTPError as e:
            # print "HTTP Error:", e.code , url
            # fail to download this file
            return (1, None)
        except URLError as e:
            # print "URL Error:", e.reason , url
            # fail to download this file
            return (1, None)
        except Exception as e:
            # fail to save the downloaded file
            # print("Error:", e)
            return (2, None)

    class QLStdIO(ql_file):
        def __init__(self, path, fd):
            super().__init__(path, fd)
            self.__fd = fd

        def write(self, write_buf):
            super().write(write_buf)
            msg(write_buf.decode('utf-8'))

        def flush(self):
            pass

        def isatty(self):
            return False

### Qiling wrapper

class QlEmuQiling:
    def __init__(self):
        self.path = None
        self.rootfs = None
        self.ql: Qiling = None
        self.status = None
        self.exit_addr = None
        self.baseaddr = None
        self.env = {}

    def start(self, *args, **kwargs):
        self.ql = Qiling(argv=self.path, rootfs=self.rootfs, verbose=QL_VERBOSE.DEBUG, env=self.env, log_plain=True, *args, **kwargs)

        if sys.platform != 'win32':
            self.ql.os.stdin = QlEmuMisc.QLStdIO('stdin', sys.__stdin__.fileno())
            self.ql.os.stdout = QlEmuMisc.QLStdIO('stdout', sys.__stdout__.fileno())
            self.ql.os.stderr = QlEmuMisc.QLStdIO('stderr', sys.__stderr__.fileno())

        self.exit_addr = self.ql.os.exit_point
        if self.ql.os.type == QL_OS.LINUX:
            f = open(self.ql.path, 'rb')
            elffile = ELFFile(f)
            elf_header = elffile.header
            if elf_header['e_type'] == 'ET_EXEC':
                self.baseaddr = self.ql.os.elf_mem_start
            elif elf_header['e_type'] == 'ET_DYN':
                if self.ql.arch.bits == 32:
                    self.baseaddr = int(self.ql.os.profile.get("OS32", "load_address"), 16)
                elif self.ql.arch.bits == 64:
                    self.baseaddr = int(self.ql.os.profile.get("OS64", "load_address"), 16)
        else:
            self.baseaddr = 0x0

    def run(self, begin=None, end=None):
        self.ql.run(begin, end)

    def set_reg(self):
        reglist = QlEmuMisc.get_reg_map(self.ql)
        regs = [ [ row, int(self.ql.arch.regs.read(row)), self.ql.arch.bits ] for row in reglist ]
        regs_len = len(regs)
        RegDig = QlEmuRegDialog(regs)
        if RegDig.show():
            for idx, val in enumerate(RegDig.items[0:regs_len-1]):
                self.ql.arch.regs.write(reglist[idx], val[1])
            return True
        else:
            return False

    def save(self):
        savedlg = QlEmuSaveDialog()
        savedlg.Compile()

        if savedlg.Execute() != 1:
            return False

        savepath = savedlg.path_name.value

        self.ql.save(reg=True, mem=True,fd=True, cpu_context=True, snapshot=savepath)
        logging.info('Save to ' + savepath)
        return True

    def load(self):
        loaddlg = QlEmuLoadDialog()
        loaddlg.Compile()

        if loaddlg.Execute() != 1:
            return False

        loadname = loaddlg.file_name.value

        self.ql.restore(snapshot=loadname)
        logging.info('Restore from ' + loadname)
        return True

    def remove_ql(self):
        if self.ql is not None:
            del self.ql
            self.ql = None

    def ql_addr_from_ida(self, addr):
        return addr - get_imagebase() + self.baseaddr

    def ida_addr_from_ql_addr(self, addr):
        return addr - self.baseaddr + get_imagebase()

### Plugin

class QlEmuPlugin(plugin_t, UI_Hooks):
    ### Ida Plugin Data

    popup_menu_hook = None

    flags = PLUGIN_HIDE
    comment = ""

    help = "Qiling Emulator"
    wanted_name = "Qiling Emulator"
    wanted_hotkey = ""

    ### View Data

    qlemuregview = None
    qlemustackview = None
    qlemumemview = {}

    def __init__(self):
        super(QlEmuPlugin, self).__init__()
        self.plugin_name = "Qiling Emulator"
        self.qlemu = None
        self.ql = None
        self.stepflag = True
        self.stephook = None
        self.qlinit = False
        self.lastaddr = None
        self.is_change_addr = -1
        self.userobj = None
        self.customscriptpath = None
        self.bb_mapping = {}
        self.ks = None

    ### Main Framework

    def init(self):
        # init data
        logging.info('---------------------------------------------------------------------------------------')
        logging.info('Qiling Emulator Plugin For IDA, by Qiling Team. Version {0}, 2020'.format(QLVERSION))
        logging.info('Based on Qiling v{0}'.format(QLVERSION))
        logging.info('Find more information about Qiling at https://qiling.io')
        logging.info('---------------------------------------------------------------------------------------')
        self.qlemu = QlEmuQiling()
        self.ql_hook_ui_actions()
        return PLUGIN_KEEP

    def run(self, arg = 0):
        logging.info(f"Registering actions.")
        self.ql_register_menu_actions()
        self.ql_attach_main_menu_actions()

    def ready_to_run(self):
        logging.info(f"UI is ready, register our menu actions.")
        self.run()

    def term(self):
        self.qlemu.remove_ql()
        self.ql_unhook_ui_actions()
        self.ql_detach_main_menu_actions()
        self.ql_unregister_menu_actions()

    ### Actions

    def ql_start(self):
        if self.qlemu is None:
            self.qlemu = QlEmuQiling()
        if self.ql_set_rootfs():
            logging.info(f'Rootfs: {self.qlemu.rootfs}')
            logging.info(f"Custom user script: {self.customscriptpath}")
            logging.info(f"Custom env: {self.qlemu.env}")
            show_wait_box("Qiling is processing ...")
            try:
                self.qlemu.start()
                self.qlinit = True
                self.lastaddr = None
            finally:
                hide_wait_box()
                logging.info("Qiling is initialized successfully.")
        if self.customscriptpath is not None:
            self.ql_load_user_script()
            self.userobj.custom_prepare(self.qlemu.ql)

        ida_ida.ql_plugin = self

    def ql_load_user_script(self):
        if self.qlinit :
            self.ql_get_user_script(is_reload=True, is_start=True)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_reload_user_script(self):
        if self.qlinit:
            self.ql_get_user_script(is_reload=True)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_continue(self):
        if self.qlinit:
            userhook = None
            pathhook = self.qlemu.ql.hook_code(self.ql_path_hook)
            if self.userobj is not None:
                userhook = self.userobj.custom_continue(self.qlemu.ql)
            if self.qlemu.status is not None:
                self.qlemu.ql.restore(self.qlemu.status)
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run(begin=self.qlemu.ql.arch.regs.arch_pc, end=self.qlemu.exit_addr)
                finally:
                    hide_wait_box()
            else:
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run()
                finally:
                    hide_wait_box()
            self.qlemu.ql.hook_del(pathhook)
            if userhook and userhook is not None:
                for hook in userhook:
                    self.qlemu.ql.hook_del(hook)
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def _color_path(self, color):
        def _cb(ql, addr, size):
            real_addr = addr - self.qlemu.baseaddr + get_imagebase()
            set_color(real_addr, CIC_ITEM, color)
        return _cb

    def ql_run_selection(self):
        if self.qlinit:
            _, start, end = IDA.get_last_selection()
            # Drop all previous status
            self.qlemu.status = None
            if self.userobj is not None:
                userhook = self.userobj.custom_execute_selection(self.qlemu.ql)
            colorhook = self.qlemu.ql.hook_code(self._color_path(Colors.Pink.value))
            show_wait_box("Qiling is processing ...")
            try:
                self.qlemu.run(begin=start, end=end)
            finally:
                hide_wait_box()
            self.qlemu.ql.hook_del(colorhook)
            if userhook and userhook is not None:
                for hook in userhook:
                    self.qlemu.ql.hook_del(hook)
            self.qlemu.status = self.qlemu.ql.save()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_set_pc(self):
        if self.qlinit:
            ea = IDA.get_current_address()
            self.qlemu.ql.arch.regs.arch_pc = ea
            logging.info(f"QIling PC set to {hex(ea)}")
            self.qlemu.status = self.qlemu.ql.save()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_run_to_here(self):
        if self.qlinit:
            curr_addr = get_screen_ea()
            untillhook = self.qlemu.ql.hook_code(self.ql_untill_hook)
            if self.qlemu.status is not None:
                self.qlemu.ql.restore(self.qlemu.status)
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run(begin=self.qlemu.ql.arch.regs.arch_pc, end=curr_addr+self.qlemu.baseaddr-get_imagebase())
                finally:
                    hide_wait_box()
            else:
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run(end=curr_addr+self.qlemu.baseaddr-get_imagebase())
                finally:
                    hide_wait_box()

            set_color(curr_addr, CIC_ITEM, 0x00B3CBFF)
            self.qlemu.ql.hook_del(untillhook)
            self.qlemu.status = self.qlemu.ql.save()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_step(self):
        if self.qlinit:
            userhook = None
            self.stepflag = True
            self.qlemu.ql.restore(saved_states=self.qlemu.status)
            self.stephook = self.qlemu.ql.hook_code(callback=self.ql_step_hook)
            if self.userobj is not None:
                userhook = self.userobj.custom_step(self.qlemu.ql)
            self.qlemu.run(begin=self.qlemu.ql.arch.regs.arch_pc, end=self.qlemu.exit_addr)
            if userhook and userhook is not None:
                for hook in userhook:
                    self.qlemu.ql.hook_del(hook)
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_save(self):
        if self.qlinit:
            if self.qlemu.save() != True:
                logging.error('Fail to save the snapshot.')
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_load(self):
        if self.qlinit:
            if self.qlemu.load() != True:
                logging.error('Fail to load the snapshot.')
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_chang_reg(self):
        if self.qlinit:
            self.qlemu.set_reg()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
            self.qlemu.status = self.qlemu.ql.save()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_reset(self):
        if self.qlinit:
            self.ql_close()
            self.qlemu = QlEmuQiling()
            self.ql_start()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_close(self):
        if self.qlinit:
            heads = Heads(get_segm_start(get_screen_ea()), get_segm_end(get_screen_ea()))
            for i in heads:
                set_color(i, CIC_ITEM, 0xFFFFFF)
            self.qlemu.remove_ql()
            del self.qlemu
            self.qlemu = None
            self.qlinit = False
            logging.info('Qiling is deleted.')
        else:
            logging.error('Qiling is not started.')

    def ql_show_reg_view(self):
        if self.qlinit:
            if self.qlemuregview is None:
                self.qlemuregview = QlEmuRegView(self)
                QlEmuRegView(self)
                self.qlemuregview.Create()
                self.qlemuregview.SetReg(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
                self.qlemuregview.Show()
                self.qlemuregview.Refresh()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_show_stack_view(self):
        if self.qlinit:
            if self.qlemustackview is None:
                self.qlemustackview = QlEmuStackView(self)
                self.qlemustackview.Create()
                self.qlemustackview.SetStack(self.qlemu.ql)
                self.qlemustackview.Show()
                self.qlemustackview.Refresh()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_show_mem_view(self, addr=get_screen_ea(), size=0x10):
        if self.qlinit:
            memdialog = QlEmuMemDialog()
            memdialog.Compile()
            memdialog.mem_addr.value = addr
            memdialog.mem_size.value = size
            ok = memdialog.Execute()
            if ok == 1:
                mem_addr = memdialog.mem_addr.value - self.qlemu.baseaddr + get_imagebase()
                mem_size = memdialog.mem_size.value
                mem_cmnt = memdialog.mem_cmnt.value

                if mem_addr not in self.qlemumemview:
                    if not self.qlemu.ql.mem.is_mapped(mem_addr, mem_size):
                        ok = ask_yn(1, "Memory [%X:%X] is not mapped!\nDo you want to map it?\n   YES - Load Binary\n   NO - Fill page with zeroes\n   Cancel - Close dialog" % (mem_addr, mem_addr + mem_size))
                        if ok == 0:
                            self.qlemu.ql.mem.map(mem_addr, mem_size)
                            self.qlemu.ql.mem.write(mem_addr, b"\x00"*mem_size)
                        elif ok == 1:
                            # TODO: map_binary
                            return
                        else:
                            return
                    self.qlemumemview[mem_addr] = QlEmuMemView(self, mem_addr, mem_size)
                    if mem_cmnt == []:
                        self.qlemumemview[mem_addr].Create("QL Memory")
                    else:
                        self.qlemumemview[mem_addr].Create("QL Memory [ " + mem_cmnt + " ]")
                    self.qlemumemview[mem_addr].SetMem(self.qlemu.ql)
                self.qlemumemview[mem_addr].Show()
                self.qlemumemview[mem_addr].Refresh()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_unload_plugin(self):
        heads = Heads(get_segm_start(get_screen_ea()), get_segm_end(get_screen_ea()))
        for i in heads:
            set_color(i, CIC_ITEM, 0xFFFFFF)
        self.ql_close()
        self.ql_detach_main_menu_actions()
        self.ql_unregister_menu_actions()
        logging.info('Unload plugin successfully!')

    def ql_menu_null(self):
        pass

    def ql_about(self):
        self.aboutdlg = QlEmuAboutDialog(QLVERSION)
        self.aboutdlg.Execute()
        self.aboutdlg.Free()

    def ql_check_update(self):
        (r, content) = QlEmuMisc.url_download(QilingStableVersionURL)
        content = content.decode("utf-8")
        if r == 0:
            try:
                version_stable = re.findall(r"\"([\d\.]+)\"", content)[0]
            except (TypeError, IndexError):
                warning("ERROR: Failed to find the Qiling version string from response.")
                logging.warning("Failed to find the Qiling version string from response.")

            # compare with the current version
            if version_stable == QLVERSION:
                self.updatedlg = QlEmuUpdateDialog(QLVERSION, "Good, you are already on the latest stable version!")
                self.updatedlg.Execute()
                self.updatedlg.Free()
            else:
                self.updatedlg = QlEmuUpdateDialog(QLVERSION, "Download latest stable version {0} from https://github.com/qilingframework/qiling/blob/master/qiling/extensions/idaplugin".format(version_stable))
                self.updatedlg.Execute()
                self.updatedlg.Free()
        else:
            # fail to download
            warning("ERROR: Failed to connect to Github. Try again later.")
            logging.warning("Failed to connect to Github when checking for the latest update. Try again later.")

    def _remove_from_bb_lists(self, bbid):
        if bbid in self.real_blocks:
            self.real_blocks.remove(bbid)
        elif bbid in self.fake_blocks:
            self.fake_blocks.remove(bbid)
        elif bbid in self.retn_blocks:
            self.retn_blocks.remove(bbid)

    def ql_mark_real(self):
        if len(self.bb_mapping) == 0:
            logging.error(f"Please perform auto analysis before marking blocks manually!")
            return
        cur_addr = IDA.get_current_address()
        cur_block = IDA.get_block(cur_addr)
        self._remove_from_bb_lists(cur_block.id)
        self.real_blocks.append(cur_block.id)
        IDA.color_block(cur_block, Colors.Green.value)

    def ql_mark_fake(self):
        if len(self.bb_mapping) == 0:
            logging.error(f"Please perform auto analysis before marking blocks manually!")
            return
        cur_addr = IDA.get_current_address()
        cur_block = IDA.get_block(cur_addr)
        self._remove_from_bb_lists(cur_block.id)
        self.fake_blocks.append(cur_block.id)
        IDA.color_block(cur_block, Colors.Gray.value)

    def ql_mark_retn(self):
        if len(self.bb_mapping) == 0:
            logging.error(f"Please perform auto analysis before marking blocks manually!")
            return
        cur_addr = IDA.get_current_address()
        cur_block = IDA.get_block(cur_addr)
        self._remove_from_bb_lists(cur_block.id)
        self.retn_blocks.append(cur_block.id)
        IDA.color_block(cur_block, Colors.Pink.value)

    def __in_bb(self, addr, bb):
        return addr < bb.end_ea and addr >= bb.start_ea

    # Identify if the given microcode block is a jmp_mbb.
    # Core patterns:
    #     ...
    #     ...
    #     goto @dispatcher
    def __is_jmp_mbb(self, mbb):
        ins_list = list(IDA.micro_code_from_mbb(mbb))
        if len(ins_list) == 0:
            logging.warning(f"Get an empty mbb at {hex(mbb.start)}?!")
            return False
        last_ins = ins_list[-1]
        if last_ins.opcode != ida_hexrays.m_goto:
            logging.warning(f"jmp_mbb at {hex(mbb.start)} the opcode of last instruction {last_ins._print()} isn't goto")
            return False
        if last_ins.l.t != ida_hexrays.mop_b:
            logging.warning(f"jmp_mbb at {hex(mbb.start)} the l of last instruction {last_ins._print()} doesn't have a microcode block reference")
            return False
        goto_mbb = self.mbbs[last_ins.l.b]
        mbb_start = goto_mbb.start
        dispatcher_bb = self.bb_mapping[self.dispatcher]
        pre_dispatcher_bb = self.bb_mapping[self.pre_dispatcher]
        if self.__in_bb(mbb_start, dispatcher_bb) or self.__in_bb(mbb_start, pre_dispatcher_bb):
            return True
        logging.warning(f"The address {hex(mbb_start)} where jmp_mbb goes isn't pre_dispatcher or dispatcher block!")
        return False

    # Identify if the given microcode block is a next_mbb.
    # Core patterns:
    #     mov #imm, reg
    def __is_next_mbb(self, mbb):
        ins_list = list(IDA.micro_code_from_mbb(mbb))
        if len(ins_list) == 0:
            logging.warning(f"Get an empty mbb at {hex(mbb.start)}?!")
            return False
        first_ins = ins_list[0]
        if first_ins.opcode != ida_hexrays.m_mov:
            logging.warning(f"next_mbb at {hex(mbb.start)} the opcode of first instruction {first_ins._print()} isn't mov")
            return False
        if first_ins.l.t != ida_hexrays.mop_n:
            logging.warning(f"next_mbb at {hex(mbb.start)} the l of first instruction {first_ins._print()} isn't an immediate number")
            return False
        if first_ins.d.t != ida_hexrays.mop_r:
            logging.warning(f"next_mbb at {hex(mbb.start)} the d of first instruction {first_ins._print()} isn't a reg")
            return False
        if len(ins_list) == 1:
            logging.info(f"A block with only one instruction which is `mov #imm, reg` at {hex(mbb.start)}.")
        return True

    def _get_jmp_ins(self, ida_addr, insns):
        # This cloud really happen! See issue #804. TODO: Investigate or re-design insns structure or replace it with ESIL.
        # So we have to fallback to legacy path.
        if ida_addr not in insns:
            return (None, None)
        ins_list = insns[ida_addr]
        result = []
        for bbid, ins in ins_list:
            if ida_hexrays.is_mcode_jcond(ins.opcode):
                result.append((bbid, ins))
        if len(result) > 1:
            logging.warning(f"More than one conditional jmp detected at {hex(ida_addr)}!")
        elif len(result) == 0:
            logging.warning(f"No conditional jmp found at {hex(ida_addr)}!")
            return (None, None)
        return result[0]

    # cmov/it eq:
    #   < ... condiontal jump > --> < mov imm, reg > (next_bb) --> < dispatcher >
    #                           |-> < dispatcher > (jmp_mbb)
    # NOTE:
    #   LLVM IR -> ASM -> IDA IR
    #   I guess this pattern is stable enough for us to rely on.
    def _force_execution_with_microcode(self, ql, ida_addr):
        _, insns, mbbs = self._prepare_microcodes(maturity=7)
        bbid, ins = self._get_jmp_ins(ida_addr, insns)
        if ins is None:
            return False
        # According to comments in hexrays.hpp, it may be a mop_v. I guess that shouldn't exist
        # so we add a sanity check here.
        if ins.d.t != ida_hexrays.mop_b:
            logging.warning(f"Sanity check: microcode {ins._print()} doesn't refer a block!")
        jmp_mbb = mbbs[ins.d.b]
        next_mbb = mbbs[bbid].nextb
        if not (self.__is_next_mbb(next_mbb) and self.__is_jmp_mbb(jmp_mbb)):
            # Switch the branch and try again?
            logging.info("Switch the jmp_bb and next_bb and try again...")
            if self.__is_jmp_mbb(next_mbb) and self.__is_next_mbb(jmp_mbb):
                jmp_mbb, next_mbb = next_mbb, jmp_mbb
            else:
                logging.error(f"Fail to identify microcode blocks at {hex(ida_addr)}")
                return False
        ins_list = list(IDA.micro_code_from_mbb(next_mbb))
        first_ins = ins_list[0]
        imm = first_ins.l.nnn.value
        reg_name = ida_hexrays.get_mreg_name(first_ins.d.r, ql.arch.pointersize)
        logging.info(f"Froce set {reg_name} to {hex(imm)}")
        ql.arch.regs.__setattr__(reg_name, imm)
        return True

    def _ida_address_after_branch(self, ida_addr):
        _, ins = self._get_jmp_ins(ida_addr, self.insns)
        if ins is None:
            return None
        return self.mbbs[ins.d.b].start

    # legacy approach for x86_64 and arm64
    # TODO & NOTE: We still need this fallback since the microcode pattern we match now sometimes may change. A better
    #              approach is that we fully drop pattern machting and emulates the microcode while recording the correspoding
    #              branch condition.
    #              For example:
    #                CSEL            W8, W8, W9, LT
    #                =>
    #                'LT' -> W9->W8 (the value of W9 can be get from emulation)
    #                'GT' -> W8->W8
    #                We assume that the comparison (like cmp W10, 0) is the same between real assembly and microcode and the only
    #                difference is the condition of the jump instruction. So the 'LT' and 'GT' here represent condition both in microcode
    #                and original assembly.
    #              For current implementation we don't record 'LT' or 'GT' which makes it hard to patch code after force execution.
    def _force_execution_by_parsing_assembly(self, ql, ida_addr):
        if "x86" in IDA.get_ql_arch_string(): # cmovlg eax, ebx
            reg1 = IDA.print_operand(ida_addr, 0).lower()
            reg2 = IDA.print_operand(ida_addr, 1).lower()
            reg2_val = ql.arch.regs.__getattr__(reg2)
            logging.info(f"Force set {reg1} to {hex(reg2_val)}")
            ql.arch.regs.__setattr__(reg1, reg2_val)
            return True
        elif "arm" in IDA.get_ql_arch_string():
            instr = IDA.get_instruction(ida_addr).lower()
            logging.info(f"Going to force execute: {instr}")
            if instr.startswith("it"): # itt eq\n moveqw low\n movteq high\n
                ida_addr = ida_addr + IDA.get_instruction_size(ida_addr)
                low = IDA.get_operand(ida_addr, 1)
                ida_addr = ida_addr + IDA.get_instruction_size(ida_addr)
                high = IDA.get_operand(ida_addr, 1)
                reg = IDA.print_operand(ida_addr, 0).lower()
                val = (high << 16) + low
                logging.info(f"Force set {reg} to {hex(val)}")
                ql.arch.regs.__setattr__(reg, val)
                return True
            elif "csel" in instr: # csel dst, src1, src2, cond
                dst = IDA.print_operand(ida_addr, 0).lower()
                src = IDA.print_operand(ida_addr, 2).lower()
                src_val = ql.arch.regs.__getattr__(src)
                logging.info(f"Force set {dst} to {hex(src_val)}")
                ql.arch.regs.__setattr__(dst, src_val)
                return True
        return False

    def _has_call_insn(self, ida_addr):
        if ida_addr not in self.insns:
            return False
        ins_list = self.insns[ida_addr]
        for _, ins in ins_list:
            if ida_hexrays.is_mcode_call(ins.opcode):
                return True
        return False

    def _guide_hook(self, ql, addr, size):
        start_bb_id = self.hook_data['startbb']
        ida_addr = self.deflatqlemu.ida_addr_from_ql_addr(addr)
        func = self.hook_data['func']
        if ida_addr < func.start_ea or ida_addr >= func.end_ea:
            logging.error(f"Address {hex(ida_addr)} out of function boundaries!")
            ql.emu_stop()
            self.hook_data['result'] = False
            return
        cur_bb = IDA.get_block(ida_addr)
        if "force" in self.hook_data and ida_addr in self.hook_data['force']:
            if self.hook_data['force'][ida_addr]:
                logging.info(f"Force execution at {hex(ida_addr)}")
                result = self._force_execution_with_microcode(ql, ida_addr)
                if not result:
                    logging.warning(f"Fail to force execution by microcode at {hex(ida_addr)}, trying legacy approach")
                    result = self._force_execution_by_parsing_assembly(ql, ida_addr)
                    if not result:
                        logging.error(f"Fail to force execution by legacy approach at {hex(ida_addr)}, stop now...")
                        self.hook_data['result'] = False
                        ql.emu_stop()
                        return
                    self.hook_data['force_legacy'] = True
            else:
                pass
            if not self.hook_data['force_legacy']:
                next_ida_addr = self._ida_address_after_branch(ida_addr)
            else:
                next_ida_addr = ida_addr + IDA.get_instruction_size(ida_addr)
            logging.info(f"Goto {hex(next_ida_addr)} after branch...")
            ql.arch.regs.arch_pc = self.deflatqlemu.ql_addr_from_ida(next_ida_addr) + self.append
            ida_addr = next_ida_addr
        # TODO: Maybe we can detect whether the program will access unmapped
        #       here so that we won't map the memory.
        if self._has_call_insn(ida_addr):
            ql.arch.regs.arch_pc += IDA.get_instruction_size(ida_addr) + self.append
            return
        if start_bb_id == cur_bb.id:
            return
        if cur_bb.id in self.real_blocks or cur_bb.id in self.retn_blocks:
            if cur_bb.id not in self.paths[start_bb_id]:
                self.paths[start_bb_id].append(cur_bb.id)
            ql.emu_stop()

    def _skip_unmapped_rw(self, ql: Qiling, type, addr, size, value):
        map_addr = ql.mem.align(addr)
        map_size = ql.mem.align_up(size)

        if not ql.mem.is_mapped(map_addr, map_size):
            logging.warning(f"Invalid memory R/W, trying to map {hex(map_size)} at {hex(map_addr)}")

            ql.mem.map(map_addr, map_size)
            ql.mem.write(map_addr, b'\x00' * map_size)

        return True

    def _find_branch_in_real_block(self, bb):
        #paddr = bb.start_ea
        #while paddr < bb.end_ea:
        #    ins = IDA.get_instruction(paddr)
        #    sz = IDA.get_instruction_size(paddr)
        #    if ins.lower().startswith("cmov"):
        #        return paddr
        #    paddr += sz

        # microcode implementation
        #  Is it possible for the address of a microcode instruction not equal
        #  to any start address of a real instruction? IDK, only Hexrays knows.
        #  So we iterate all addresses so that we won't miss any posssible branchs.
        for addr in range(bb.start_ea, bb.end_ea):
            if addr in self.insns:
                for _, insn in self.insns[addr]:
                    opcode = insn.opcode
                    if ida_hexrays.is_mcode_jcond(opcode):
                        return addr
        return None

    def _log_paths_str(self):
        for bbid, succs in self.paths.items():
            if len(succs) == 1:
                logging.info(f"{self._block_str(bbid)} -> {self._block_str(succs[0])}")
            elif len(succs) == 2:
                logging.info(f"{self._block_str(bbid)} --(force jump)--> {self._block_str(succs[0])}")
                logging.info(f"|----(skip jump)----> {self._block_str(succs[1])}")
            elif len(succs) > 2:
                logging.warning(f"succs: {succs} found from {self._block_str(bbid)}!")

    # Is this correct?
    def _thumb_detect(self, ida_addr):
        return IDA.get_instruction_size(ida_addr) == 2

    def _log_verbose(self, ql, addr, size):
        logging.debug(f"addr: {hex(addr)} ida_addr: {hex(self.deflatqlemu.ida_addr_from_ql_addr(addr))}")
        registers = [ k for k in ql.arch.regs.register_mapping.keys() if type(k) is str ]
        for idx in range(0, len(registers), 3):
            regs = registers[idx:idx+3]
            s = "\t".join(map(lambda v: f"{v:4}: {ql.arch.regs.__getattr__(v):016x}", regs))
            logging.debug(s)

    # Q: Why we need emulation to help us find real control flow considering there are some
    #    switch-case patterns in mircocode which can be analysed statically?
    # A: Emulation makes the de-obf much more robust and general and can work under less assumptions, like
    #    the optimization level, the complexity of original control flow or even the verson of the
    #    Hexrays decompiler because we try to avoid matching fixed microcode patterns. In fact, static
    #    analysis is always possible since the brain of a reverse engineer is the best decompiler
    #    and emulator, isn't it? :D
    def _search_path(self):
        self.paths = {bbid: [] for bbid in self.bb_mapping.keys()}
        reals = [self.first_block, *self.real_blocks]
        self.deflatqlemu = QlEmuQiling()
        self.deflatqlemu.path = self.qlemu.path
        self.deflatqlemu.rootfs = self.qlemu.rootfs
        first_block = self.bb_mapping[self.first_block]
        if IDA.get_ql_arch_string() == "arm32":
            if self._thumb_detect(first_block.start_ea):
                logging.info(f"Thumb detected, enable it.")
                self.deflatqlemu.start(archtype=QL_ARCH.ARM, thumb=True)
                self.deflatqlemu.ql.arch.regs.cpsr |= 0x20
                self.append = 1
        else:
            self.deflatqlemu.start()
            self.append = 0
        ql = self.deflatqlemu.ql
        if logging.root.level <= logging.DEBUG:
            ql.hook_code(self._log_verbose)
        self.hook_data = None
        ql.hook_mem_read_invalid(self._skip_unmapped_rw)
        ql.hook_mem_write_invalid(self._skip_unmapped_rw)
        ql.hook_mem_unmapped(self._skip_unmapped_rw)
        # set up stack before we really run.
        ql.run(begin=self.deflatqlemu.ql_addr_from_ida(first_block.start_ea) + self.append, end=self.deflatqlemu.ql_addr_from_ida(first_block.end_ea), count=0xFFF)
        # okay, we can set up our core hook now.
        ql.hook_code(self._guide_hook)
        for bbid in reals:
            logging.debug(f"Search control flow for block: {self._block_str(bbid)}")
            bb = self.bb_mapping[bbid]
            braddr = self._find_branch_in_real_block(bb)
            self.hook_data = {
                "startbb": bbid,
                "func": IDA.get_function(first_block.start_ea),
                "result": True,
                "force_legacy": False
            }
            ql_bb_start_ea = self.deflatqlemu.ql_addr_from_ida(bb.start_ea) + self.append
            ctx = ql.save()
            # Skip force execution in the first block.
            # `end=0` is a workaround for ql remembering last exit_point.
            if braddr is None or bb.id == self.first_block:
                ql.run(begin=ql_bb_start_ea, end=0, count=0xFFF)
            else:
                self.hook_data['force'] = {braddr: True}
                ctx2 = ql.save()
                ql.run(begin=ql_bb_start_ea, end=0, count=0xFFF)
                ql.restore(ctx2)
                if not self.hook_data['result']:
                    return False
                self.hook_data['force'] = {braddr: False}
                ql.run(begin=ql_bb_start_ea, end=0, count=0xFFF)
            ql.restore(ctx)
            if not self.hook_data['result']:
                return False
        self._log_paths_str()
        return True


    # IDA doesn't support arm assembling, so it's a good chance to replace IDA
    # assembler implmentation with keystone.
    def _initialize_keystone(self):
        if self.ks is None:
            self.ks = self.deflatqlemu.ql.arch.assembler


    def _asm(self, *args, **kwargs):
        self._initialize_keystone()
        return self.ks.asm(*args, **kwargs)

    def _assemble(self, instr, addr):
        logging.debug(f"Keystone: Assemble {instr} at {hex(addr)}")
        bs, count = self._asm(instr, addr)
        IDA.patch_bytes(addr, bytes(bs))
        IDA.perform_analysis(addr, addr + len(bs))
        return bs, count

    # Patching microcode is TOO complex.
    # I would rahter write another 1e10 llvm passes than a single hexrays decompiler pass.
    def _arch_jmp_instruction(self, addr):
        arch = IDA.get_ql_arch_string()
        op = None
        if "x86" in arch:
            op = "jmp"
        elif "arm" in arch:
            op = "B"
        elif "mips" in arch:
            op = "j"
        return f"{op} {addr}"

    # See comments above.
    def _arch_cond_jmp_instruction(self, cond, addr):
        arch = IDA.get_ql_arch_string()
        op = None
        if "x86" in arch:
            op = f"j{cond}"
        elif "arm" in arch:
            op = f"b{cond}"
        elif "mips" in arch:
            op = f"j{cond}"
        return f"{op} {addr}"

    # See comments above.
    def _arch_parse_cond_from_addr(self, braddr):
        arch = IDA.get_ql_arch_string()
        instr = IDA.get_instruction(braddr).lower()
        if "x86" in arch: # cmovge
            return instr[4:]
        elif "arm" in arch:
            if instr.startswith("it"): # itt eq
                tks = instr.split(" ")
                if len(tks) != 2:
                    logging.error(f"Can't get condition from {instr}")
                    return None
                return tks[-1]
            elif "csel" in instr:
                return IDA.print_operand(braddr, 3)
        # TODO: mips
        return None

    # Really FORCE ida to analyse again.
    def _force_analysis(self, start, end):
        ida_bytes.del_items(start, 0, end-start)
        ida_auto.plan_and_wait(start, end)
        return

    # The only way to make IDA generate right assembly is to undefine bytes patched just now
    # and call IDA to analyse it again. Nice job again as always, Hexrays!
    def _patch_bytes_with_force_analysis(self, start, bs):
        IDA.patch_bytes(start, bs)
        return self._force_analysis(start, start+len(bs))

    def _arch_branch_patch(self, braddr, bbid):
        bb = self.bb_mapping[bbid]
        force_addr = self.bb_mapping[self.paths[bbid][0]].start_ea
        normal_addr = self.bb_mapping[self.paths[bbid][1]].start_ea
        # Temporary dirty fix.
        # See comments for _force_execution_by_parsing_assembly.
        if "arm64" == IDA.get_ql_arch_string():
            force_addr, normal_addr = normal_addr, force_addr
        # Parse condition before patching nop.
        cond = self._arch_parse_cond_from_addr(braddr)
        buffer = [0] * (bb.end_ea - braddr)
        instr_to_assemble = self._arch_cond_jmp_instruction(cond, f"{hex(force_addr)}h")
        logging.info(f"Assemble {instr_to_assemble} at {hex(force_addr)}")
        bs1, _ = self._asm(instr_to_assemble, braddr)
        buffer[:len(bs1)] = bs1
        next_instr_address = braddr + len(bs1)
        instr_to_assemble = self._arch_jmp_instruction(f"{hex(normal_addr)}h")
        logging.info(f"Assemble {instr_to_assemble} at {hex(normal_addr)}")
        bs2, _ = self._asm(instr_to_assemble, next_instr_address)
        buffer[len(bs1):len(bs1) + len(bs2)] = bs2
        logging.info(f"Patch real block with branch from {hex(braddr)} to {hex(bb.end_ea)}")
        self._patch_bytes_with_force_analysis(braddr, bytes(buffer))


    def _patch_codes(self):
        if len(self.paths[self.first_block]) != 1:
            logging.error(f"Found wrong ways in first block: {self._block_str(self.bb_mapping[self.first_block])}, should be 1 path but get {len(self.paths[self.first_block])}, exit.")
            return
        logging.info("NOP dispatcher block")
        dispatcher_bb = self.bb_mapping[self.dispatcher]
        # Some notes:
        #    Patching b'\x00' instead of 'nop' can help IDA decompile a better result. Don't know why...
        #    Besides
        buffer = [0] * (dispatcher_bb.end_ea - dispatcher_bb.start_ea)
        first_jmp_addr = dispatcher_bb.start_ea
        instr_to_assemble = self._arch_jmp_instruction(f"{hex(self.bb_mapping[self.paths[self.first_block][0]].start_ea)}h")
        logging.info(f"Assemble {instr_to_assemble} at {hex(first_jmp_addr)}")
        bs, _ = self._asm(instr_to_assemble, first_jmp_addr)
        buffer[:len(bs)] = bs
        logging.info(f"Patch first jump at {hex(first_jmp_addr)}")
        self._patch_bytes_with_force_analysis(first_jmp_addr, bytes(buffer))
        for bbid in self.real_blocks:
            logging.debug(f"Patching real block: {self._block_str(bbid)}")
            bb = self.bb_mapping[bbid]
            braddr = self._find_branch_in_real_block(bb)
            if braddr is None:
                last_instr_address = IDA.get_prev_head(bb.end_ea)
                buffer = [0x90] * (bb.end_ea - last_instr_address)
                if len(self.paths[bbid]) != 1:
                    logging.warning(f"Found wrong ways in block: {self._block_str(bb)}, should be 1 path but get {len(self.paths[bbid])}")
                    continue
                instr_to_assemble = self._arch_jmp_instruction(f"{hex(self.bb_mapping[self.paths[bbid][0]].start_ea)}h")
                logging.info(f"Assemble {instr_to_assemble} at {hex(last_instr_address)}")
                bs, _ = self._asm(instr_to_assemble, last_instr_address)
                buffer[:len(bs)] = bs
                logging.info(f"Patch real block from {hex(last_instr_address)} to {hex(bb.end_ea)}")
                self._patch_bytes_with_force_analysis(last_instr_address, bytes(buffer))
            else:
                if len(self.paths[bbid]) != 2:
                    logging.warning(f"Found wrong ways in block: {self._block_str(bb)}, should be 2 paths but get {len(self.paths[bbid])}")
                    continue
                self._arch_branch_patch(braddr, bbid)
        for bbid in self.fake_blocks:
            bb = self.bb_mapping[bbid]
            logging.info(f"Patch NOP for block: {self._block_str(bb)}")
            self._patch_bytes_with_force_analysis(bb.start_ea, b"\x00"*(bb.end_ea-bb.start_ea))
        logging.info(f"Patch NOP for pre_dispatcher.")
        bb = self.bb_mapping[self.pre_dispatcher]
        self._patch_bytes_with_force_analysis(bb.start_ea, b"\x00"*(bb.end_ea-bb.start_ea))

    def _prepare_microcodes(self, decomp_flags=ida_hexrays.DECOMP_WARNINGS | ida_hexrays.DECOMP_NO_WAIT, maturity=7):
        dispatcher_bb = self.bb_mapping[self.dispatcher]
        target_function = IDA.get_function(dispatcher_bb.start_ea)
        # Reduce optimization to make pattern more stable.
        logging.info(f"Generate microcode from {hex(target_function.start_ea)} to {hex(target_function.end_ea)}")
        mba = IDA.get_micro_code_mba(target_function.start_ea, target_function.end_ea, decomp_flags, maturity)
        insns = {}
        mbbs = {}
        for i in range(mba.qty):
            mbb = mba.get_mblock(i)
            mbbs[i] = mbb
            cur_insn = mbb.head
            while cur_insn != None:
                insn_ea = cur_insn.ea
                if insn_ea not in insns:
                    insns[insn_ea] = []
                insns[insn_ea].append((i, cur_insn))
                cur_insn = cur_insn.next
        return mba, insns, mbbs

    def ql_deflat(self):
        if len(self.bb_mapping) == 0:
            self.ql_parse_blocks_for_deobf()
        if not self.qlinit:
            logging.info("Qiling should be setup firstly!")
            return
        self.mba, self.insns, self.mbbs = self._prepare_microcodes(maturity=3)
        logging.debug("Microcode generation done. Going to search path.")
        if not self._search_path():
            logging.info(f"Fail to search path. Please fire an issue to us at https://github.com/qilingframework/qiling with relevant logs!")
            return
        logging.debug("Real control flows search done. Going to patch codes.")
        self._patch_codes()
        logging.debug("Codes patched. Let's tell IDA to analyse the whole function again.")
        IDA.perform_analysis(self.deflat_func.start_ea, self.deflat_func.end_ea)
        del self.deflatqlemu
        self.deflatqlemu = None
        self.ks = None

    def _block_str(self, bb):
        if type(bb) is int:
            bb = self.bb_mapping[bb]
        return f"Block id: {bb.id}, start_address: {bb.start_ea:x}, end_address: {bb.end_ea:x}, type: {bb.type}"

    def ql_parse_blocks_for_deobf(self):
        cur_addr = IDA.get_current_address()
        flowchart = IDA.get_flowchart(cur_addr)
        self.deflat_func = IDA.get_function(cur_addr)
        self.bb_mapping = {bb.id:bb for bb in flowchart}
        if flowchart is None:
            return
        bb_count = {}
        for bb in flowchart:
            for succ in bb.succs():
                if succ.id not in bb_count:
                    bb_count[succ.id] = 0
                bb_count[succ.id] += 1
        max_ref_bb_id = None
        max_ref = 0
        for bb_id, ref in bb_count.items():
            if ref > max_ref:
                max_ref = ref
                max_ref_bb_id = bb_id
        self.pre_dispatcher = max_ref_bb_id
        try:
            self.dispatcher = list(self.bb_mapping[self.pre_dispatcher].succs())[0].id
            self.first_block = flowchart[0].id
        except IndexError:
            logging.error("Fail to get dispatcher and first_block.")
            return
        self.real_blocks = []
        self.fake_blocks = []
        self.retn_blocks = []
        for bb in flowchart:
            if self.pre_dispatcher in [b.id for b in bb.succs()] and IDA.get_instructions_count(bb.start_ea, bb.end_ea) > 1:
                self.real_blocks.append(bb.id)
            elif IDA.block_is_terminating(bb):
                self.retn_blocks.append(bb.id)
            elif bb.id != self.first_block and bb.id != self.pre_dispatcher and bb.id != self.dispatcher:
                self.fake_blocks.append(bb.id)
        for bbid in self.real_blocks:
            IDA.color_block(self.bb_mapping[bbid], Colors.Green.value)
        for bbid in self.fake_blocks:
            IDA.color_block(self.bb_mapping[bbid], Colors.Gray.value)
        for bbid in self.retn_blocks:
            IDA.color_block(self.bb_mapping[bbid], Colors.Pink.value)
        IDA.color_block(self.bb_mapping[self.dispatcher], Colors.Blue.value)
        IDA.color_block(self.bb_mapping[self.pre_dispatcher], Colors.Blue.value)
        IDA.color_block(self.bb_mapping[self.first_block], Colors.Beige.value)
        logging.info(f"First block: {self._block_str(self.first_block)}")
        logging.info(f"Dispatcher: {self._block_str(self.dispatcher)}")
        logging.info(f"Pre dispatcher: {self._block_str(self.pre_dispatcher)}")
        logging.info(f"Real blocks:")
        for s in map(self._block_str, self.real_blocks): logging.info(s)
        logging.info(f"Fake blocks:")
        for s in map(self._block_str, self.fake_blocks): logging.info(s)
        logging.info(f"Return blocks:")
        for s in map(self._block_str, self.retn_blocks): logging.info(s)
        logging.info(f"Auto analysis finished, please check whether the result is correct.")
        logging.info(f"You may change the property of each block manually if necessary.")

    # jb addr
    # jnb addr
    # <=>
    # jmp addr / nop
    # nop
    def _junk_useless_jcc(self, start, end):
        bs = IDA.get_bytes(start, end-start)
        patterns = []
        # Volume 2 Table A-2
        for opc in range(0x70, 0x7F, 2):
            # jb -> jnb
            r_opc = opc + 1
            patterns.append( re.compile(re.escape(bytes([opc])) + rb'.' + re.escape(bytes([r_opc])) + rb'.') )
            patterns.append( re.compile(re.escape(bytes([r_opc])) + rb'.' + re.escape(bytes([opc])) + rb'.') )

        for pattern in patterns:
            logging.debug(f"Start with pattern: {pattern}")
            tmpbs = bs
            result = re.search(pattern, tmpbs)
            while result is not None:
                l = result.span()[0]
                r = result.span()[1]
                l_offset = tmpbs[l+1]
                r_offset = tmpbs[l+3]
                tmpbs = tmpbs[r:]
                result = re.search(pattern, tmpbs)
                if l_offset == r_offset + 2:
                    logging.info(f"Get a junk jcc at [{hex(start+l)}, {hex(start+r)}] with offset {hex(l_offset)}.")
                    # If it jumps down, check if we can fill the codes with nops safely.
                    if l_offset <= 0x7F:
                        can_fill_with_nops = True
                        for addr in range(start+l, start + l +  2 + l_offset):
                            if len(IDA.get_xrefsto(addr, flags=ida_xref.XREF_FAR)) != 0:
                                logging.info(f"Find multiple Xrefs at {hex(addr)}, patch a jmp...")
                                can_fill_with_nops = False
                                break
                        if can_fill_with_nops:
                            logging.info(f"Fill NOPs from {hex(start+l)} to {hex(start + l + 2 + l_offset)}.")
                            IDA.fill_bytes(start+l, start + l + 2 + l_offset)
                            continue
                    # Or we simply patch a jmp.
                    logging.info(f"Patch a jmp at {start+l:x}.")
                    IDA.fill_bytes(start+l, start+r)
                    IDA.patch_bytes(start+l, b"\xeb" + bytes([l_offset]))


    # Remove junk code by fixed patterns.
    # If you find new patterns, please fire an issue or PR!
    # TODO: Find some samples to add more patterns.
    def ql_remove_junk_code_by_patterns(self):
        _, start, end = IDA.get_last_selection()
        self._junk_useless_jcc(start, end)


    def ql_nop_items_without_colors(self):
        _, start, end = IDA.get_last_selection()
        cur_start, cur_end = IDA.get_item(start)
        while cur_start < end:
            if not IDA.is_colored_item(cur_start):
                logging.info(f"Item at [{hex(cur_start)}, {hex(cur_end)}) doesn't have a color, nop it.")
                IDA.fill_bytes(cur_start, cur_end)
            cur_start, cur_end = IDA.get_item(cur_end)

    ### Hook

    def ql_step_hook(self, ql, addr, size):
        self.stepflag = not self.stepflag
        addr = addr - self.qlemu.baseaddr + get_imagebase()
        if self.stepflag:
            set_color(addr, CIC_ITEM, 0x00FFD700)
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, ql)
            self.qlemu.status = ql.save()
            ql.os.stop()
            self.qlemu.ql.hook_del(self.stephook)
            jumpto(addr)

    def ql_path_hook(self, ql, addr, size):
        addr = addr - self.qlemu.baseaddr + get_imagebase()
        set_color(addr, CIC_ITEM, 0x007FFFAA)
        bp_count = get_bpt_qty()
        bp_list = []
        if bp_count > 0:
            for num in range(0, bp_count):
                bp_list.append(get_bpt_ea(num))

            if addr in bp_list and (addr != self.lastaddr or self.is_change_addr>1):
                self.qlemu.status = ql.save()
                ql.os.stop()
                self.lastaddr = addr
                self.is_change_addr = -1
                jumpto(addr)

            self.is_change_addr += 1


    def ql_untill_hook(self, ql, addr, size):
        addr = addr - self.qlemu.baseaddr + get_imagebase()
        set_color(addr, CIC_ITEM, 0x00B3CBFF)

    ### User Scripts

    def ql_get_user_script(self, is_reload=False, is_start=False):
        def get_user_scripts_obj(scriptpath:str, classname:str, is_reload:bool):
            try:
                import os, sys
                import importlib

                modulepath,filename = os.path.split(scriptpath)
                logging.info(modulepath)
                logging.info(filename)
                scriptname,_ = os.path.splitext(filename)
                logging.info(scriptname)
                sys.path.append(modulepath)
                module = importlib.import_module(scriptname)

                if is_reload:
                    importlib.reload(module)
                cls = getattr(module, classname)
                return cls()
            except Exception as e:
                logging.exception("")
                return None

        self.userobj = get_user_scripts_obj(self.customscriptpath, 'QILING_IDA', is_reload)
        if self.userobj is not None:
            if is_reload and not is_start:
                logging.info('Custom user script is reloaded.')
            else:
                logging.info('Custom user script is loaded successfully.')
        else:
            logging.info('Custom user script not found.')

    ### Dialog

    def ql_set_rootfs(self):
        setupdlg = QlEmuSetupDialog()
        setupdlg.Compile()

        if setupdlg.Execute() != 1:
            return False

        rootfspath = setupdlg.path_name.value
        customscript = setupdlg.script_name.value
        parameter = setupdlg.parameters.value
        env = setupdlg.env_var.value

        if customscript != '':
            self.customscriptpath = customscript

        para_array = shlex.split(parameter)
        self.qlemu.path = [get_input_file_path()] + para_array
        logging.info(self.qlemu.path)

        if env != '':
            try:
                env_f = open(env)
                self.qlemu.env = load(env_f)
                env_f.close()
            except Exception as e:
                logging.error("Error parsing the env file!!")
                logging.error(e.message)
                return False

        if self.qlemu is not None:
            self.qlemu.rootfs = rootfspath
            return True
        return False

    ### Menu

    menuitems = []

    def ql_register_new_action(self, act_name, act_text, act_handler, shortcut, tooltip, icon):
        new_action = action_desc_t(
            act_name,       # The action name. This acts like an ID and must be unique
            act_text,       # The action text.
            act_handler,    # The action handler.
            shortcut,       # Optional: the action shortcut
            tooltip,        # Optional: the action tooltip (available in menus/toolbar)
            icon)           # Optional: the action icon (shows when in menus/toolbars)
        register_action(new_action)

    def ql_handle_menu_action(self, action):
        [x.handler() for x in self.menuitems if x.action == action]

    def ql_register_menu_actions(self):
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":start",             self.ql_start,                 "Setup",                      "Setup",                     None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":reloaduserscripts", self.ql_reload_user_script,      "Reload User Scripts",        "Reload User Scripts",       None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":runtohere",         self.ql_run_to_here,             "Execute Till",               "Execute Till",              None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":runselection",      self.ql_run_selection,           "Execute Selection",         "Execute Selection",              None,                True))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":runfromhere",       self.ql_continue,              "Continue",                   "Continue",                  None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":setpc",             self.ql_set_pc,                "Set PC",                     "Set PC",              None,                True))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":step",              self.ql_step,                  "Step",                       "Step (CTRL+SHIFT+F9)",      "CTRL+SHIFT+F9",        True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":changreg",          self.ql_chang_reg,              "Edit Register",              "Edit Register",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":reset",             self.ql_reset,                 "Restart",                    "Restart Qiling",            None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":close",             self.ql_close,                 "Close",                      "Close Qiling",              None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":reg view",          self.ql_show_reg_view,           "View Register",              "View Register",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":stack view",        self.ql_show_stack_view,         "View Stack",                 "View Stack",                None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":memory view",       self.ql_show_mem_view,           "View Memory",                "View Memory",               None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":save",              self.ql_save,                  "Save Snapshot",              "Save Snapshot",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":load",              self.ql_load,                  "Load Snapshot",              "Load Snapshot",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        if UseAsScript:
            self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":unload",            self.ql_unload_plugin,           "Unload Plugin",              "Unload Plugin",             None,                   False  ))
            self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":about",             self.ql_about,                 "About",                      "About",                     None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":checkupdate",       self.ql_check_update,           "Check Update",               "Check Update",              None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":parseblocks",       self.ql_parse_blocks_for_deobf,           "Auto Analysis For Deflat",               "Auto Analysis For Deflat",              None,                   True  ))
        # TODO: Dispatcher blocks should support assignment.
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":markreal",       self.ql_mark_real,           "Mark as Real Block",               "Mark as Real Block",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":markfake",       self.ql_mark_fake,           "Mark as Fake Block",               "Mark as Fake Block",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":markretn",       self.ql_mark_retn,           "Mark as Return Block",               "Mark as Return Block",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":deflat",       self.ql_deflat,           "Deflat",               "Deflat",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":removejunkcodebypatterns",       self.ql_remove_junk_code_by_patterns,           "Remove Junk Code by Patterns",               "Remove Junk Code by Patterns",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":nopinstructions",       self.ql_nop_items_without_colors,           "Nop Items without Color",               "Nop Items without Color",              None,                   True  ))

        for item in self.menuitems:
            if item.action == "-":
                continue
            self.ql_register_new_action(item.action, item.title, QlEmuMisc.menu_action_handler(self, item.action), item.shortcut, item.tooltip,  -1)

    def ql_unregister_menu_actions(self):
        for item in self.menuitems:
            unregister_action(item.action)

    def ql_attach_main_menu_actions(self):
        for item in self.menuitems:
            attach_action_to_menu("Edit/Plugins/" + self.plugin_name + "/" + item.title, item.action, SETMENU_APP)

    def ql_detach_main_menu_actions(self):
        for item in self.menuitems:
            detach_action_from_menu("Edit/Plugins/" + self.plugin_name + "/" + item.title, item.action)

    ### POPUP MENU

    def ql_hook_ui_actions(self):
        self.popup_menu_hook = self
        self.popup_menu_hook.hook()

    def ql_unhook_ui_actions(self):
        if self.popup_menu_hook != None:
            self.popup_menu_hook.unhook()

    # IDA 7.x

    def finish_populating_widget_popup(self, widget, popup_handle):
        if get_widget_type(widget) == BWN_DISASM:
            for item in self.menuitems:
                if item.popup:
                    attach_action_to_popup(widget, popup_handle, item.action, self.plugin_name + "/")

    ### Close View

    def ql_close_reg_view(self):
        self.qlemuregview = None

    def ql_close_stack_view(self):
        self.qlemustackview = None

    def ql_close_mem_view(self, viewid):
        del self.qlemumemview[viewid]

    def ql_close_all_views(self):
        if self.qlemuregview is not None:
            self.qlemuregview.Close()
        if self.qlemustackview is not None:
            self.qlemustackview.Close()

        for viewid in self.qlemumemview:
            self.qlemumemview[viewid].Close()
            self.qlemumemview = None

    def ql_update_views(self, addr, ql):
        if self.qlemuregview is not None:
            self.qlemuregview.SetReg(addr, ql)

        if self.qlemustackview is not None:
            self.qlemustackview.SetStack(self.qlemu.ql)

        for id in self.qlemumemview:
            self.qlemumemview[id].SetMem(self.qlemu.ql)


def PLUGIN_ENTRY():
    qlEmu = QlEmuPlugin()
    return qlEmu

if UseAsScript:
    if __name__ == "__main__":
        qlEmu = QlEmuPlugin()
        qlEmu.init()
        qlEmu.run()


repository: 0xba1100n/VxhunterIDAForkForPython3HighVersion
file: vxhunter_ida_py3.py
description: 修复停止维护的Vxhunter IDA插件项目在python3高些的版本下运行的若干问题
wanted_name: VxHunter
comment: VxHunter plugin for IDA Pro
created_at: 2024-08-25 12:47:43+00:00
pushed_at: 2024-08-26 03:26:32+00:00
url:
https://github.com/0xba1100n/VxhunterIDAForkForPython3HighVersion/blob/776dfcea9a73bac0e09623f2976f7aabc2c5075d/vxhunter_ida_py3.py

content:
# coding=utf-8
import logging
import struct
import idc
import ida_nalt
import ida_bytes
import ida_funcs
import ida_kernwin
import idaapi
import time


default_check_count = 100

known_address = [0x80002000, 0x10000, 0x1000, 0xf2003fe4, 0x100000, 0x107fe0]

function_name_key_words = ['bzero', 'usrInit', 'bfill']

# VxWorks 5.5
vx_5_sym_types = [
    # 0x00,      # Undefined Symbol
    # 0x01,      # Global (external)
    # 0x02,      # Local Absolute
    0x03,      # Global Absolute
    0x04,      # Local .text
    0x05,      # Global .text
    0x06,      # Local Data
    0x07,      # Global Data
    0x08,      # Local BSS
    0x09,      # Global BSS
    0x12,      # Local Common symbol
    0x13,      # Global Common symbol
    0x40,      # Local Symbols related to a PowerPC SDA section
    0x41,      # Global Symbols related to a PowerPC SDA section
    0x80,      # Local symbols related to a PowerPC SDA2 section
    0x81,      # Local symbols related to a PowerPC SDA2 section
]

# VxWorks 6.8
vx_6_sym_types = [
    # 0x00,  # Undefined Symbol
    # 0x01,  # Global (external)
    # 0x02,  # Local Absolute
    0x03,  # Global Absolute
    0x04,  # Local .text
    0x05,  # Global .text
    0x08,  # Local Data
    0x09,  # Global Data
    0x10,  # Local BSS
    0x11,  # Global BSS
    0x20,  # Local Common symbol
    0x21,  # Global Common symbol
    0x40,  # Local Symbols
    0x41,  # Global Symbols
]

need_create_function = [0x04, 0x05]


class VxTarget(object):
    def __init__(self, firmware, vx_version=5, is_big_endian=False, logger=None):
        """
        :param firmware: data of firmware
        :param vx_version: 5 = VxWorks 5.x; 6= VxWorks 6.x
        :param is_big_endian: True = big endian; False = little endian
        :param logger: logger for the target (default: None)
        """
        self.is_big_endian = is_big_endian
        self._vx_version = vx_version
        self.symbol_table_start = None
        self.symbol_table_end = None
        self._string_table = []
        self._symbol_table = []
        self.symbols = []
        self.load_address = None
        self._firmware = firmware
        self._has_symbol = None
        if self._vx_version == 5:
            self._symbol_interval = 16
        elif self._vx_version == 6:
            self._symbol_interval = 20
        self.start_time = None
        self._performance_status = []

        if logger is None:
            self.logger = logging.getLogger(__name__)
            self.logger.setLevel(logging.INFO)
            console_handler = logging.StreamHandler()
            console_format = logging.Formatter('[%(levelname)-8s][%(module)s.%(funcName)s] %(message)s')
            console_handler.setFormatter(console_format)
            self.logger.addHandler(console_handler)
        else:
            self.logger = logger
        self._firmware_info = {
            "load_address": -1,
            "has_symbol": False,
            "vx_version": self._vx_version,
            "big_endian": self.is_big_endian,
            "symbol_table_start": None,
            "symbol_table_end": None,
            "symbols": None
        }

        self.prepare()

    def reset_timer(self):
        self.start_time = time.time()

    def start_timer(self):
        if self.start_time:
            return False
        else:
            self.start_time = time.time()
            return self.start_time

    def get_timer(self):
        if self.start_time:
            return time.time() - self.start_time
        return False

    def prepare(self):
        """ Trying to find symbol from image.

        :return: True if found symbol, False otherwise.
        """
        self.find_symbol_table()
        if self._has_symbol is False:
            return None
        self.logger.debug("has_symbol: {}".format(self._has_symbol))
        self.get_symbol_table()

    def _check_vxworks_endian(self):
        """ Get image endian from image file.

        :return:
        """
        data1 = self._firmware[self.symbol_table_start + 4:self.symbol_table_start + 4 + self._symbol_interval]
        data2 = self._firmware[self.symbol_table_start + 4 + self._symbol_interval:self.symbol_table_start +
                                                                                   4 + self._symbol_interval * 2]
        if data1[0:2] == data2[0:2]:
            self.logger.info("VxWorks endian: Big endian.")
            self.is_big_endian = True
        elif data1[2:4] == data2[2:4]:
            self.logger.info("VxWorks endian: Little endian.")
            self.is_big_endian = False
        else:
            self.logger.info("VxWorks endian unknown. Assuming little endian.")
            self.is_big_endian = False

    def _check_symbol_format(self, offset):
        """ Check offset is symbol table.

        :param offset: offset from image.
        :return: True if offset is symbol table, False otherwise.
        """
        start_offset = offset
        end_offset = offset + self._symbol_interval * default_check_count
        if end_offset > len(self._firmware):
            return False

        check_data = self._firmware[start_offset:end_offset]
        is_big_endian = True
        is_little_endian = True
        # check symbol data match struct
        for i in range(default_check_count):
            check_data_1 = check_data[i * self._symbol_interval:(i + 1) * self._symbol_interval]
            if len(check_data_1) < self._symbol_interval:
                self.logger.debug("check_data_1 length is too small: {}".format(len(check_data_1)))
                break

            if self._check_symbol_format_simple(check_data_1) is False:
                return False

        if self._vx_version == 5:
            self.logger.debug("Check VxWorks 5 symbol format")
            # check is big endian
            for i in range(9):
                check_data_1 = check_data[4 + i * self._symbol_interval:6 + i * self._symbol_interval]
                data2 = check_data[4 + (i + 1) * self._symbol_interval:6 + (i + 1) * self._symbol_interval]
                if check_data_1 != data2:
                    self.logger.debug("VxWorks binary is not big endian.")
                    is_big_endian = False
                    break

            # check is little endian
            for i in range(9):
                check_data_1 = check_data[6 + i * self._symbol_interval:8 + i * self._symbol_interval]
                data2 = check_data[6 + (i + 1) * self._symbol_interval:8 + (i + 1) * self._symbol_interval]
                if check_data_1 != data2:
                    self.logger.debug("VxWorks binary is not little endian.")
                    is_little_endian = False
                    break

            if is_big_endian and is_little_endian:
                return False

            return is_big_endian ^ is_little_endian

        return True

    def _check_symbol_format_simple(self, data):
        """ Check single symbol format is correct.

        :param data: single symbol data.
        :return: True if data is symbol, False otherwise.
        """
        if self._vx_version == 5:
            # Check symbol type is valid
            sym_type = data[14]
            if sym_type not in vx_5_sym_types:
                return False

            # symbol should end with '\x00'
            if data[15] != 0:
                return False

            # Check symbol group is '\x00\x00'
            if data[12:14] != b'\x00\x00':
                return False

            # symbol_name point should not be zero
            if data[4:8] == b'\x00\x00\x00\x00':
                return False

            # symbol value point should not be zero
            if data[8:12] == b'\x00\x00\x00\x00':
                return False

            return True

        elif self._vx_version == 6:
            # Check symbol type is valid
            sym_type = data[18]
            if sym_type not in vx_6_sym_types:
                return False

            # symbol should end with '\x00'
            if data[19] != 0:
                return False

            # Check symbol group is '\x00\x00'
            if data[16:18] != b'\x00\x00':
                return False

            # symbol_name point should not be zero
            if data[4:8] == b'\x00\x00\x00\x00':
                return False

            # TODO: Need handle this problem
            # sometime data[8:12] will be '\x00\x00\x00\x00'
            # if data[8:12] == '\x00\x00\x00\x00':
            #     return False
            return True

        return False

    def find_symbol_table(self):
        """ Find symbol table from image.

        :return:
        """
        self.reset_timer()
        for offset in range(len(self._firmware) - self._symbol_interval):
            if self.symbol_table_start is None:
                # Get first data valid the symbol_format
                if not self._check_symbol_format_simple(self._firmware[offset: offset + self._symbol_interval]):
                    continue

                elif self._check_symbol_format(offset):
                    self.logger.info("symbol table start offset: {:010x}".format(offset))
                    self.symbol_table_start = offset
                    self._has_symbol = True
                    self._firmware_info["has_symbol"] = True
                    break
            else:
                break
        performance_data = "Find symbol table takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

        if self.symbol_table_start:
            self.reset_timer()
            for i in range(self.symbol_table_start, len(self._firmware), self._symbol_interval):
                check_data = self._firmware[i:i + self._symbol_interval]

                if len(check_data) < self._symbol_interval:
                    self.logger.debug("check_data length is too small: {}".format(check_data))
                    break

                if self._check_symbol_format_simple(check_data):
                    self.symbol_table_end = i + self._symbol_interval
                    self.logger.debug("self.symbol_table_end: {:010x}".format(self.symbol_table_end))

                else:
                    self.logger.info("Symbol table end offset: {:010x}".format(self.symbol_table_end))
                    break

        else:
            self.logger.error("Didn't find symbol table in this image")
            self._has_symbol = False

        performance_data = "Analyzer symbol table takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

    def get_symbol_table(self):
        """ get symbol table data.

        :return: True if get symbol table data successful, False otherwise.
        """
        if self.symbol_table_start and self.symbol_table_end:
            self._check_vxworks_endian()

        else:
            return False

        for i in range(self.symbol_table_start, self.symbol_table_end, self._symbol_interval):
            symbol_name_addr = self._firmware[i + 4:i + 8]
            symbol_dest_addr = self._firmware[i + 8:i + 12]
            symbol_flag = self._firmware[i + self._symbol_interval - 2]
            if self.is_big_endian:
                unpack_format = '>I'
            else:
                unpack_format = '<I'
            symbol_name_addr = int(struct.unpack(unpack_format, symbol_name_addr)[0])
            symbol_dest_addr = int(struct.unpack(unpack_format, symbol_dest_addr)[0])
            self.logger.debug("symbol_name_addr: {}; symbol_dest_addr: {}".format(symbol_name_addr, symbol_dest_addr))
            self._symbol_table.append({'symbol_name_addr': symbol_name_addr, 'symbol_name_length': None, 'symbol_dest_addr': symbol_dest_addr, 'symbol_flag': symbol_flag, 'offset': i})
        self.logger.debug("len(self._symbol_table): %s".format(len(self._symbol_table)))
        self._symbol_table = sorted(self._symbol_table, key=lambda x: x['symbol_name_addr'])
        for i in range(len(self._symbol_table) - 1):
            self._symbol_table[i]['symbol_name_length'] = self._symbol_table[i + 1]['symbol_name_addr'] - \
                                                          self._symbol_table[i]['symbol_name_addr']
        self.logger.debug("len(self._symbol_table): {}".format(len(self._symbol_table)))
        return True

    @staticmethod
    def _is_printable(c):
        """ Check Char is printable.

        :param c: char to check.
        :return: True if char is printable, False otherwise.
        """
        if type(c) is int:
            return 32 <= c <= 126
        elif type(c) is str:
            return 32 <= ord(c) <= 126

    def _is_func_name(self, string):
        """ Check target string is match function name format.

        :param string: string to check.
        :return: True if string is match function name format, False otherwise.
        """
        #
        bad_str = [b'\\', b'%', b'+', b',', b'&', b'/', b')', b'(', b'[', b']']
        # function name length should less than 512 byte
        if len(string) > 512:
            return False

        for data in bad_str:
            if data in string:
                return False

        for c in string:
            if self._is_printable(c) is False:
                return False
        return True

    def _get_prev_string_data(self, offset):
        """ Get previous string from giving offset.

        :param offset: offset of image.
        :return: string data, string start offset, string end offset.
        """
        while offset > 0:
            if self._firmware[offset] != 0:
                start_address = offset
                end_address = offset + 1
                while offset > 0:
                    if self._firmware[offset - 1] == 0:
                        start_address = offset
                        break
                    offset -= 1
                data = self._firmware[start_address:end_address]
                self.logger.debug("data: {}; start_address: {:010x}; end_address: {:010x}".format(data, start_address, end_address))
                return data, start_address, end_address
            else:
                offset -= 1
        self.logger.debug("Done looking for previous string data.")
        return None, None, None

    def _get_next_string_data(self, offset):
        """ Get next string from giving offset.

        :param offset: offset of image.
        :return: string data, string start offset, string end offset.
        """
        while offset < len(self._firmware):
            if self._firmware[offset] != 0:
                start_address = offset
                end_address = offset
                while offset <= len(self._firmware):
                    offset += 1
                    if self._firmware[offset] == 0:
                        end_address = offset
                        break
                data = self._firmware[start_address:end_address]
                return data, start_address, end_address
            else:
                offset += 1
        return None, None, None

    def find_string_table_by_key_function_index(self, key_offset):
        """ Find string table by VxWorks key function name offset in VxWorks image.

        :param key_offset: key function name offset in VxWorks image.
        :return:
        """
        self.logger.debug("Attempting to find string table by key function index with offset {:010x}".format(key_offset))
        temp_str_tab_data = []
        if len(self._symbol_table) > default_check_count:
            count = default_check_count
        else:
            count = len(self._symbol_table)
        start_offset = key_offset
        end_offset = key_offset
        self.logger.debug("Initializing with start_offset = end_offset = {:010x}".format(key_offset))

        while start_offset > 0:
            if self._is_printable(self._firmware[start_offset]) is True:
                # get string from offset
                string, start_address, end_address = self._get_prev_string_data(start_offset)
                self.logger.debug("string: {}; start_address: {:010x}; end_address: {:010x}".format(string, start_address, end_address))
                # check string is function name
                if self._is_func_name(string) is False:
                    if len(temp_str_tab_data) < count:
                        self.logger.error("Can't find any string table with key index: {}".format(string))
                        return None, None
                    else:
                        self.logger.info("Found string table start address at {:010x}".format(start_address))
                        break
                else:
                    temp_str_tab_data.append((string, start_address, end_address))

                # get previous string from offset
                prev_string, prev_start_address, prev_end_address = self._get_prev_string_data(start_address - 1)
                self.logger.debug("prev_string: {}, prev_start_address: {:010x}, prev_end_address: {:010x}".format(prev_string, prev_start_address, prev_end_address))
                if prev_start_address:
                    # strings interval should less than 4
                    if 4 < (start_address - prev_end_address):
                        if len(temp_str_tab_data) < count:
                            self.logger.error("Can't find any string table with key index: {}".format(string))
                            return None, None
                        else:
                            self.logger.info("found string table start address at {:010x}".format(start_address))
                            break
                    else:
                        start_offset = start_address - 1
                        self.logger.debug("start_offset: {}".format(start_offset))
                else:
                    break
            else:
                start_offset -= 1

        while end_offset < len(self._firmware):
            # find first printable char
            if self._is_printable(self._firmware[end_offset]) is True:
                # get string from offset
                string, start_address, end_address = self._get_next_string_data(end_offset)
                # check string is function name
                if self._is_func_name(string) is False:
                    if len(temp_str_tab_data) < count:
                        temp_str_tab_data = []
                        end_offset = end_address
                        continue
                    else:
                        self.logger.info("found string table end at {:010x}".format(end_address))
                        break

                else:
                    temp_str_tab_data.append((string, start_address, end_address))

                # get next string from offset
                next_string, next_start_address, next_end_address = self._get_next_string_data(end_address)
                if next_start_address:
                    # strings interval should less than 4
                    if 4 < (next_start_address - end_address):
                        if len(temp_str_tab_data) < count:
                            self.logger.error("Can't find any string table with key index.")
                            return None, None
                        else:
                            self.logger.info("Found string table end at {:010x}".format(end_address))
                            break
                    else:
                        end_offset = end_address
            else:
                end_offset += 1

        temp_str_tab_data = sorted(temp_str_tab_data, key=lambda x: (x[1]))
        table_start_offset = temp_str_tab_data[0][1]
        table_end_offset = temp_str_tab_data[-1][2]
        self.logger.info("Found a string table at: {:010x} to {:010x}".format(table_start_offset, table_end_offset))
        return table_start_offset, table_end_offset

    def get_string_table(self, str_start_address, str_end_address):
        """ Get string table data from VxWorks image with string table start and end offset.

        :param str_start_address: string table start address.
        :param str_end_address: string table end address.
        :return:
        """
        self._string_table = []
        offset = str_start_address
        address = offset
        str_tab_data = []
        while offset <= str_end_address:
            if self._firmware[offset] == 0:
                while offset <= str_end_address:
                    offset += 1
                    if self._firmware[offset] != 0:
                        next_address = offset
                        string = self._firmware[address:next_address]
                        length = next_address - address
                        str_tab_data.append({'address': address, 'string': string, 'length': length})
                        offset = next_address
                        address = next_address
                        break
            else:
                offset += 1
        self._string_table = str_tab_data

    def _check_fix(self, func_index, str_index):
        """

        :param func_index:
        :param str_index:
        :return:
        """
        try:
            fault_count = 0
            self.logger.debug("Symbol table's first element: {}".format(self._symbol_table[0]))
            if len(self._symbol_table) <= default_check_count:
                count = len(self._symbol_table)
                self.logger.debug("Length of symbol table, {}, is less than default. Setting iteration count to actual length of table, {}.".format(len(self._symbol_table), count))
            else:
                count = default_check_count
                self.logger.debug("Length of symbol table, {}, is greater than default. Setting iteration count to default, {}.".format(len(self._symbol_table), count))
            for i in range(count):

                if (func_index >= len(self._symbol_table)) or (str_index >= len(self._string_table)):
                    self.logger.debug("_check_fix False: func_index greater than length of _symbol_table, or str_index greater than length of _string_table.")
                    return False
                self.logger.debug("str_index: {}; _string_table[str_index]: {}".format(str_index, self._string_table[str_index]))
                self.logger.debug("func_index: {}; _symbol_table[func_index]: {}".format(func_index, self._symbol_table[func_index]))
                if i == count - 1:
                    if fault_count < 10:
                        self.logger.debug("_check_fix True")
                        return True
                    else:
                        self.logger.debug("_check_fix False: Too many faults.")
                        return False

                if self._string_table[str_index]['length'] == self._symbol_table[func_index]['symbol_name_length']:
                    func_index += 1
                    str_index += 1
                    self.logger.debug("_check_fix continue")

                elif self._symbol_table[func_index]['symbol_name_length'] < self._string_table[str_index]['length']:
                    # Sometime Symbol name might point to mid of string.
                    fault_count += 1
                    func_index += 1
                else:
                    self.logger.debug("_check_fix False: symbol_name_length from func_index larger than length from str_index.")
                    return False
        except Exception as e:
            self.logger.debug(e)
            raise

    def find_loading_address(self):
        """ Find VxWorks image load address by automatic analysis.

        :return: Load address if found, None otherwise.
        """

        # Search function keyword in firmware to locate the function string tables.
        self.reset_timer()
        if self._has_symbol is False:
            return None

        for key_word in function_name_key_words:
            prefix_keyword = b'\x00_' + key_word + b'\x00'
            key_word = b'\x00' + key_word + b'\x00'
            if key_word in self._firmware is False and prefix_keyword in self._firmware is False:
                self.logger.info("Firmware does not contain a function named {}".format(key_word))
                return None

        performance_data = "Search function keyword in firmware takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

        # Search function keyword in firmware to locate the function string tables.
        self.reset_timer()
        try:
            key_function_index = self._firmware.index(b'\x00' + function_name_key_words[0] + b'\x00')
            self.logger.debug("key_function_index: {}".format(key_function_index))
        except Exception as err:
            # Handler _ prefix symbols
            self.logger.debug("Exception occurred while loading key_function_index: {}. Continuing...".format(err))
            key_function_index = self._firmware.index(b'\x00_' + function_name_key_words[0] + b'\x00')
            self.logger.debug("key_function_index: {}".format(key_function_index))

        str_start_address, str_end_address = self.find_string_table_by_key_function_index(key_function_index)
        self.get_string_table(str_start_address, str_end_address)
        performance_data = "Get function string table takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

        # TODO: Need improve performance
        self.reset_timer()
        self.logger.info("Starting loading address analysis")
        for str_index in range(len(self._string_table)):
            for func_index in range(len(self._symbol_table)):
                self.logger.debug("self._string_table[str_index]['length']: {}".format(self._string_table[str_index]['length']))
                self.logger.debug("self._symbol_table[func_index]['symbol_name_length']: {}".format(self._symbol_table[func_index]['symbol_name_length']))
                if self._string_table[str_index]['length'] == self._symbol_table[func_index]['symbol_name_length']:
                    if self._check_fix(func_index, str_index) is True:
                        self.logger.debug("self._symbol_table[func_index]['symbol_name_addr']: {}".format(self._symbol_table[func_index]['symbol_name_addr']))
                        self.logger.debug("self._string_table[str_index]['address']: %s" % self._string_table[str_index]['address'])
                        self.load_address = self._symbol_table[func_index]['symbol_name_addr'] - \
                                            self._string_table[str_index]['address']
                        self._firmware_info["load_address"] = self.load_address
                        self.logger.info('load address is {:010x}'.format(self.load_address))
                        performance_data = "Analyze loading address takes {:.3f} seconds".format(
                            self.get_timer())
                        self._performance_status.append(performance_data)
                        self.logger.debug(performance_data)
                        return self.load_address
                else:
                    continue

        self.logger.error("We didn't find load address in this firmware, sorry!")
        performance_data = "Analyze loading address takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

    def _check_load_address(self, address):
        """

        :param address:
        :return:
        """
        if not self._has_symbol:
            return False
        if len(self._symbol_table) > default_check_count:
            self.logger.debug("Length of symbol table greater than default. Setting iteration count to default of {}.".format(default_check_count))
            count = default_check_count
        else:
            count = len(self._symbol_table)
        self.logger.debug("symbol_table length is {}".format(count))
        for i in range(count):
            offset = self._symbol_table[i]['symbol_name_addr'] - address
            if offset <= 0:
                return False
            # TODO: Need improve, currently use string point to check.
            string, str_start_address, str_end_address = self._get_next_string_data(offset)
            if str_start_address != offset:
                self.logger.info("String {} at offset {} didn't match symbol table.".format(string, offset))
                return False
        self.logger.info('Load address is {:010x}'.format(address))
        return True

    def quick_test(self):
        """ Using known load address list to match VxWorks image.

        :return: Load address if match known address, None otherwise.
        """
        if self._has_symbol is False:
            return None
        self.logger.debug("has_symbol: {}".format(self._has_symbol))
        for address in known_address:
            if self._check_load_address(address):
                self.load_address = address
                self._firmware_info["load_address"] = self.load_address
                return self.load_address
            else:
                self.logger.info('Load address is not {:010x}'.format(address))

    def cleanup(self):
        """ Clean up variables.

        :return:
        """
        self.is_big_endian = False
        self.symbol_table_start = None
        self.symbol_table_end = None
        self._string_table = []
        self._symbol_table = []
        self.load_address = None
        self._has_symbol = None

    def get_string_from_firmware_by_offset(self, string_offset):
        symbol_name = ""
        while True:
            if self._firmware[string_offset] != 0x00:
                symbol_name += chr(self._firmware[string_offset])
                string_offset += 1

            else:
                break

        return symbol_name

    def get_symbols(self):
        self.symbols = []
        if self.load_address:
            for symbol in self._symbol_table:
                symbol_name_addr = symbol["symbol_name_addr"]
                symbol_dest_addr = symbol["symbol_dest_addr"]
                symbol_flag = symbol["symbol_flag"]
                symbol_name_firmware_addr = symbol_name_addr - self.load_address
                symbol_name = self.get_string_from_firmware_by_offset(symbol_name_firmware_addr)
                self.symbols.append({
                    "symbol_name": symbol_name,
                    "symbol_name_addr": symbol_name_addr,
                    "symbol_dest_addr": symbol_dest_addr,
                    "symbol_flag": symbol_flag
                })
            return self.symbols

        else:
            return None

    def get_firmware_info(self):
        symbols = self.get_symbols()
        self._firmware_info["symbols"] = symbols
        return self._firmware_info

    def get_performance_status(self):
        return self._performance_status

def demangle_function(demangle_string):
    function_return = None
    function_parameters = None
    function_name_end = len(demangle_string)

    # get parameters
    index = len(demangle_string) - 1
    if demangle_string[-1] == ')':
        # have parameters
        parentheses_count = 0
        while index >= 0:
            if demangle_string[index] == ')':
                parentheses_count += 1

            elif demangle_string[index] == '(':
                parentheses_count -= 1

            index -= 1

            if parentheses_count == 0:
                break

        function_parameters = demangle_string[index + 2:-1]
        function_name_end = index

    # get function name
    while index >= 0:
        if demangle_string[index] == ' ':
            break
        else:
            index -= 1
    function_name_start = index
    function_name = demangle_string[function_name_start + 1:function_name_end + 1]

    # get function return
    function_return = demangle_string[:function_name_start]
    return function_return, function_name, function_parameters


def is_vx_symbol_file(file_data, is_big_endian=True):
    # Check key function names
    for key_function in function_name_key_words:
        if key_function not in file_data:
            print("key function not found")
            return False

    if is_big_endian:
        return struct.unpack('>I', file_data[:4])[0] == len(file_data)

    else:
        return struct.unpack('<I', file_data[:4])[0] == len(file_data)


# --------------------------------------------------------------------------
# Plugin
# --------------------------------------------------------------------------
class AutoFixIDBForm(idaapi.Form):
    def __init__(self):
        self.invert = False
        self.vx_version = 5
        super(AutoFixIDBForm, self).__init__(
            r"""BUTTON YES* Start analyze

VxHunter Auto Fix IDB will fix IDB with VxWorks symbol table.

Please choose VxWorks main version
{FormChangeCb}
<VxWorks Main Version     :{c_vxversion}>
            """, {
                'FormChangeCb': self.FormChangeCb(self.OnFormChange),
                'c_vxversion': self.DropdownListControl(
                    items=("5", "6"),
                    readonly=True,
                    selval=0),

            }
        )

        self.Compile()

    def OnFormChange(self, fid):
        if fid == -2:
            self.vx_version = (5, 6)[self.GetControlValue(self.c_vxversion)]
            return 1


class FixCodeForm(idaapi.Form):
    def __init__(self):
        self.invert = False
        self.start_address = 0
        self.end_address = 0
        super(FixCodeForm, self).__init__(
            r"""BUTTON YES* Fix Code from image

VxHunter fix code will make all data as code from start_address to end_address. 

Please input start address and end address
{FormChangeCb}
<Start address     :{c_StartAddress}>
<End address     :{c_EndAddress}>
            """, {
                'FormChangeCb': self.FormChangeCb(self.OnFormChange),
                'c_StartAddress': self.NumericInput(value=self.start_address, swidth=40, tp=self.FT_ADDR),
                'c_EndAddress': self.NumericInput(value=self.end_address, swidth=40, tp=self.FT_ADDR),
            }
        )
        self.Compile()

    def OnFormChange(self, fid):
        if fid == -2:
            self.start_address = self.GetControlValue(self.c_StartAddress)
            self.end_address = self.GetControlValue(self.c_EndAddress)
            return 1


class FixAsciiForm(idaapi.Form):
    def __init__(self):
        self.invert = False
        self.string_address = 0
        super(FixAsciiForm, self).__init__(
            r"""BUTTON YES* Fix Code from image

VxHunter fix ascii string. 

Please input string table start address.
{FormChangeCb}
<String Address     :{c_Address}>
            """, {
                'FormChangeCb': self.FormChangeCb(self.OnFormChange),
                'c_Address': self.NumericInput(value=self.string_address, swidth=40, tp=self.FT_ADDR),
            }
        )
        self.Compile()

    def OnFormChange(self, fid):
        if fid == -2:
            self.string_address = self.GetControlValue(self.c_Address)
            return 1


class VxHunter_Plugin_t(idaapi.plugin_t):
    comment = "VxHunter plugin for IDA Pro"
    help = ""
    wanted_name = "VxHunter"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        # register popup menu handlers
        try:
            # Register Auto Fix IDB handler
            VxHunterMCFixIDB.register(self, "Auto Fix IDB With symbol table")
            # Register Fix Code handler
            VxHunterMCFixCode.register(self, "Fix Code from start address to end address")
            # Register Fix Ascii handler
            VxHunterMCFixAscii.register(self, "Fix Ascii string table with giving address")
            # Register Load symbol file handler
            VxHunterMCLoadSymbolFile.register(self, "Load VxWorks symbol file")

        except Exception as err:
            print("Got Error!!!: %s" % err)

        # setup popup menu
        if idaapi.IDA_SDK_VERSION >= 700:
            # Add menu IDA >= 7.0
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCFixIDB.get_name(), idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCFixCode.get_name(), idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCFixAscii.get_name(), idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCLoadSymbolFile.get_name(), idaapi.SETMENU_APP)
        else:
            # add Vxhunter menu
            menu = idaapi.add_menu_item("Edit/VxHunter/", "Auto Fix IDB1", "", 1, self.handler_auto_fix_idb, None)
            if menu is not None:
                pass

        print("=" * 80)
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    # null handler
    def menu_null(self):
        pass

    def handler_auto_fix_idb(self):
        form = AutoFixIDBForm()
        ok = form.Execute()
        if ok == 1:
            vx_version = int(form.vx_version)
            print("vx_version:%s" % vx_version)
            firmware_path = idaapi.get_input_file_path()
            firmware = open(firmware_path, 'rb').read()
            target = VxTarget(firmware=firmware, vx_version=vx_version)
            # target.logger.setLevel(logging.DEBUG)
            target.quick_test()

            if target.load_address:
                print("Load Address is:%s" % target.load_address)
            else:
                target.find_loading_address()
                if target.load_address:
                    print("Load Address is:%s" % target.load_address)
            if not target.load_address:
                return
            symbol_table_start = target.symbol_table_start
            symbol_table_end = target.symbol_table_end
            load_address = target.load_address
            self.fix_vxworks_idb(load_address, vx_version, symbol_table_start, symbol_table_end)
        form.Free()

    def handler_fix_code(self):
        form = FixCodeForm()
        ok = form.Execute()
        if ok == 1:
            start_address = int(form.start_address)
            end_address = int(form.end_address)
            self.fix_code(start_address, end_address)
        form.Free()

    def handler_fix_ascii(self):
        form = FixAsciiForm()
        ok = form.Execute()
        if ok == 1:
            string_address = int(form.string_address)
            self.fix_ascii(string_address)

        form.Free()

    def handler_load_symbol_file(self):
        self.load_symbol_file()

    @staticmethod
    def fix_vxworks_idb(load_address, vx_version, symbol_table_start, symbol_table_end):
        current_image_base = idaapi.get_imagebase()
        symbol_interval = 16
        if vx_version == 6:
            symbol_interval = 20
        symbol_table_start += load_address
        symbol_table_end += load_address
        ea = symbol_table_start
        shift_address = load_address - current_image_base
        while shift_address >= 0x70000000:
            idaapi.rebase_program(0x70000000, 0x0008)
            shift_address -= 0x70000000
        idaapi.rebase_program(shift_address, 0x0008)
        while ea < symbol_table_end:
            # for VxWorks 6 unknown symbol format
            if idc.get_wide_byte(ea + symbol_table_end - 2) == 3:
                ea += symbol_interval
                continue
            offset = 4
            idc.create_strlit(idc.get_wide_dword(ea + offset), idc.BADADDR)
            sName = idc.get_strlit_contents(idc.get_wide_dword(ea + offset), -1, ida_nalt.STRTYPE_C)
            print("Found %s in symbol table" % sName)
            if sName:
                sName_dst = idc.get_wide_dword(ea + offset + 4)
                if vx_version == 6:
                    sName_type = idc.get_wide_byte(ea + 0x12)
                else:
                    sName_type = idc.get_wide_byte(ea + 0x0e)

                idc.set_name(sName_dst, sName.decode('utf-8'), idc.SN_CHECK)
                if sName_type in need_create_function:
                    print("Start fix Function %s at %s" % (sName, hex(sName_dst)))
                    idc.create_insn(sName_dst)  # might not need
                    ida_funcs.add_func(sName_dst, idc.BADADDR)
            ea += symbol_interval
        print("Fix function by symbol table finish.")
        print("Start IDA auto analysis, depending on the size of the firmware this might take a few minutes.")
        idaapi.auto_wait()

    @staticmethod
    def fix_code(start_address, end_address):
        # Todo: There might be some data in the range of codes.
        offset = start_address
        while offset <= end_address:
            offset = ida_bytes.next_addr(offset)
            flags = ida_bytes.get_full_flags(offset)
            if not ida_bytes.is_code(flags):
                # Todo: Check should use MakeCode or MakeFunction
                # idc.create_insn(offset)
                ida_funcs.add_func(offset)

    @staticmethod
    def get_prev_ascii_string_address(address):
        """

        :param address: must be current ascii string start address.
        :return:
        """
        prev_string_start_address = address
        # string table interval should less than 5 bytes.
        if idc.get_wide_dword(address - 5) == 0:
            return None
        else:
            prev_string_start_address -= 5
            # TODO: Need handle short string.
            while idaapi.get_byte(prev_string_start_address) != 0:
                prev_string_start_address -= 1
            return prev_string_start_address + 1

    @staticmethod
    def get_next_ascii_string_address(address):
        """

        :param address: must be current ascii string start address.
        :return:
        """
        next_string_start_address = address
        # find current string end address
        while idaapi.get_byte(next_string_start_address) != 0:
            next_string_start_address += 1

        # string table interval should less than 5 bytes.
        # TODO: need handle short string.
        if idc.get_wide_dword(next_string_start_address + 1) == 0:
            return None

        while idaapi.get_byte(next_string_start_address) == 0:
            next_string_start_address += 1

        return next_string_start_address

    def get_string_table_start_address(self, address):
        string_table_start_address = address
        # find current string start address
        while idaapi.get_byte(string_table_start_address - 1) != 0:
            string_table_start_address -= 1

        while self.get_prev_ascii_string_address(string_table_start_address):
            string_table_start_address = self.get_prev_ascii_string_address(string_table_start_address)

        return string_table_start_address

    def fix_ascii(self, address):
        string_table_start_address = self.get_string_table_start_address(address)
        string_address = string_table_start_address
        while True:
            if string_address:
                print("Start Make string at address: %s" % hex(string_address))
                idc.create_strlit(string_address, idc.BADADDR)
                string_address = self.get_next_ascii_string_address(string_address)
            else:
                break

    def load_symbols(self, file_data, is_big_endian=True):
        symbol_list = []
        if is_big_endian:
            unpack_format = '>I'
        else:
            unpack_format = '<I'

        symbol_count = struct.unpack(unpack_format, file_data[4:8])[0]
        print("symbol_count: %s" % symbol_count)
        symbol_offset = 8
        string_table_offset = 8 + 8 * symbol_count
        print("string_table_offset: %s" % string_table_offset)
        # get symbols
        for i in range(symbol_count):
            offset = i * 8
            symbol_data = file_data[symbol_offset + offset:symbol_offset + offset + 8]
            flag = ord(symbol_data[0])
            string_offset = struct.unpack(unpack_format, '\x00' + symbol_data[1:4])[0]
            string_offset += string_table_offset
            print("string_offset: %s" % string_offset)
            symbol_name = ""
            while True:
                if file_data[string_offset] != '\x00':
                    symbol_name += file_data[string_offset]
                    string_offset += 1

                else:
                    break
            print("symbol_name: %s" % symbol_name)
            symbol_address = struct.unpack(unpack_format, symbol_data[-4:])[0]
            symbol_list.append([flag, symbol_name, symbol_address])
            # Find TP-Link device loading address with symbols
            if "wrs_kernel_text_start" in symbol_name:
                load_address = symbol_address
                current_image_base = idaapi.get_imagebase()
                shift_address = load_address - current_image_base
                while shift_address >= 0x70000000:
                    idaapi.rebase_program(0x70000000, 0x0008)
                    shift_address -= 0x70000000
                idaapi.rebase_program(shift_address, 0x0008)

        # load symbols
        for symbol_data in symbol_list:
            flag, symbol_name, symbol_address = symbol_data
            idc.set_name(symbol_address, symbol_name.decode('utf-8'), idc.SN_CHECK)
            if flag == 0x54:
                if symbol_name:
                    print("Start fix Function %s at %s" % (symbol_name, hex(symbol_address)))
                    idc.create_insn(symbol_address)  # might not need
                    ida_funcs.add_func(symbol_address, idc.BADADDR)

    def load_symbol_file(self):
        symbol_file_path = ida_kernwin.ask_file(0, "*", "Please chose the VxWorks symbol file")
        print("symbol_file_path: {}".format(symbol_file_path))
        symbol_file_data = open(symbol_file_path, 'rb').read()
        if is_vx_symbol_file(symbol_file_data):
            self.load_symbols(symbol_file_data)
            idaapi.auto_wait()

        else:
            return

    def run(self, arg):
        self.handler_auto_fix_idb()


try:
    class VxHunterMenuContext(idaapi.action_handler_t):

        @classmethod
        def get_name(self):
            return self.__name__

        @classmethod
        def get_label(self):
            return self.label

        @classmethod
        def register(self, plugin, label):
            self.plugin = plugin
            self.label = label
            instance = self()
            return idaapi.register_action(idaapi.action_desc_t(
                self.get_name(),  # Name. Acts as an ID. Must be unique.
                instance.get_label(),  # Label. That's what users see.
                instance  # Handler. Called when activated, and for updating
            ))

        @classmethod
        def unregister(self):
            """Unregister the action.
            After unregistering the class cannot be used.
            """
            idaapi.unregister_action(self.get_name())

        @classmethod
        def activate(self, ctx):
            # dummy method
            return 1

        @classmethod
        def update(self, ctx):
            try:
                if ctx.form_type == idaapi.BWN_DISASM:
                    return idaapi.AST_ENABLE_FOR_FORM
                else:
                    return idaapi.AST_DISABLE_FOR_FORM
            except:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS

    # context menu for Fix idb
    class VxHunterMCFixIDB(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_auto_fix_idb()
            return 1

    class VxHunterMCFixCode(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_fix_code()
            return 1

    class VxHunterMCFixAscii(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_fix_ascii()
            return 1

    class VxHunterMCLoadSymbolFile(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_load_symbol_file()
            return 1


except Exception as err:
    # TODO: Add some handle later.
    print("err: {}".format(err))


# register IDA plugin
def PLUGIN_ENTRY():
    return VxHunter_Plugin_t()


repository: Goatman13/IDA-Mips-saved-registers-renamer
file: mips_stack.py
description: Renames variables for MIPS registers saved on the stack
wanted_name: Mips saved registers renamer
comment: Renames MIPS registers saved on the stack
created_at: 2022-07-11 11:38:46+00:00
pushed_at: 2022-07-11 11:49:41+00:00
url:
https://github.com/Goatman13/IDA-Mips-saved-registers-renamer/blob/ca7552116c688d71ea0c8197981be81f83a3c460/mips_stack.py

content:
# IDA plugin to name stack variables that are simply used to store register
# values until a function returns ($ra, $s0-$s7, $fp, $gp).
#
# Invoke by going to Plugins->Mips saved registers renamer,
# or by pressing Alt + Shift + 2.
#
# Craig Heffner
# Tactical Network Solutions

import idaapi
import idautils
import idc
import ida_auto

class NameMIPSSavedRegisters(object):
	INSIZE = 4
	SEARCH_DEPTH = 35

	ARCH = {
			'savedregs': ['$s0', '$s1', '$s2', '$s3', '$s4', '$s5', '$s6',
						  '$s7', '$fp', '$gp', '$ra', '$f20', '$f21', '$f22',
						  '$f23', '$f24', '$f25', '$f26', '$f27', '$f28', '$f29', '$f30', '$f31'],
	}

	def __init__(self):
		print("Naming saved register locations...", end=' ')

		for ea in idautils.Functions():
			mea = ea
			named_regs = []
			last_iteration = False

			while mea < (ea + (self.INSIZE * self.SEARCH_DEPTH)):
				mnem = idc.print_insn_mnem(mea)

				if mnem in ['sw', 'sd', 'sq', 'swc1', 'sdc1']:
					reg = idc.print_operand(mea, 0)
					dst = idc.print_operand(mea, 1)

					if reg in self.ARCH['savedregs'] and \
							reg not in named_regs and \
							dst.endswith('($sp)') and 'var_' in dst:
						split_string = 'var_'
						stack_position = "[sp-%d]"
						if '_s' in dst:
							split_string += 's'
							stack_position = "[sp+%d]"

						offset = int(dst.split(split_string)[1].split('(')[0],
									 16)
						idc.define_local_var(
							ea, idc.find_func_end(ea),
							stack_position % offset, "saved_%s" % reg[1:])
						named_regs.append(reg)

				if last_iteration:
					break
				elif mnem.startswith('j') or mnem.startswith('b'):
					last_iteration = True

				mea += self.INSIZE

		print("Mips saved registers renamer done.")

def name_saved_registers():
	NameMIPSSavedRegisters()

class mips_saved_registers_t(idaapi.plugin_t):
	flags = 0
	comment = "Renames MIPS registers saved on the stack"
	help = "Renames MIPS registers saved on the stack"
	wanted_name = "Mips saved registers renamer"
	wanted_hotkey = "Alt-Shift-2"
	
	def init(self):
		if (idaapi.ph.id == idaapi.PLFM_MIPS):
			idaapi.msg("Mips saved registers renamer loaded.\n")
			return idaapi.PLUGIN_KEEP
			
		return idaapi.PLUGIN_SKIP
		
	def run(self, arg):
		if ida_auto.auto_is_ok() == 1:
			name_saved_registers()
		else:
			idaapi.warning("Aborted! Run plugin again after initial autoanalyze finished.")
			
	def term(self):
		pass

def PLUGIN_ENTRY():
	return mips_saved_registers_t()


repository: zengfr/ida-pro-idb-database
file: demo/4/cyclomatic_complexity.py
description: many of ida pro idb database for game hacking and reverse analysis engineering.good luck.
wanted_name: Cyclomatic Complexity
comment: Cyclomatic Complexity
created_at: 2022-11-20 02:35:57+00:00
pushed_at: 2024-11-02 01:54:30+00:00
url:
https://github.com/zengfr/ida-pro-idb-database/blob/b2a948cf8f60865932cb2eb76bb73decc4b5b34c/demo/4/cyclomatic_complexity.py

content:
#!/usr/bin/python

from idc import *
from idaapi import *
from idautils import *
from sets import Set

def cyclomatic_complexity(function_ea):
  """
  Cyclomatic Complexity function - by Ero Carrera  

  Calculate the cyclomatic complexity measure for a function.
  Given the starting address of a function, it will find all the basic block's boundaries and edges
  between them and will return the cyclomatic complexity, defined as:
    CC = Edges - Nodes + 2
  """

  f_start = function_ea
  f_end = FindFuncEnd(function_ea)
  edges = Set()
  boundaries = Set((f_start,))

  for head in Heads(f_start, f_end):
    if isCode(GetFlags(head)):
      refs = CodeRefsFrom(head, 0)
      refs = Set(filter(lambda x: x>=f_start and x<=f_end, refs))
      
      if refs:
        next_head = NextHead(head, f_end)
        if isFlow(GetFlags(next_head)):
          refs.add(next_head)
        boundaries.union_update(refs)

        for r in refs:
          if isFlow(GetFlags(r)):
            edges.add((PrevHead(r, f_start), r))
          edges.add((head, r))

  return len(edges) - len(boundaries) + 2

class CyclomaticComplexityChoose(Choose2):
  def __init__(self, title):
    Choose2.__init__(self, title, [ 
        ["Address", 8 | Choose2.CHCOL_HEX], 
        ["Name", 30 | Choose2.CHCOL_PLAIN], 
        ["Complexity", 6 | Choose2.CHCOL_DEC],
        ["Library func", 6 | Choose2.CHCOL_PLAIN] ])

    self.title  = title
    self.colors = (0xF3E2A9, 0xCEF6E3, 0xFA2ECC, 0xFFFF)
    self.items  = []
    self.icon   = 41
    self.PopulateItems()

  def OnClose(self):
    return True

  def OnSelectLine(self, n):
    item = self.items[int(n)]
    jumpto(int(item[0], 16))

  def OnGetLine(self, index):
    return self.items[index]

  def OnGetSize(self):
    return len(self.items)

  def OnDeleteLine(self, n):
    del self.items[n]
    return n

  def OnGetLineAttr(self,n):
    complexity = int(self.items[n][2])
    color = self.colors[0]
    if 10 < complexity <= 20:
      color = self.colors[1]
    elif 20 < complexity <= 50:
      color = self.colors[2]
    elif complexity > 50:
      color = self.colors[3]
    return [color, 0]

  def OnRefresh(self, n):
    return n

  def OnCommand(self, n, cmd_id):
    if cmd_id == self.cmd_exc_lib_funcs:
      self.exclude_lib_funcs()
    return n

  def exclude_lib_funcs(self):
    if not len(self.items):
      return False 
    self.items = [i for i in self.items if i[3] != 'True']
    return True

  def show(self):
    n_functions = len(list(Functions()))
    if n_functions > 0:      
      b = self.Show()
      if b == 0:
        self.cmd_exc_lib_funcs = self.AddCommand("Exclude library functions")
        return True
    else:
      warning("IDA has not identified functions.")
      return False

  def PopulateItems(self):
    for function_ea in Functions():
      self.items.append([
        "%08x" % function_ea, 
        GetFunctionName(function_ea), 
        "%d" % cyclomatic_complexity(function_ea),
        "%s" % ((GetFunctionFlags(function_ea) & FUNC_LIB) != 0)
      ])

def show_choose():
  choose = CyclomaticComplexityChoose("Cyclomatic complexity")
  choose.show()

class CyclomaticComplexity_t(plugin_t):
  flags = PLUGIN_UNL
  comment = "Cyclomatic Complexity"
  help = ""
  wanted_name = "Cyclomatic Complexity"
  wanted_hotkey = ""

  def init(self):
    self.icon_id = 0
    return PLUGIN_OK

  def run(self, arg=0):
    show_choose()

  def term(self):
    pass

def PLUGIN_ENTRY():
  return CyclomaticComplexity_t()

if __name__ == '__main__':
  show_choose()
  


repository: tmr232/ida-plugins
file: autostruct/autostruct.py
description: IDA Pro Plugins
wanted_name: AutoStruct
comment: AutoStruct struct creator
created_at: 2015-05-27 12:35:29+00:00
pushed_at: 2015-05-31 13:05:09+00:00
url:
https://github.com/tmr232/ida-plugins/blob/c333ebecc9b3d641cdfdc9ab0d66f73e0b551fc3/autostruct/autostruct.py

content:
import idaapi
import idc
import sark


class AutoStruct(idaapi.plugin_t):
    flags = 0
    comment = "AutoStruct struct creator"
    help = "Automagically Create and Apply Structs"
    wanted_name = "AutoStruct"
    wanted_hotkey = "Shift+T"

    def init(self):
        self._prev_struct_name = ""
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        start, end = sark.get_selection()

        struct_name = idc.AskStr(self._prev_struct_name, "Struct Name")
        if not struct_name:
            return
        self._prev_struct_name = struct_name

        common_reg = sark.structure.get_common_register(start, end)
        reg_name = idc.AskStr(common_reg, "Register")
        if not reg_name:
            return

        offsets, operands = sark.structure.infer_struct_offsets(start, end, reg_name)
        try:
            sark.structure.create_struct_from_offsets(struct_name, offsets)
        except sark.exceptions.SarkStructAlreadyExists:
            yes_no_cancel = idc.AskYN(idaapi.ASKBTN_NO,
                                      "Struct already exists. Modify?\n"
                                      "Cancel to avoid applying the struct.")
            if yes_no_cancel == idaapi.ASKBTN_CANCEL:
                return

            elif yes_no_cancel == idaapi.ASKBTN_YES:
                sid = sark.structure.get_struct(struct_name)
                sark.structure.set_struct_offsets(offsets, sid)

            else:  # yes_no_cancel == idaapi.ASKBTN_NO:
                pass

        sark.structure.apply_struct(start, end, reg_name, struct_name)


def PLUGIN_ENTRY():
    return AutoStruct()


repository: polymorf/findcrypt-yara
file: findcrypt3.py
description: IDA pro plugin to find crypto constants (and more)
wanted_name: Findcrypt
comment: Findcrypt plugin for IDA Pro (using yara framework)
created_at: 2017-03-20 16:02:35+00:00
pushed_at: 2024-11-19 11:36:18+00:00
url:
https://github.com/polymorf/findcrypt-yara/blob/044644b9c52ae3e7b2305bac15b0c12c9e31a282/findcrypt3.py

content:
# -*- coding: utf-8 -*-

import idaapi
import idautils
import ida_bytes
import ida_diskio
import idc
import operator
import yara
import os
import glob

VERSION = "0.2"
YARARULES_CFGFILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "findcrypt3.rules")

try:
    class Kp_Menu_Context(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

        @classmethod
        def get_name(self):
            return self.__name__

        @classmethod
        def get_label(self):
            return self.label

        @classmethod
        def register(self, plugin, label):
            self.plugin = plugin
            self.label = label
            instance = self()
            return idaapi.register_action(idaapi.action_desc_t(
                self.get_name(),  # Name. Acts as an ID. Must be unique.
                instance.get_label(),  # Label. That's what users see.
                instance  # Handler. Called when activated, and for updating
            ))

        @classmethod
        def unregister(self):
            """Unregister the action.
            After unregistering the class cannot be used.
            """
            idaapi.unregister_action(self.get_name())

        @classmethod
        def activate(self, ctx):
            # dummy method
            return 1

        @classmethod
        def update(self, ctx):
            if ctx.widget_type == idaapi.BWN_DISASM:
                return idaapi.AST_ENABLE_FOR_WIDGET
            return idaapi.AST_DISABLE_FOR_WIDGET

    class Searcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1

except:
    pass


p_initialized = False


class YaraSearchResultChooser(idaapi.Choose):
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False):
        idaapi.Choose.__init__(
            self,
            title,
            [
                ["Address", idaapi.Choose.CHCOL_HEX|10],
                ["Rules file", idaapi.Choose.CHCOL_PLAIN|12],
                ["Name", idaapi.Choose.CHCOL_PLAIN|25],
                ["String", idaapi.Choose.CHCOL_PLAIN|25],
                ["Value", idaapi.Choose.CHCOL_PLAIN|40],
                ["Hex", idaapi.Choose.CHCOL_PLAIN|45],
            ],
            flags=flags,
            width=width,
            height=height,
            embedded=embedded)
        self.items = items
        self.selcount = 0
        self.n = len(items)

    def OnClose(self):
        return

    def OnSelectLine(self, n):
        self.selcount += 1
        idc.jumpto(self.items[n][0])

    def OnGetLine(self, n):
        res = self.items[n]
        res = [idc.atoa(res[0]), res[1], res[2], res[3], res[4], res[5]]
        return res

    def OnGetSize(self):
        n = len(self.items)
        return n

    def show(self):
        return self.Show() >= 0

#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class Findcrypt_Plugin_t(idaapi.plugin_t):
    comment = "Findcrypt plugin for IDA Pro (using yara framework)"
    help = "todo"
    wanted_name = "Findcrypt"
    wanted_hotkey = "Ctrl-Alt-F"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        global p_initialized

        # register popup menu handlers
        try:
            Searcher.register(self, "Findcrypt")
        except:
            pass

        if p_initialized is False:
            p_initialized = True
            self.user_directory = self.get_user_directory()
            idaapi.register_action(idaapi.action_desc_t(
                "Findcrypt",
                "Find crypto constants",
                Searcher(),
                None,
                None,
                0))
            idaapi.attach_action_to_menu("Search", "Findcrypt", idaapi.SETMENU_APP)
            print("=" * 80)
            print("Findcrypt v{0} by David BERARD, 2017".format(VERSION))
            print("Findcrypt search shortcut key is Ctrl-Alt-F")
            print("Global rules in %s" % YARARULES_CFGFILE)
            print("User-defined rules in %s/*.rules" % self.user_directory)
            print("=" * 80)

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass


    def toVirtualAddress(self, offset, segments):
        va_offset = 0
        for seg in segments:
            if seg[1] <= offset < seg[2]:
                va_offset = seg[0] + (offset - seg[1])
        return va_offset


    def get_user_directory(self):
        user_dir = ida_diskio.get_user_idadir()
        plug_dir = os.path.join(user_dir, "plugins")
        res_dir = os.path.join(plug_dir, "findcrypt-yara")
        if not os.path.exists(res_dir):
            os.makedirs(res_dir, 0o755)
        return res_dir


    def get_rules_files(self):
        rules_filepaths = {"global":YARARULES_CFGFILE}
        for fpath in glob.glob(os.path.join(self.user_directory, "*.rules")):
            name = os.path.basename(fpath)
            rules_filepaths.update({name:fpath})
        return rules_filepaths


    def search(self):
        memory, offsets = self._get_memory()
        rules = yara.compile(filepaths=self.get_rules_files())
        values = self.yarasearch(memory, offsets, rules)
        c = YaraSearchResultChooser("Findcrypt results", values)
        r = c.show()

    def yarasearch(self, memory, offsets, rules):
        print(">>> start yara search")
        values = list()
        matches = rules.match(data=memory)
        for match in matches:
            for string in match.strings:
                for instance in string.instances:
                    name = match.rule
                    if name.endswith("_API"):
                        try:
                            name = name + "_" + idc.GetString(self.toVirtualAddress(instance.offset, offsets))
                        except:
                            pass
                    value = [
                        self.toVirtualAddress(instance.offset, offsets),
                        match.namespace,
                        name + "_" + hex(self.toVirtualAddress(instance.offset, offsets)).lstrip("0x").rstrip("L").upper(),
                        string.identifier,
                        repr(instance.matched_data),
                        instance.matched_data.hex().upper(),
                    ]
                    idaapi.set_name(value[0], name
                             + "_"
                             + hex(self.toVirtualAddress(instance.offset, offsets)).lstrip("0x").rstrip("L").upper()
                             , 0)
                    values.append(value)
        print("<<< end yara search")
        return values

    def _get_memory(self):
        result = bytearray()
        segment_starts = [ea for ea in idautils.Segments()]
        offsets = []
        start_len = 0
        for start in segment_starts:
            end = idc.get_segm_attr(start, idc.SEGATTR_END)
            result += ida_bytes.get_bytes(start, end - start)
            offsets.append((start, start_len, len(result)))
            start_len = len(result)
        return bytes(result), offsets

    def run(self, arg):
        self.search()


# register IDA plugin
def PLUGIN_ENTRY():
    return Findcrypt_Plugin_t()


repository: pfalcon/idapython
file: Scripts/VirusTotal.py
description: Automatically exported from code.google.com/p/idapython
wanted_name: VirusTotal report
comment: VirusTotal plugin for IDA
created_at: 2015-04-14 07:52:06+00:00
pushed_at: 2015-04-14 07:54:41+00:00
url:
https://github.com/pfalcon/idapython/blob/081b988a03b88867786ad4131269db6930637a5b/Scripts/VirusTotal.py

content:
# -----------------------------------------------------------------------
# VirusTotal IDA Plugin
# By Elias Bachaalany <elias at hex-rays.com>
# (c) Hex-Rays 2011
#
# Special thanks:
# - VirusTotal team
# - Bryce Boe for his VirusTotal Python code
#
import idaapi
import idc
from idaapi import Choose2, plugin_t
import BboeVt as vt
import webbrowser
import urllib
import os


PLUGIN_TEST = 0

# -----------------------------------------------------------------------
# Configuration file
VT_CFGFILE = os.path.join(idaapi.get_user_idadir(), "virustotal.cfg")

# -----------------------------------------------------------------------
# VirusTotal Icon in PNG format
VT_ICON = (
    "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A\x00\x00\x00\x0D\x49\x48\x44\x52"
    "\x00\x00\x00\x10\x00\x00\x00\x10\x04\x03\x00\x00\x00\xED\xDD\xE2"
    "\x52\x00\x00\x00\x30\x50\x4C\x54\x45\x03\x8B\xD3\x5C\xB4\xE3\x9C"
    "\xD1\xED\xF7\xFB\xFD\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\xD3\xF2\x42\x61\x00\x00\x00"
    "\x4B\x49\x44\x41\x54\x78\x9C\x2D\xCA\xC1\x0D\x80\x30\x0C\x43\x51"
    "\x27\x2C\x50\x89\x05\x40\x2C\x40\xEB\xFD\x77\xC3\x76\xC9\xE9\xEB"
    "\xC5\x20\x5F\xE8\x1A\x0F\x97\xA3\xD0\xE4\x1D\xF9\x49\xD1\x59\x29"
    "\x4C\x43\x9B\xD0\x15\x01\xB5\x4A\x9C\xE4\x70\x14\x39\xB3\x31\xF8"
    "\x15\x70\x04\xF4\xDA\x20\x39\x02\x8A\x0D\xA8\x0F\x94\xA7\x09\x0E"
    "\xC5\x16\x2D\x54\x00\x00\x00\x00\x49\x45\x4E\x44\xAE\x42\x60\x82")


# -----------------------------------------------------------------------
class VirusTotalConfig(object):
    def __init__(self):
        self.Default()


    def Default(self):
        self.md5sum = GetInputMD5()
        self.infile = idaapi.dbg_get_input_path()
        if not self.infile:
            self.infile = ""

        # Persistent options
        self.apikey = ""
        self.options = 1 | 2


    def Read(self):
        """
        Read configuration from file
        """
        if not os.path.exists(VT_CFGFILE):
            return
        f = open(VT_CFGFILE, 'r')
        lines = f.readlines()
        for i, line in enumerate(lines):
            line = line.strip()
            if i == 0:
                self.apikey = line
            elif i == 1:
                self.options = int(line)
            else:
                break


    def Write(self):
        """
        Write configuration to file
        """
        lines = (self.apikey.strip(), str(self.options))
        try:
            f = open(VT_CFGFILE, 'w')
            f.write("\n".join(lines))
            f.close()
        except:
            pass


# -----------------------------------------------------------------------
def VtReport(apikey, filename=None, md5sum=None):
    if filename is None and md5sum is None:
        return (False, "No parameters passed!")

    # Check filename existance
    if filename is not None and not os.path.exists(filename):
        return (False, "Input file '%s' does not exist!" % filename)

    #print("fn=%s md5=%s" % (filename, md5sum))
    # Get file report from VirusTotal
    try:
        vt.set_apikey(apikey)
        result = vt.get_file_report(filename=filename, md5sum=md5sum)
    except Exception as e:
        return (False, "Exception:\n%s" % str(e))

    # Already analyzed?
    if result is not None:
        # Transform the results
        items = []
        for av, mwname in result.items():
            mwname = str(mwname) if mwname else "n/a"
            av = str(av)
            items.append([av, mwname])
        result = items

    return (True, result)


# -----------------------------------------------------------------------
class VirusTotalChooser(Choose2):
    """
    Chooser class to display results from VT
    """
    def __init__(self, title, items, icon, embedded=False):
        Choose2.__init__(self,
                         title,
                         [ ["Antivirus", 20], ["Result", 40] ],
                         embedded=embedded)
        self.items = items
        self.icon = icon


    def GetItems(self):
        return self.items


    def SetItems(self, items):
        self.items = [] if items is None else items


    def OnClose(self):
        pass


    def OnGetLine(self, n):
        return self.items[n]


    def OnGetSize(self):
        return len(self.items)


    def OnSelectLine(self, n):
        # Google search for the malware name and the antivirus name
        s = urllib.urlencode({"q" : " ".join(self.items[n])})
        webbrowser.open_new_tab("http://www.google.com/search?%s" % s)


# --------------------------------------------------------------------------
class VirusTotalForm(Form):
    def __init__(self, icon):
        self.EChooser = VirusTotalChooser("E1", [], icon, embedded=True)
        Form.__init__(self, r"""STARTITEM {id:txtInput}
VirusTotal - IDAPython plugin v1.0 (c) Hex-Rays

{FormChangeCb}
<#API key#~A~pi key:{txtApiKey}>

Options:
<#Open results in a chooser when form closes#~P~opout results on close:{rOptRemember}>
<#Use MD5 checksum#~M~D5Sum:{rOptMD5}>
<#Use file on disk#~F~ile:{rOptFile}>{grpOptions}>

<#Type input (file or MD5 string)#~I~nput:{txtInput}>
<Results:{cEChooser}>
<#Get reports from VT#~R~eport:{btnReport}>
""", {
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
            'txtApiKey'   : Form.StringInput(swidth=80),
            'grpOptions'  : Form.ChkGroupControl(("rOptRemember", "rOptMD5", "rOptFile")),
            'txtInput'    : Form.FileInput(open=True),
            'btnReport'   : Form.ButtonInput(self.OnReportClick),
            'cEChooser'   : Form.EmbeddedChooserControl(self.EChooser)
        })



    def OnReportClick(self, code=0):
        pass


    def OnFormChange(self, fid):
        if fid == self.rOptMD5.id or fid == self.rOptFile.id:
            input = (self.cfg.md5sum, self.cfg.infile)
            if fid == self.rOptMD5.id:
                c1 = self.rOptMD5
                c2 = self.rOptFile
                idx = 0
            else:
                c1 = self.rOptFile
                c2 = self.rOptMD5
                idx = 1

            v = not self.GetControlValue(c1)
            if v: idx = not idx

            # Uncheck the opposite input type
            self.SetControlValue(c2, v)

            # Set input field depending on input type
            self.SetControlValue(self.txtInput, input[idx])
        #
        # Report button
        #
        elif fid == self.btnReport.id:
            input = self.GetControlValue(self.txtInput)
            as_file = self.GetControlValue(self.rOptFile)
            apikey = self.GetControlValue(self.txtApiKey)

            ok, r = VtReport(self.cfg.apikey,
                        filename=input if as_file else None,
                        md5sum=None if as_file else input)

            # Error?
            if not ok:
                idc.Warning(r)
                return 1

            # Pass the result
            self.EChooser.SetItems(r)

            # We have results and it was a file? Print its MD5
            if r and as_file:
                print("%s: %s" % (vt.LAST_FILE_HASH, input))

            # Refresh the embedded chooser control
            # (Could also clear previous results if not were retrieved during this run)
            self.RefreshField(self.cEChooser)

            # Store the input for the caller
            self.cfg.input = input

            # No results and file as input was supplied?
            if r is None:
                if as_file:
                    # Propose to upload
                    if idc.AskYN(0, "HIDECANCEL\nNo previous results. Do you want to submit the file:\n\n'%s'\n\nto VirusTotal?" % input) == 0:
                        return 1

                    try:
                        r = vt.scan_file(input)
                    except Exception as e:
                        idc.Warning("Exceptio during upload: %s" % str(e))
                    else:
                        if r is None:
                            idc.Warning("Failed to upload the file!")
                        else:
                            idc.Warning("File uploaded. Check again later to get the analysis report. Scan id: %s" % r)
                else:
                    idc.Warning("No results found for hash: %s" % input)

        return 1


    def Show(self, cfg):
        # Compile the form once
        if not self.Compiled():
            _, args = self.Compile()
            #print args[0]

        # Populate the form
        self.txtApiKey.value  = cfg.apikey
        self.grpOptions.value = cfg.options
        self.txtInput.value   = cfg.infile if self.rOptFile.checked else cfg.md5sum

        # Remember the config
        self.cfg = cfg

        # Execute the form
        ok = self.Execute()

        # Forget the cfg
        del self.cfg

        # Success?
        if ok != 0:
            # Update config
            cfg.options = self.grpOptions.value
            cfg.apikey  = self.txtApiKey.value

            # Popup results?
            if self.rOptRemember.checked:
                ok = 2

        return ok


# -----------------------------------------------------------------------
class VirusTotalPlugin_t(plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "VirusTotal plugin for IDA"
    help = ""
    wanted_name = "VirusTotal report"
    wanted_hotkey = "Alt-F8"


    def init(self):
        # Some initialization
        self.icon_id = 0
        return idaapi.PLUGIN_OK


    def run(self, arg=0):
        # Load icon from the memory and save its id
        self.icon_id = idaapi.load_custom_icon(data=VT_ICON, format="png")
        if self.icon_id == 0:
            raise RuntimeError("Failed to load icon data!")

        # Create config object
        cfg = VirusTotalConfig()

        # Read previous config
        cfg.Read()

        # Create form
        f = VirusTotalForm(self.icon_id)

        # Show the form
        ok = f.Show(cfg)
        if ok == 0:
            f.Free()
            return

        # Save configuration
        cfg.Write()

        # Spawn a non-modal chooser w/ the results if any
        if ok == 2 and f.EChooser.GetItems():
            VirusTotalChooser(
                "VirusTotal results [%s]" % cfg.input,
                f.EChooser.GetItems(),
                self.icon_id).Show()

        f.Free()
        return


    def term(self):
        # Free the custom icon
        if self.icon_id != 0:
            idaapi.free_custom_icon(self.icon_id)


# -----------------------------------------------------------------------
def PLUGIN_ENTRY():
    return VirusTotalPlugin_t()

# --------------------------------------------------------------------------
if PLUGIN_TEST:
    # Create form
    f = PLUGIN_ENTRY()
    f.init()
    f.run()
    f.term()




repository: a1ext/auto_re
file: auto_re.py
description: IDA PRO auto-renaming plugin with tagging support
url: https://github.com/a1ext/auto_re/blob/2ea88f0c18aa06ccb2fe073ab5aece27736dce15/auto_re.py
wanted_name: Auto RE
created_at: 2016-09-16 17:26:41+00:00
pushed_at: 2024-09-05 21:19:14+00:00
content:
# -*- coding: utf-8 -*
__author__ = 'Trafimchuk Aliaksandr'
__version__ = '1.9'

from collections import defaultdict
import idaapi
from idautils import FuncItems, CodeRefsTo
from idaapi import o_reg, o_imm, o_far, o_near, o_mem, o_displ
import os
import re
import sys
import traceback


HAS_PYSIDE = idaapi.IDA_SDK_VERSION < 690
if HAS_PYSIDE:
    from PySide import QtGui, QtCore
    from PySide.QtGui import QTreeView, QVBoxLayout, QLineEdit, QMenu, QInputDialog, QAction, QTabWidget
else:
    from PyQt5 import QtGui, QtCore
    from PyQt5.QtWidgets import QTreeView, QVBoxLayout, QLineEdit, QMenu, QInputDialog, QAction, QTabWidget


try:
    # Python 2.
    xrange
except NameError:
    # Python 3.
    xrange = range


# enable to allow PyCharm remote debug
RDEBUG = False
# adjust this value to be a full path to a debug egg
RDEBUG_EGG = r'c:\Program Files\JetBrains\PyCharm 2017.1.4\debug-eggs\pycharm-debug.egg'
RDEBUG_HOST = 'localhost'
RDEBUG_PORT = 12321


TAGS_IGNORE_LIST = {
    'OpenProcessToken',
    'DisconnectNamedPipe'
}

IGNORE_CALL_LIST = {
    'RtlNtStatusToDosError',
    'GetLastError',
    'SetLastError'
}

TAGS = {
    'net': ['WSAStartup', 'socket', 'recv', 'recvfrom', 'send', 'sendto', 'acccept', 'bind', 'listen', 'select',
            'setsockopt', 'ioctlsocket', 'closesocket', 'WSAAccept', 'WSARecv', 'WSARecvFrom', 'WSASend', 'WSASendTo',
            'WSASocket', 'WSAConnect', 'ConnectEx', 'TransmitFile', 'HTTPOpenRequest', 'HTTPSendRequest',
            'URLDownloadToFile', 'InternetCrackUrl', 'InternetOpen', 'InternetOpen', 'InternetConnect',
            'InternetOpenUrl', 'InternetQueryOption', 'InternetSetOption', 'InternetReadFile', 'InternetWriteFile',
            'InternetGetConnectedState', 'InternetSetStatusCallback', 'DnsQuery', 'getaddrinfo', 'GetAddrInfo',
            'GetAdaptersInfo', 'GetAdaptersAddresses', 'HttpQueryInfo', 'ObtainUserAgentString', 'WNetGetProviderName',
            'GetBestInterfaceEx', 'gethostbyname', 'getsockname', 'connect', 'WinHttpOpen', 'WinHttpSetTimeouts',
            'WinHttpSendRequest', 'WinHttpConnect', 'WinHttpCrackUrl', 'WinHttpReadData', 'WinHttpOpenRequest',
            'WinHttpReceiveResponse', 'WinHttpQueryHeaders', 'HttpSendRequestW', 'HttpSendRequestA', 'HttpAddRequestHeadersW', 'HttpAddRequestHeadersA', 'HttpOpenRequestW', 'HttpOpenRequestA', 'NetServerGetInfo', 'NetApiBufferFree', 'NetWkstaGetInfo'],
    'spawn': ['CreateProcess', 'ShellExecute', 'ShellExecuteEx', 'system', 'CreateProcessInternal', 'NtCreateProcess',
              'ZwCreateProcess', 'NtCreateProcessEx', 'ZwCreateProcessEx', 'NtCreateUserProcess', 'ZwCreateUserProcess',
              'RtlCreateUserProcess', 'NtCreateSection', 'ZwCreateSection', 'NtOpenSection', 'ZwOpenSection',
              'NtAllocateVirtualMemory', 'ZwAllocateVirtualMemory', 'NtWriteVirtualMemory', 'ZwWriteVirtualMemory',
              'NtMapViewOfSection', 'ZwMapViewOfSection', 'OpenSCManager', 'CreateService', 'OpenService',
              'StartService', 'ControlService', 'ShellExecuteExA', 'ShellExecuteExW'],
    'inject': ['OpenProcess-disabled', 'ZwOpenProcess', 'NtOpenProcess', 'WriteProcessMemory', 'NtWriteVirtualMemory',
               'ZwWriteVirtualMemory', 'CreateRemoteThread', 'QueueUserAPC', 'ZwUnmapViewOfSection', 'NtUnmapViewOfSection'],
    'com': ['CoCreateInstance', 'CoInitializeSecurity', 'CoGetClassObject', 'OleConvertOLESTREAMToIStorage', 'CreateBindCtx', 'CoSetProxyBlanket', 'VariantClear'],
    'crypto': ['CryptAcquireContext', 'CryptProtectData', 'CryptUnprotectData', 'CryptProtectMemory',
               'CryptUnprotectMemory', 'CryptDecrypt', 'CryptEncrypt', 'CryptHashData', 'CryptDecodeMessage',
               'CryptDecryptMessage', 'CryptEncryptMessage', 'CryptHashMessage', 'CryptExportKey', 'CryptGenKey',
               'CryptCreateHash', 'CryptDecodeObjectEx', 'EncryptMessage', 'DecryptMessage'],
    'kbd': ['SendInput', 'VkKeyScanA', 'VkKeyScanW'],
    'file': ['_open64', 'open64', 'open', 'open64', 'fopen', 'fread', 'fclose', 'fwrite', 'flock', 'read', 'write',
             'fstat', 'lstat', 'stat', 'chmod', 'chown', 'lchown', 'link', 'symlink', 'readdir', 'readdir64'],
    'reg': ['RegOpenKeyExW', 'RegQueryValueExW', 'RegSetValueExW', 'RegCreateKeyExW', 'RegDeleteValueW', 'RegEnumKeyW', 'RegCloseKey', 'RegQueryInfoKeyW', 'RegOpenKeyExA', 'RegQueryValueExA', 'RegSetValueExA', 'RegCreateKeyExA', 'RegDeleteValueA', 'RegEnumKeyA',  'RegQueryInfoKeyA'],
    'dev': ['DeviceIoControl'],
    'wow': ['Wow64DisableWow64FsRedirection', 'Wow64RevertWow64FsRedirection']
}

STRICT_TAG_NAME_CHECKING = {'file'}

blacklist = {'@__security_check_cookie@4', '__SEH_prolog4', '__SEH_epilog4'}
replacements = [
    ('??3@YAXPAX@Z', 'alloc'),
    ('?', '')
]

def inf_is_64bit():
    return (idaapi.inf_is_64bit if idaapi.IDA_SDK_VERSION >= 900 else idaapi.cvar.inf.is_64bit)()


def get_addr_width():
    return '16' if inf_is_64bit() else '8'


def decode_insn(ea):
    if idaapi.IDA_SDK_VERSION >= 700 and sys.maxsize > 2**32:
        insn = idaapi.insn_t()
        if idaapi.decode_insn(insn, ea) > 0:
            return insn
    else:
        if idaapi.decode_insn(ea):
            return idaapi.cmd.copy()


def force_name(ea, new_name):
    if not ea or ea == idaapi.BADADDR:
        return
    if idaapi.IDA_SDK_VERSION >= 700:
        return idaapi.force_name(ea, new_name, idaapi.SN_NOCHECK)
    return idaapi.do_name_anyway(ea, new_name, 0)


class AutoReIDPHooks(idaapi.IDP_Hooks):
    """
    Hooks to keep view updated if some function is updated
    """
    def __init__(self, view, *args):
        super(AutoReIDPHooks, self).__init__(*args)
        self._view = view

    def __on_rename(self, ea, new_name):
        if not self._view:
            return
        items = self._view._model.findItems(('%0' + get_addr_width() + 'X') % ea, QtCore.Qt.MatchRecursive)
        if len(items) != 1:
            return

        item = items[0]
        index = self._view._model.indexFromItem(item)
        if not index.isValid():
            return

        name_index = index.sibling(index.row(), 1)
        if not name_index.isValid():
            return

        self._view._model.setData(name_index, new_name)

    def ev_rename(self, ea, new_name):
        """ callback for IDA >= 700 """
        self.__on_rename(ea, new_name)
        return super(AutoReIDPHooks, self).ev_rename(ea, new_name)

    def rename(self, ea, new_name):
        """ callback for IDA < 700 """
        self.__on_rename(ea, new_name)
        return super(AutoReIDPHooks, self).rename(ea, new_name)


class AutoREView(idaapi.PluginForm):
    ADDR_ROLE = QtCore.Qt.UserRole + 1

    OPT_FORM_PERSIST = idaapi.PluginForm.FORM_PERSIST if hasattr(idaapi.PluginForm, 'FORM_PERSIST') else idaapi.PluginForm.WOPN_PERSIST
    OPT_FORM_NO_CONTEXT = idaapi.PluginForm.FORM_NO_CONTEXT if hasattr(idaapi.PluginForm, 'FORM_NO_CONTEXT') else idaapi.PluginForm.WCLS_NO_CONTEXT

    def __init__(self, data):
        super(AutoREView, self).__init__()
        self._data = data
        self.tv = None
        self._model = None
        self._idp_hooks = None

    def Show(self):
        return idaapi.PluginForm.Show(self, 'AutoRE', options=self.OPT_FORM_PERSIST)

    def _get_parent_widget(self, form):
        if HAS_PYSIDE:
            return self.FormToPySideWidget(form)
        return self.FormToPyQtWidget(form)

    def OnCreate(self, form):
        self.parent = self._get_parent_widget(form)

        self._idp_hooks = AutoReIDPHooks(self)
        if not self._idp_hooks.hook():
            print('IDP_Hooks.hook() failed')

        self.tv = QTreeView()
        self.tv.setExpandsOnDoubleClick(False)

        root_layout = QVBoxLayout(self.parent)
        # self.le_filter = QLineEdit(self.parent)

        # root_layout.addWidget(self.le_filter)
        root_layout.addWidget(self.tv)

        self.parent.setLayout(root_layout)

        self._model = QtGui.QStandardItemModel()
        self._init_model()
        self.tv.setModel(self._model)

        self.tv.setColumnWidth(0, 200)
        self.tv.setColumnWidth(1, 300)
        self.tv.header().setStretchLastSection(True)

        self.tv.expandAll()

        self.tv.doubleClicked.connect(self.on_navigate_to_method_requested)
        # self.le_filter.textChanged.connect(self.on_filter_text_changed)
        self.tv.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        self.tv.customContextMenuRequested.connect(self._tree_customContextMenuRequesssted)

        rename_action = QAction('Rename...', self.tv)
        rename_action.setShortcut('n')
        rename_action.triggered.connect(self._tv_rename_action_triggered)
        self.tv.addAction(rename_action)

    def _tree_customContextMenuRequesssted(self, pos):
        idx = self.tv.indexAt(pos)
        if not idx.isValid():
            return

        addr = idx.data(role=self.ADDR_ROLE)
        if not addr:
            return

        name_idx = idx.sibling(idx.row(), 1)
        old_name = name_idx.data()

        menu = QMenu()
        rename_action = menu.addAction('Rename `%s`...' % old_name)
        rename_action.setShortcut('n')
        action = menu.exec_(self.tv.mapToGlobal(pos))
        if action == rename_action:
            return self._rename_ea_requested(addr, name_idx)

    def _tv_rename_action_triggered(self):
        selected = self.tv.selectionModel().selectedIndexes()
        if not selected:
            return

        idx = selected[0]
        if not idx.isValid():
            return

        addr = idx.data(role=self.ADDR_ROLE)
        if not addr:
            return

        name_idx = idx.sibling(idx.row(), 1)
        if not name_idx.isValid():
            return

        return self._rename_ea_requested(addr, name_idx)

    def _rename_ea_requested(self, addr, name_idx):
        old_name = name_idx.data()

        if idaapi.IDA_SDK_VERSION >= 700:
            new_name = idaapi.ask_str(str(old_name), 0, 'New name:')
        else:
            new_name = idaapi.askstr(0, str(old_name), 'New name:')

        if new_name is None:
            return

        force_name(addr, new_name)
        renamed_name = idaapi.get_ea_name(addr)
        name_idx.model().setData(name_idx, renamed_name)

    def OnClose(self, form):
        if self._idp_hooks:
            self._idp_hooks.unhook()

    def _tv_init_header(self, model):
        item_header = QtGui.QStandardItem("EA")
        item_header.setToolTip("Address")
        model.setHorizontalHeaderItem(0, item_header)

        item_header = QtGui.QStandardItem("Function name")
        model.setHorizontalHeaderItem(1, item_header)

        item_header = QtGui.QStandardItem("API called")
        model.setHorizontalHeaderItem(2, item_header)

    # noinspection PyMethodMayBeStatic
    def _tv_make_tag_item(self, name):
        rv = QtGui.QStandardItem(name)

        rv.setEditable(False)
        return [rv, QtGui.QStandardItem(), QtGui.QStandardItem()]

    def _tv_make_ref_item(self, tag, ref):
        ea_item = QtGui.QStandardItem(('%0' + get_addr_width() + 'X') % ref['ea'])
        ea_item.setEditable(False)
        ea_item.setData(ref['ea'], self.ADDR_ROLE)

        name_item = QtGui.QStandardItem(ref['name'])
        name_item.setEditable(False)
        name_item.setData(ref['ea'], self.ADDR_ROLE)

        apis = ', '.join(ref['tags'][tag])
        api_name = QtGui.QStandardItem(apis)
        api_name.setEditable(False)
        api_name.setData(ref['ea'], self.ADDR_ROLE)
        api_name.setToolTip(apis)

        return [ea_item, name_item, api_name]

    def _init_model(self):
        self._model.clear()

        root_node = self._model.invisibleRootItem()
        self._tv_init_header(self._model)

        for tag, refs in self._data.items():
            item_tag_list = self._tv_make_tag_item(tag)
            item_tag = item_tag_list[0]

            root_node.appendRow(item_tag_list)

            for ref in refs:
                ref_item_list = self._tv_make_ref_item(tag, ref)

                item_tag.appendRow(ref_item_list)

    def on_navigate_to_method_requested(self, index):
        addr = index.data(role=self.ADDR_ROLE)
        if addr is not None:
            idaapi.jumpto(addr)

    # def on_filter_text_changed(self, text):
    #     print('on_text_changed: %s' % text)


class auto_re_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = ""

    help = ""
    wanted_name = "Auto RE"
    wanted_hotkey = "Ctrl+Shift+M"

    _PREFIX_NAME = 'au_re_'
    _MIN_MAX_MATH_OPS_TO_ALLOW_RENAME = 10

    _CALLEE_NODE_NAMES = {
        idaapi.PLFM_MIPS: '$ mips',
        idaapi.PLFM_ARM: '$ arm'
    }
    _DEFAULT_CALLEE_NODE_NAME = '$ vmm functions'

    _JMP_TYPES = {idaapi.NN_jmp, idaapi.NN_jmpni, idaapi.NN_jmpfi, idaapi.NN_jmpshort}

    def __init__(self):
        super(auto_re_t, self).__init__()
        self._data = None
        self.view = None

    def init(self):
        # self._cfg = None
        self.view = None
        # self._load_config()

        return idaapi.PLUGIN_OK

    # def _load_config(self):
    #     self._cfg = {'auto_rename': False}

    # def _store_config(self, cfg):
    #     pass

    def _handle_tags(self, fn, fn_an, known_refs):
        if known_refs:
            known_refs = dict(known_refs)
            for k, names in known_refs.items():
                existing = set(fn_an['tags'][k])
                new = set(names) - existing
                if new:
                    fn_an['tags'][k] += list(new)

        tags = dict(fn_an['tags'])
        if not tags:
            return
        print('fn: %#08x tags: %s' % (self.start_ea_of(fn), tags))
        cmt = idaapi.get_func_cmt(fn, True)
        if cmt:
            cmt += '\n'
        s = str(tags.keys())
        name = idaapi.get_ea_name(self.start_ea_of(fn))
        item = {'ea': self.start_ea_of(fn), 'name': name, 'tags': tags}
        if not cmt or s not in cmt:
            idaapi.set_func_cmt(fn, '%sTAGS: %s' % (cmt or '', s), True)
        # self.mark_position(self.start_ea_of(fn), 'TAGS: %s' % s)
        for tag in tags:
            if tag not in self._data:
                self._data[tag] = list()
            self._data[tag].append(item)

    def _handle_calls(self, fn, fn_an):
        num_calls = len(fn_an['calls'])
        if num_calls != 1:
            return

        dis = fn_an['calls'][0]
        if dis.Op1.type not in (o_imm, o_far, o_near, o_mem):
            return

        ea = dis.Op1.value
        if not ea and dis.Op1.addr:
            ea = dis.Op1.addr

        if idaapi.has_dummy_name(self.get_flags_at(ea)):
            return

        # TODO: check is there jmp, push+retn then don't rename the func
        if fn_an['strange_flow']:
            return

        possible_name = idaapi.get_ea_name(ea)
        if not possible_name or possible_name in blacklist:
            return

        normalized = self.normalize_name(possible_name)

        # if self._cfg.get('auto_rename'):
        if len(fn_an['math']) < self._MIN_MAX_MATH_OPS_TO_ALLOW_RENAME:
            force_name(self.start_ea_of(fn), normalized)
        # TODO: add an API to the view
        print('fn: %#08x: %d calls, %d math%s possible name: %s, normalized: %s' % (
            self.start_ea_of(fn), len(fn_an['calls']), len(fn_an['math']), 'has bads' if fn_an['has_bads'] else '',
            possible_name, normalized))

    # noinspection PyMethodMayBeStatic
    def _check_is_jmp_wrapper(self, dis):
        # checks instructions like `jmp API`
        if dis.itype not in self._JMP_TYPES:
            return

        # handle call wrappers like jmp GetProcAddress
        if dis.Op1.type == idaapi.o_mem and dis.Op1.addr:
            # TODO: check is there better way to determine is the function a wrapper
            v = dis.Op1.addr
            flags = self.get_flags_at(v)
            if v and dis.itype == idaapi.NN_jmpni and self.is_data(flags) and self.__is_ptr_val(flags):
                v = self.__get_ptr_val(v)
            return v

    # noinspection PyMethodMayBeStatic
    def _check_is_push_retn_wrapper(self, dis0, dis1):
        """
        Checks for sequence of push IMM32/retn
        :param dis0: the first insn
        :param dis1: the second insn
        :return: value of IMM32
        """
        if dis0.itype != idaapi.NN_push or dis0.Op1.type != idaapi.o_imm or not dis0.Op1.value:
            return

        if dis1.itype not in (idaapi.NN_retn,):
            return

        return dis0.Op1.value

    def _preprocess_api_wrappers(self, fnqty):
        rv = defaultdict(dict)

        for i in xrange(fnqty):
            fn = idaapi.getn_func(i)
            items = list(FuncItems(self.start_ea_of(fn)))
            if len(items) not in (1, 2):
                continue

            dis0 = decode_insn(items[0])
            if dis0 is None:
                continue
            addr = self._check_is_jmp_wrapper(dis0)

            if not addr and len(items) > 1:
                dis1 = decode_insn(items[1])
                if dis1 is not None:
                    addr = self._check_is_push_retn_wrapper(dis0, dis1)

            if not addr:
                continue

            name = idaapi.get_ea_name(addr)
            name = name.replace(idaapi.FUNC_IMPORT_PREFIX, '')
            if not name:
                continue

            imp_stripped_name = name.lstrip('_')

            for tag, names in TAGS.items():
                for tag_api in names:
                    if tag in STRICT_TAG_NAME_CHECKING:
                        match = tag_api in (name, imp_stripped_name)
                    else:
                        match = tag_api in name
                    if not match:
                        continue

                    refs = list(CodeRefsTo(self.start_ea_of(fn), 1))

                    for ref in refs:
                        ref_fn = idaapi.get_func(ref)
                        if not ref_fn:
                            # idaapi.msg('AutoRE: there is no func for ref: %08x for api: %s' % (ref, name))
                            continue
                        if tag not in rv[self.start_ea_of(ref_fn)]:
                            rv[self.start_ea_of(ref_fn)][tag] = list()
                        if name not in rv[self.start_ea_of(ref_fn)][tag]:
                            rv[self.start_ea_of(ref_fn)][tag].append(name)
        return dict(rv)

    def run(self, arg):
        if RDEBUG and RDEBUG_EGG:
            if not os.path.isfile(RDEBUG_EGG):
                idaapi.msg('AutoRE: Remote debug is enabled, but I cannot find the debug egg: %s' % RDEBUG_EGG)
            else:
                import sys

                if RDEBUG_EGG not in sys.path:
                    sys.path.append(RDEBUG_EGG)

                import pydevd
                pydevd.settrace(RDEBUG_HOST, port=RDEBUG_PORT, stdoutToServer=True, stderrToServer=True)

        try:
            self._data = dict()
            count = idaapi.get_func_qty()

            # pre-process of api wrapper functions
            known_refs_tags = self._preprocess_api_wrappers(count)

            for i in xrange(count):
                fn = idaapi.getn_func(i)
                fn_an = self.analyze_func(fn)

                # if fn_an['math']:
                # 	print('fn: %#08x has math' % self.start_ea_of(fn))

                if idaapi.has_dummy_name(self.get_flags_at(self.start_ea_of(fn))):
                    self._handle_calls(fn, fn_an)

                known_refs = known_refs_tags.get(self.start_ea_of(fn))
                self._handle_tags(fn, fn_an, known_refs)

            if self.view:
                self.view.Close(AutoREView.OPT_FORM_NO_CONTEXT)
            self.view = AutoREView(self._data)
            self.view.Show()
        except:
            idaapi.msg('AutoRE: error: %s\n' % traceback.format_exc())

    def term(self):
        self._data = None

    @classmethod
    def disasm_func(cls, fn):
        rv = list()
        items = list(FuncItems(cls.start_ea_of(fn)))
        for item_ea in items:
            obj = {'ea': item_ea, 'fn_ea': cls.start_ea_of(fn), 'dis': None}
            insn = decode_insn(item_ea)
            if insn is not None:
                obj['dis'] = insn
            rv.append(obj)
        return rv

    @classmethod
    def get_callee_netnode(cls):
        node_name = cls._CALLEE_NODE_NAMES.get(idaapi.ph.id, cls._DEFAULT_CALLEE_NODE_NAME)
        n = idaapi.netnode(node_name)
        return n

    @classmethod
    def get_callee(cls, ea):
        n = cls.get_callee_netnode()
        v = n.altval(ea)
        v -= 1
        if v == idaapi.BADNODE:
            return
        return v

    @classmethod
    def _analysis_handle_call_insn(cls, dis, rv):
        rv['calls'].append(dis)
        if dis.Op1.type != o_mem or not dis.Op1.addr:
            callee = cls.get_callee(dis.ip)
            if not callee:
                return
        else:
            callee = dis.Op1.addr

        cls._apply_tag_on_callee(callee, rv, is_call=True)

    @classmethod
    def _apply_tag_on_callee(cls, callee_ea, rv, is_call=False):
        name = idaapi.get_ea_name(callee_ea)
        name = name.replace(idaapi.FUNC_IMPORT_PREFIX, '')

        if '@' in name:
            name = name.split('@')[0]

        if not name:
            return

        if name in IGNORE_CALL_LIST:
            if is_call:
                rv['calls'].pop()
            return

        if name in TAGS_IGNORE_LIST:
            return

        for tag, names in TAGS.items():
            for tag_api in names:
                if tag in STRICT_TAG_NAME_CHECKING:
                    match = tag_api in (name, name.lstrip('_'))
                else:
                    match = tag_api in name
                if not match or name in rv['tags'][tag]:
                    continue

                # print('%#08x: %s, tag: %s' % (dis.ea, name, tag))
                rv['tags'][tag].append(name)
                break

    @classmethod
    def __is_ptr_val(cls, flags):
        if idaapi.IDA_SDK_VERSION >= 700:
            return (idaapi.is_qword if inf_is_64bit() else idaapi.is_dword)(flags)
        return (idaapi.isQwrd if inf_is_64bit() else idaapi.isDwrd)(flags)

    @classmethod
    def __get_ptr_val(cls, ea):
        if inf_is_64bit():
            return idaapi.get_qword(ea)

        return (idaapi.get_dword if idaapi.IDA_SDK_VERSION >= 700 else idaapi.get_long)(ea)
        
    @classmethod
    def start_ea_of(cls, o):
        return getattr(o, 'start_ea' if idaapi.IDA_SDK_VERSION >= 700 else 'startEA')

    @classmethod
    def end_ea_of(cls, o):
        return getattr(o, 'end_ea' if idaapi.IDA_SDK_VERSION >= 700 else 'endEA')

    @classmethod
    def get_flags_at(cls, ea):
        return getattr(idaapi, 'get_flags' if idaapi.IDA_SDK_VERSION >= 700 else 'getFlags')(ea)

    @classmethod
    def is_data(cls, flags):
        return getattr(idaapi, 'is_data' if idaapi.IDA_SDK_VERSION >= 700 else 'isData')(flags)

    @classmethod
    def analyze_func(cls, fn):
        rv = {
            'fn': fn,
            'calls': [],
            'math': [],
            'has_bads': False,
            'strange_flow': False,
            'tags': defaultdict(list)
        }
        items = cls.disasm_func(fn)
        items_set = set(map(lambda x: x['ea'], items))

        for item in items:
            dis = item['dis']
            if dis is None:
                rv['has_bads'] = True
                continue

            if dis.itype in (idaapi.NN_call, idaapi.NN_callfi, idaapi.NN_callni):
                cls._analysis_handle_call_insn(dis, rv)
            elif dis.itype == idaapi.NN_xor:
                if dis.Op1.type == o_reg and dis.Op2.type == o_reg and dis.Op1.reg == dis.Op2.reg:
                    continue
                rv['math'].append(dis)
            elif dis.itype in (idaapi.NN_shr, idaapi.NN_shl, idaapi.NN_sal, idaapi.NN_sar, idaapi.NN_ror,
                               idaapi.NN_rol, idaapi.NN_rcl, idaapi.NN_rcl):
                # TODO
                rv['math'].append(dis)
            elif dis.itype in cls._JMP_TYPES:
                if dis.Op1.type not in (o_far, o_near, o_mem, o_displ):
                    continue

                if dis.Op1.type == o_displ:
                    rv['strange_flow'] = True
                    continue

                ea = dis.Op1.value
                if not ea and dis.Op1.addr:
                    ea = dis.Op1.addr
                if ea not in items_set:
                    rv['strange_flow'] = True

                # flags = self.get_flags_at(ea)
                # if dis.itype == idaapi.NN_jmpni and dis.Op1.type == o_mem and ea and self.is_data(flags):
                #     if cls.__is_ptr_val(flags):
                #         val = cls.__get_ptr_val(ea)
                #         if val:
                #             cls._apply_tag_on_callee(val, rv, is_call=False)

        return rv

    @classmethod
    def normalize_name(cls, n):
        for repl in replacements:
            n = n.replace(*repl)
        if '@' in n:
            n = n.split('@')[0]
        if len(n) < 3:
            return ''
        if not n.startswith(cls._PREFIX_NAME):
            n = cls._PREFIX_NAME + n
        return n

    # @classmethod
    # def mark_position(cls, ea, name, slot=[0]):
    #     curloc = idaapi.curloc()
    #     curloc.ea = ea
    #     curloc.lnnum = 0
    #     curloc.x = 0
    #     curloc.y = 0
    #     slot[0] += 1
    #     curloc.mark(slot[0], name, name)


# noinspection PyPep8Naming
def PLUGIN_ENTRY():
    return auto_re_t()


repository: dshikashio/idarest
file: idarest.py
description: Expose some basic IDA Pro interactions through a REST API for JSONP
url: https://github.com/dshikashio/idarest/blob/a59144a74044847b9d2b19349aeb244b2bfca26f/idarest.py
wanted_name: IDA Rest API
created_at: 2015-03-03 04:49:53+00:00
pushed_at: 2015-03-21 02:07:22+00:00
content:
from BaseHTTPServer import BaseHTTPRequestHandler,HTTPServer
from SocketServer import ThreadingMixIn
import re
import threading
import cgi
import urlparse
import json

try:
    import idaapi
    import idautils
    import idc
except:
    pass

API_PREFIX = '/ida/api/v1.0'

class HTTPRequestError(BaseException):
    def __init__(self, msg, code):
        self.msg = msg
        self.code = code

class UnknownApiError(HTTPRequestError):
    pass

class HTTPRequestHandler(BaseHTTPRequestHandler):
    routes = {}
    prefns = {}
    postfns = {}

    @staticmethod
    def build_route_pattern(route):
        return re.compile("^{0}$".format(route))

    @staticmethod
    def route(route_str):
        def decorator(f):
            route_path = API_PREFIX + '/' + route_str + '/?'
            route_pattern = HTTPRequestHandler.build_route_pattern(route_path)
            HTTPRequestHandler.routes[route_str] = (route_pattern, f)
            return f
        return decorator

    @staticmethod
    def prefn(route_str):
        def decorator(f):
            HTTPRequestHandler.prefns.setdefault(route_str, []).append(f)
            return f
        return decorator

    @staticmethod
    def postfn(route_str):
        def decorator(f):
            HTTPRequestHandler.postfns.setdefault(route_str, []).append(f)
            return f
        return decorator

    def _get_route_match(self, path):
        for (key, (route_pattern,view_function)) in self.routes.items():
            m = route_pattern.match(path)
            if m:
                return key,view_function
        return None

    def _get_route_prefn(self, key):
        try:
            return self.prefns[key]
        except:
            return []

    def _get_route_postfn(self, key):
        try:
            return self.postfns[key]
        except:
            return []

    def _serve_route(self, args):
        path = urlparse.urlparse(self.path).path
        route_match = self._get_route_match(path)
        if route_match:
            key,view_function = route_match
            for prefn in self._get_route_prefn(key):
                args = prefn(self, args)
            results = view_function(self, args)
            for postfn in self._get_route_postfn(key):
                results = postfn(self, results)
            return results
        else:
            raise UnknownApiError('Route "{0}" has not been registered'.format(path), 404)

    def _serve(self, args):
        try:
            response = {
                'code' : 200,
                'msg'  : 'OK',
                'data' : self._serve_route(args)
            }
        except UnknownApiError as e:
            self.send_error(e.code, e.msg)
            return
        except HTTPRequestError as e:
            response = {'code': e.code, 'msg' : e.msg}
        except ValueError as e:
            response = {'code': 400, 'msg': 'ValueError : ' + str(e)}
        except KeyError as e:
            response = {'code': 400, 'msg': 'KeyError : ' + str(e)}

        jsonp_callback = self._extract_callback()
        if jsonp_callback:
            content_type = 'application/javascript'
            response_fmt = jsonp_callback + '({0});'
        else:
            content_type = 'application/json'
            response_fmt = '{0}'

        self.send_response(200)
        self.send_header('Content-Type', content_type)
        self.end_headers()

        response = json.dumps(response)
        self.wfile.write(response_fmt.format(response))

    def _extract_post_map(self):
        content_type,_t = cgi.parse_header(self.headers.getheader('content-type'))
        if content_type != 'application/json':
            raise HTTPRequestError(
                    'Bad content-type, use application/json',
                    400)
        length = int(self.headers.getheader('content-length'))
        try:
            return json.loads(self.rfile.read(length))
        except ValueError as e:
            raise HTTPRequestError(
                    'Bad or malformed json content',
                    400)

    def _extract_query_map(self):
        query = urlparse.urlparse(self.path).query
        qd = urlparse.parse_qs(query)
        args = {}
        for k,v in qd.iteritems():
            if len(v) != 1:
                raise HTTPRequestError(
                    "Query param specified multiple times : " + k,
                    400)
            args[k.lower()] = v[0]
        return args

    def _extract_callback(self):
        try:
            args = self._extract_query_map()
            return args['callback']
        except:
            return ''

    def do_POST(self):
        try:
            args = self._extract_post_map() 
        except TypeError as e:
            # thrown on no content, just continue on
            args = '{}'
        except HTTPRequestError as e:
            self.send_error(e.code, e.msg)
            return
        self._serve(args)

    def do_GET(self):
        try:
            args = self._extract_query_map() 
        except HTTPRequestError as e:
            self.send_error(e.code, e.msg)
            return
        self._serve(args)


"""
API handlers for IDA

"""
def check_ea(f):
    def wrapper(self, args):
        if 'ea' in args:
            try:
                ea = int(args['ea'], 16)
            except ValueError:
                raise IDARequestError(
                        'ea parameter malformed - must be 0xABCD', 400)
            if ea > idc.MaxEA():
                raise IDARequestError(
                        'ea out of range - MaxEA is 0x%x' % idc.MaxEA(), 400)
            args['ea'] = ea
        return f(self, args)
    return wrapper

def check_color(f):
    def wrapper(self, args):
        if 'color' in args:
            color = args['color']
            try:
                color = color.lower().lstrip('#').rstrip('h')
                if color.startswith('0x'):
                    color = color[2:]
                # IDA Color is BBGGRR, we need to convert from RRGGBB
                color = color[-2:] + color[2:4] + color[:2]
                color = int(color, 16)
            except:
                raise IDARequestError(
                        'color parameter malformed - must be RRGGBB form', 400)
            args['color'] = color
        return f(self, args)
    return wrapper

def require_params(*params):
    def decorator(f):
        def wrapped(self, args):
            for x in params:
                if x not in args:
                    raise IDARequestError('missing parameter {0}'.format(x), 400)
            return f(self, args)
        return wrapped
    return decorator

class IDARequestError(HTTPRequestError):
    pass

class IDARequestHandler(HTTPRequestHandler):
    @staticmethod
    def _hex(v):
        return hex(v).rstrip('L')

    @HTTPRequestHandler.route('info')
    def info(self, args):
        # No args, Return everything we can meta-wise about the ida session
        # file crcs
        result = {
                'md5' : idc.GetInputMD5(),
                'idb_path' : idc.GetIdbPath(),
                'file_path' : idc.GetInputFilePath(),
                'ida_dir' : idc.GetIdaDirectory(),
                'min_ea' : self._hex(idc.MinEA()),
                'max_ea' : self._hexidc.MaxEA()),
                'segments' : self.segments({})['segments'],
                # idaapi.cvar.inf
                'procname' : idc.GetLongPrm(idc.INF_PROCNAME),
            }
        return result

    @HTTPRequestHandler.route('query')
    @check_ea
    def query(self, args):
        # multiple modes
        # with address return everything about that address
        # with name, return everything about that name
        return {}


    @HTTPRequestHandler.route('cursor')
    @check_ea
    def cursor(self, args):
        # XXX - Doesn't work
        #if 'window' in args:
        #    tform = idaapi.find_tform(args['window'])
        #    if tform:
        #        idaapi.switchto_tform(tform, 1)
        #    else:
        #        raise IDARequestError(
        #            'invalid window - {0}'.format(args['window']), 400)
        result = {}
        if 'ea' in args:
            ea = args['ea']
            success = idaapi.jumpto(ea)
            result['moved'] = success
        result['ea'] = self._hex(idaapi.get_screen_ea())
        return result

    def _get_segment_info(self, s):
        return {
            'name' : idaapi.get_true_segm_name(s),
            'ida_name' : idaapi.get_segm_name(s),
            'start' : self._hex(s.startEA),
            'end' : self._hex(s.endEA),
            'size' : self._hex(s.size())
        }

    @HTTPRequestHandler.route('segments')
    @check_ea
    def segments(self, args):
        if 'ea' in args:
            s = idaapi.getseg(args['ea'])
            if not s:
                raise IDARequestError('Invalid address', 400)
            return {'segment': self._get_segment_info(s)}
        else:
            m = {'segments': []}
            for i in range(idaapi.get_segm_qty()):
                s = idaapi.getnseg(i)
                m['segments'].append(self._get_segment_info(s))
            return m

    @HTTPRequestHandler.route('names')
    def names(self, args):
        m = {'names' : []}
        for n in idautils.Names():
            m['names'].append([self._hex(n[0]), n[1]])
        return m

    @HTTPRequestHandler.route('color')
    @check_color
    @check_ea
    @require_params('ea')
    def color(self, args):
        ea = args['ea']
        if 'color' in args:
            color = args['color']
            def f():
                idc.SetColor(ea, idc.CIC_ITEM, color)
            idaapi.execute_sync(f, idaapi.MFF_WRITE)
            idc.Refresh()
            return {}
        else:
            return {'color' : str(GetColor(ea, idc.CIC_ITEM))}

        
# Figure out when this is really needed
#def f():
#    idaapi.jumpto(ea) # DO STUFF
#idaapi.execute_sync(f, idaapi.MFF_FAST)

"""
Threaded HTTP Server and Worker

Use a worker thread to manage the server so that we can run inside of
IDA Pro without blocking execution.

"""
class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
    allow_reuse_address = True


class Worker(threading.Thread):
    def __init__(self, host='127.0.0.1', port=8899):
        threading.Thread.__init__(self)
        self.httpd = ThreadedHTTPServer((host, port), IDARequestHandler)

    def run(self):
        self.httpd.serve_forever()

    def stop(self):
        self.httpd.shutdown()

"""
IDA Pro Plugin Interface

Define an IDA Python plugin required class and function.
"""

MENU_PATH = 'Edit/Other'
class idarest_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = ""
    help = "IDA Rest API for basic RE tool interoperability"
    wanted_name = "IDA Rest API"
    wanted_hotkey = "Alt-7"

    def _add_menu(self, *args):
        idaapi.msg("Adding menu item\n")
        ctx = idaapi.add_menu_item(*args)
        if ctx is None:
            idaapi.msg("Add failed!\n")
            return False
        else:
            self.ctxs.append(ctx)
            return True

    def _add_menus(self):
        ret = []
        ret.append(
            self._add_menu(MENU_PATH, 'Stop IDARest', '', 1, self.stop, tuple()))
        ret.append(
            self._add_menu(MENU_PATH, 'Start IDARest', '', 1, self.start, tuple()))
        if False in ret:
            return idaapi.PLUGIN_SKIP
        else:
            return idaapi.PLUGIN_KEEP


    def init(self):
        idaapi.msg("Initializing %s\n" % self.wanted_name)
        self.ctxs = []
        self.worker = None
        self.port = 8899
        self.host = '127.0.0.1'
        ret = self._add_menus()
        idaapi.msg("Init done\n")
        return ret

    def _get_netinfo(self):
        info = idaapi.askstr(0,
                "{0}:{1}".format(self.host, self.port),
                "Enter IDA Rest Connection Info")
        if not info:
            raise ValueError("User canceled")
        host,port = info.split(':')
        port = int(port)
        return host,port

    def start(self, *args):
        idaapi.msg("Starting IDARest\n")
        if self.worker:
            idaapi.msg("Already running\n")
            return
        try:
            self.host,self.port = self._get_netinfo()
        except:
            pass

        try:
            self.worker = Worker(self.host,self.port)
        except Exception as e:
            idaapi.msg("Error starting worker : " + str(e) + "\n")
            return
        self.worker.start()
        idaapi.msg("Worker running\n")

    def stop(self, *args):
        idaapi.msg("Stopping IDARest\n")
        if self.worker:
            self.worker.stop()
            del self.worker
            self.worker = None

    def run(self, arg):
        pass

    def term(self):
        idaapi.msg("Terminating %s\n" % self.wanted_name)
        try:
            self.stop()
        except:
            pass
        for ctx in self.ctxs:
            idaapi.del_menu_item(ctx)

def PLUGIN_ENTRY():
    return idarest_plugin_t()


repository: tmr232/Sark
file: plugins/lca.py
description: IDAPython Made Easy
url: https://github.com/tmr232/Sark/blob/950e2c17e520ab637a3cdcd0da312d224b951f13/plugins/lca.py
wanted_name: Lowest Common Ancestors
comment: Lowest Common Ancestors
created_at: 2015-01-19 09:13:48+00:00
pushed_at: 2025-02-22 10:17:16+00:00
content:
from contextlib import suppress
import sark
import sark.graph
import networkx as nx
import sark.ui
import idc
import idaapi

COLOR_SOURCE = 0x364b00
COLOR_TARGET = 0x601116
COLOR_DISABLED = 0x000673
COLOR_PATH = 0x004773
COLOR_SOURCE_TARGET = 0x634107

PADDING = 1
PAD_WIDTH = 3


def pad(text, padding=PADDING):
    top_bottom = ("\n" * padding) + " "
    right_left = " " * padding * PAD_WIDTH
    return top_bottom + right_left + text + right_left + top_bottom


def remove_target_handler(lca_viewer):
    class RemoveTargetHandler(sark.ui.ActionHandler):
        TEXT = "Remove Target"

        def _activate(self, ctx):
            node_id = lca_viewer.current_node_id
            lca_viewer.remove_target(lca_viewer[node_id])
            lca_viewer.rebuild_graph()
            lca_viewer.Refresh()
            idaapi.msg("[LCA] Target Removed: {}\n".format(idaapi.get_ea_name(lca_viewer[node_id], idaapi.GN_VISIBLE)))

    return RemoveTargetHandler


def disable_source_handler(lca_viewer):
    class DisableSourceHandler(sark.ui.ActionHandler):
        TEXT = "Disable Source"

        def _activate(self, ctx):
            node_id = lca_viewer.current_node_id
            lca_viewer.disable_source(lca_viewer[node_id])
            lca_viewer.rebuild_graph()
            lca_viewer.Refresh()

    return DisableSourceHandler


def enable_source_handler(lca_viewer):
    class EnableSourceHandler(sark.ui.ActionHandler):
        TEXT = "Enable Source"

        def _activate(self, ctx):
            node_id = lca_viewer.current_node_id
            lca_viewer.enable_source(lca_viewer[node_id])
            lca_viewer.rebuild_graph()
            lca_viewer.Refresh()

    return EnableSourceHandler


def add_address_handler(lca_viewer):
    class AddAddressHandler(sark.ui.ActionHandler):
        TEXT = "Add Target Function"
        HOTKEY = "Shift+Space"

        def _activate(self, ctx):
            ea = idaapi.ask_long(0, "Add LCA Target")
            if ea is None:
                return

            lca_viewer.add_target(ea)
            lca_viewer.rebuild_graph()
            lca_viewer.Refresh()

    return AddAddressHandler


def add_function_handler(lca_viewer):
    class AddFunctionHandler(sark.ui.ActionHandler):
        TEXT = "Add Target Function"
        HOTKEY = "Space"

        def _activate(self, ctx):
            func = idaapi.choose_func("Add LCA Target", 0)
            if not func:
                return

            lca_viewer.add_target(func.start_ea)
            lca_viewer.rebuild_graph()
            lca_viewer.Refresh()
            lca_viewer.Show()

    return AddFunctionHandler


class LCAGraph(idaapi.GraphViewer):
    def __init__(self, title):
        self._title = title
        idaapi.GraphViewer.__init__(self, self._title)

        self._targets = set()
        self._sources = set()

        # This might take a while...
        self._idb_graph = sark.graph.get_idb_graph()

        self._lca_graph = nx.DiGraph()

        self._handlers = [add_function_handler(self),
                          add_address_handler(self)]

        self._current_node_id = 0

        self._disabled_sources = set()

        self._remove_target_handler = remove_target_handler(self)
        self._enable_source_handler = enable_source_handler(self)
        self._disable_source_handler = disable_source_handler(self)

        self._node_ids = {}

    @property
    def current_node_id(self):
        return self._current_node_id

    def OnGetText(self, node_id):
        return pad(idaapi.get_ea_name(self[node_id], idaapi.GN_VISIBLE))

    def _register_handlers(self):
        for handler in self._handlers:
            handler.register()

    def Show(self):
        if not idaapi.GraphViewer.Show(self):
            return False

        self._register_handlers()
        self.color_nodes()
        return True

    def disable_source(self, source):
        self._disabled_sources.add(source)

    def enable_source(self, source):
        self._disabled_sources.remove(source)

    def add_target(self, target):
        if target not in self._idb_graph.node:
            idaapi.msg(
                "[LCA] Target {} not in IDB graph. Cannot add.\n".format(idaapi.get_ea_name(target, idaapi.GN_VISIBLE)))
            raise KeyError("Target {} not in IDB graph.".format(idaapi.get_ea_name(target, idaapi.GN_VISIBLE)))

        self._targets.add(target)

    def remove_target(self, target):
        self._targets.remove(target)

    def add_targets(self, targets):
        for target in targets:
            self.add_target(target)

    def rebuild_graph(self):
        self._sources = sark.graph.lowest_common_ancestors(self._idb_graph, self._targets)

        # Remove disabled sources from the connected graph.
        active_sources = self._sources - self._disabled_sources
        if active_sources:
            self._lca_graph = sark.graph.get_lca_graph(self._idb_graph, self._targets, active_sources)
        else:
            self._lca_graph = nx.DiGraph()
            self._lca_graph.add_nodes_from(self._targets)

        # Make sure the disabled sources are still shown.
        self._lca_graph.add_nodes_from(self._disabled_sources)

    def _set_node_bg_color(self, node_id, bg_color):
        node_info = idaapi.node_info_t()
        node_info.bg_color = bg_color
        self.SetNodeInfo(node_id, node_info, idaapi.NIF_BG_COLOR)

    def color_nodes(self):
        self.clear_nodes()
        for node_id, node_ea in enumerate(self):
            if node_ea in self._targets and node_ea in self._sources:
                self._set_node_bg_color(node_id, COLOR_SOURCE_TARGET)

            elif node_ea in self._disabled_sources:
                self._set_node_bg_color(node_id, COLOR_DISABLED)

            elif node_ea in self._targets:
                self._set_node_bg_color(node_id, COLOR_TARGET)

            elif node_ea in self._sources:
                self._set_node_bg_color(node_id, COLOR_SOURCE)

    def clear_nodes(self):
        for node_id in range(self.Count()):
            self._set_node_bg_color(node_id, 0xFFFFFFFF)

    def OnRefresh(self):
        self.Clear()

        if self._targets and self._lca_graph is None:
            # This might take a while...
            self.rebuild_graph()

        try:
            lca_graph_nodes_iter = self._lca_graph.nodes()
            lca_graph_edges_iter = self._lca_graph.edges()
        except AttributeError:
            lca_graph_nodes_iter = self._lca_graph.nodes_iter()
            lca_graph_edges_iter = self._lca_graph.edges_iter()

        node_ids = {node: self.AddNode(node) for node in lca_graph_nodes_iter}

        self._node_ids = node_ids

        for frm, to in lca_graph_edges_iter:
            self.AddEdge(node_ids[frm], node_ids[to])

        self.color_nodes()

        return True

    def OnActivate(self):
        # Refresh on every activation to make sure the names are up to date.
        self.Refresh()
        self._register_handlers()
        self.color_nodes()
        return True

    def _unregister_handlers(self):
        for handler in self._handlers:
            handler.unregister()

    def OnDeactivate(self):
        self._unregister_handlers()

    def OnDblClick(self, node_id):
        # On double-click, jump to the clicked address.
        idaapi.jumpto(self[node_id])

        return True

    def _attach_to_popup(self, action_name):
        idaapi.attach_action_to_popup(self.GetTCustomControl(), None, action_name)

    def OnClick(self, node_id):
        self.color_nodes()
        self._current_node_id = node_id
        node_ea = self[node_id]

        self._remove_target_handler.unregister()
        self._disable_source_handler.unregister()
        self._enable_source_handler.unregister()

        if node_ea in self._targets:
            self._remove_target_handler.register()
            self._attach_to_popup(self._remove_target_handler.get_name())

            for ea in nx.ancestors(self._lca_graph, node_ea):
                if ea not in self._targets and ea not in self._sources:
                    self._set_node_bg_color(self._node_ids[ea], COLOR_PATH)

        if node_ea in self._sources:
            if node_ea in self._disabled_sources:
                self._enable_source_handler.register()
                self._attach_to_popup(self._enable_source_handler.get_name())
            else:
                self._disable_source_handler.register()
                self._attach_to_popup(self._disable_source_handler.get_name())

                for ea in nx.descendants(self._lca_graph, node_ea):
                    if ea not in self._targets and ea not in self._sources:
                        self._set_node_bg_color(self._node_ids[ea], COLOR_PATH)

        return False


def lca_viewer_starter(lca_plugin):
    class LCAViewerStarter(sark.ui.ActionHandler):
        TEXT = "LCA Graph"
        TOOLTIP = "Show an interactive lowest-common-ancestors graph."

        def _activate(self, ctx):
            lca_plugin.show_graph()

    return LCAViewerStarter


def idaview_add_target_handler(lca_plugin):
    class IDAViewAddTargetHandler(sark.ui.ActionHandler):
        TEXT = "Add LCA Target"

        def _activate(self, ctx):
            if lca_plugin._lca_viewer:
                with suppress(KeyError):
                    lca_plugin._lca_viewer.add_target(ctx.cur_ea)
                    lca_plugin._lca_viewer.rebuild_graph()
                    idaapi.msg("[LCA] Target Added: {}\n".format(idaapi.get_ea_name(ctx.cur_ea, idaapi.GN_VISIBLE)))

    return IDAViewAddTargetHandler


def idaview_hooks(idaview_handler):
    class Hooks(idaapi.UI_Hooks):
        def finish_populating_widget_popup(self, form, popup):
            if idaapi.get_widget_type(form) == idaapi.BWN_DISASM:
                idaapi.attach_action_to_popup(form, popup, idaview_handler.get_name(), "")

    return Hooks


class LCA(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "Lowest Common Ancestors"
    help = "Lowest Common Ancestors"
    wanted_name = "Lowest Common Ancestors"
    wanted_hotkey = ""

    def init(self):
        self._lca_viewer = None

        self._lca_starter = lca_viewer_starter(self)
        self._lca_starter.register()
        idaapi.attach_action_to_menu("View/Graph Overview", self._lca_starter.get_name(), idaapi.SETMENU_APP)

        self._idaview_handler = idaview_add_target_handler(self)
        self._idaview_handler.register()
        self._hooks = idaview_hooks(self._idaview_handler)()
        self._hooks.hook()

        return idaapi.PLUGIN_KEEP

    def term(self):
        self._lca_starter.unregister()
        self._idaview_handler.unregister()
        self._hooks.unhook()

    def show_graph(self):
        if self._lca_viewer is None:
            self._lca_viewer = LCAGraph("LCA Graph")
        self._lca_viewer.Show()

    def run(self, arg):
        self.show_graph()


def PLUGIN_ENTRY():
    return LCA()


repository: PwCUK-CTO/ScatterBee_Analysis
file: Scripts/ScatterJump.py
description: Scripts to aid analysis of files obfuscated with ScatterBee.
wanted_name: ScatterJumper
comment: ScatterJumper
created_at: 2021-11-03 12:58:22+00:00
pushed_at: 2023-01-06 14:39:45+00:00
url:
https://github.com/PwCUK-CTO/ScatterBee_Analysis/blob/a8711200622717cc64335c0b4714303692348c05/Scripts/ScatterJump.py

content:
"""
Description: IDA Python processor module extension to hook ScatterBee calls and fixup jump destinations
Author: @malworms
License:
Copyright 2021 PwC UK
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.
"""

import idaapi
import idautils
import idc
import binascii

def AddOffset(addr, offset, is64bit):
        if is64bit:
            if offset & 0x80000000:
                offset |= 0xffffffff00000000
            addr += offset
            addr &= 0xffffffffffffffff
            return addr
        else:
            addr += offset
            addr &= 0xffffffff
            return addr

class ScatterHook(idaapi.IDP_Hooks):
    def __init__(self, tramp_offset):
        idaapi.IDP_Hooks.__init__(self)
        self.is64bit = idaapi.get_inf_structure().is_64bit
        self.tramp_offset = tramp_offset

    def ev_ana_insn(self, insn):
        b = idaapi.get_byte(insn.ea)
        if b == 0xE8: # call
            offset = idaapi.get_dword(insn.ea+1)
            if AddOffset(insn.ea+5, offset, self.is64bit) == self.tramp_offset:# the call is going to end up going to the dispatcher function
                insn.itype = idaapi.NN_jmp
                insn.Op1.type = idaapi.o_near
                insn.Op1.flags = 8
                insn.Op1.dtype = idaapi.dt_dword
                insn.Op1.specval = 0x1e0000#relative jump and not an absolute jump
                if self.is64bit:
                    insn.Op1.addr = (insn.ea+5+idaapi.get_qword(insn.ea+5)) & 0xffffffffffffffff
                else:
                    insn.Op1.addr = (idaapi.get_dword(insn.ea+5) + insn.ea + 5) & 0xffffffff
                insn.size = 9
                return True
        return False

def GetAllCallLocations():
    start_ea = idaapi.cvar.inf.minEA
    end_ea = idaapi.cvar.inf.maxEA
    call_loc = idaapi.find_binary(start_ea, end_ea, "e8", 16, idaapi.SEARCH_DOWN)
    call_locs = []
    while call_loc != idaapi.BADADDR:
        call_locs.append(call_loc)
        call_loc = idaapi.find_binary(call_loc+1, end_ea, "e8", 16, idaapi.SEARCH_DOWN)
    return call_locs

def GetDestination(call_addr):
    return AddOffset(call_addr+5, idaapi.get_dword(call_addr+1), idaapi.get_inf_structure().is_64bit())

#for each call location, find the most common call destination and check it starts with an xchg instruction
#return the location if there are > 1000 calls, most code will not call the same function 1500 times
def NeedScatterJumper():
    if idaapi.find_binary(idaapi.cvar.inf.minEA, idaapi.cvar.inf.maxEA, "47 6f 20 62 75 69 6c 64 69 6e 66 3a", 16, 0) != idaapi.BADADDR:
        #We found a Go build string that is in both gobfuscated and none gobfuscated binaries, don't use the decoder.
        return idaapi.BADADDR
    max_val = 0
    max_key = idaapi.BADADDR
    calls = GetAllCallLocations()
    dests = {}
    for call in calls:
        dest = GetDestination(call)
        if dest in dests:
            cur_val = dests.get(dest)
            dests[dest] = cur_val+1
            if cur_val+1 > max_val:
                max_val = cur_val+1
                max_key = dest
        else:
            dests[dest] = 1
    print ("Trampoline function entry: %08x" % max_key)
    #now print out if the start bytes are unusual, if they are the user can see the bytes and decide to decline to load the plugin
    if idaapi.get_inf_structure().is_64bit():
        if idaapi.get_bytes(max_key, 2) != b"\x4c\x87" or idaapi.get_bytes(max_key, 2) != b"\x48\x87":
            print ("found unknown start bytes: ", idaapi.get_bytes(max_key, 5))
    else:
        if idaapi.get_bytes(max_key, 1) != b"\x87":
            print ("found unknown start bytes: ", idaapi.get_bytes(max_key, 5))
    #check there were lots of calls to this function
    #this check is quite naive, but will only fp in weird circumstances and the user can still overrule it later.
    #go binaries may be particularly susceptible to FPs but we have ruled them out by default with the first check.
    if max_val > 1500:
        return max_key
    return idaapi.BADADDR

#core IDA processor module extension code
class ScatterHook_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "ScatterJumper"
    wanted_hotkey = ""
    help = "Runs transparently"
    wanted_name = "ScatterJumper"
    hook = None

    def init(self):
        self.hook = None
        jump_loc = NeedScatterJumper()
        if jump_loc != idaapi.BADADDR:
            if idaapi.ask_yn(0, "Use the ScatterJump plugin?"):
                print ("######## Loading ScatterJumper ########")
                self.hook = ScatterHook(jump_loc)
                self.hook.hook()
                return idaapi.PLUGIN_KEEP
        return idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass

    def term(self):
        if self.hook:
            self.hook.unhook()

def PLUGIN_ENTRY():
    return ScatterHook_t()


repository: Bobo1239/RustSigGen
file: ida-rust-plugin/ida-rust-plugin.py
wanted_name: Rust Signature Generator
created_at: 2024-12-01 23:08:40+00:00
pushed_at: 2025-01-16 13:21:35+00:00
url:
https://github.com/Bobo1239/RustSigGen/blob/eefb3da1eea1bc3779c3911fa6b9e206615c1a3d/ida-rust-plugin/ida-rust-plugin.py

content:
import asyncio
import configparser
import logging
import os
from pathlib import Path

import ida_diskio
import ida_funcs
import ida_idaapi
import ida_nalt
import qasync
from PyQt5.QtWidgets import QApplication

import ida_rust_plugin


class RustSignatureGenerator(ida_idaapi.plugmod_t):
    def run(self, arg):
        print("Starting signature generation...")
        asyncio.create_task(self.generate_and_apply_signature())

    async def generate_and_apply_signature(self):
        # Get the path of the original binary
        # Alternatively we could also get the bytes from the .idb; e.g. https://github.com/polymorf/findcrypt-yara/blob/42a02818586732a1ebd1bce8a9e241b1c46781a0/findcrypt3.py#L212
        bin_path = ida_nalt.get_input_file_path()
        signatures_path = Path(ida_diskio.get_user_idadir()).joinpath("sig", "pc")
        os.makedirs(signatures_path, exist_ok=True)

        config = configparser.ConfigParser()
        config.read(os.path.join(os.path.dirname(__file__), "ida-rust-plugin.cfg"))
        flair_path = config["ida-rust-plugin"]["flair_path"]

        # NOTE: Errors from the Rust side do get propagated here
        sig_path = await ida_rust_plugin.generate_signature_for_bin(
            bin_path, flair_path, signatures_path
        )
        ida_funcs.plan_to_apply_idasgn(sig_path)


class RustSignatureGeneratorPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_UNL | ida_idaapi.PLUGIN_MULTI
    help = "TODO"
    wanted_name = "Rust Signature Generator"

    def init(self):
        LOG_FORMAT = "[%(levelname)s] %(message)s"
        logging.basicConfig(format=LOG_FORMAT, level=logging.INFO)

        # Set asyncio event loop to the QT one; without this the GUI freezes during script runtime and
        # strange Rust crashes occur
        asyncio.set_event_loop(
            qasync.QEventLoop(QApplication.instance(), already_running=True)
        )

        return RustSignatureGenerator()


def PLUGIN_ENTRY():
    return RustSignatureGeneratorPlugin()


repository: burpheart/Gepetto-ChatGPT
file: Gepetto-ChatGPT.py
description: IDA plugin which queries OpenAI's ChatGPT model to speed up reverse-engineering
wanted_name: Gepetto_CHAT
comment: Uses ChatGPT to enrich the decompiler's output
created_at: 2022-12-08 05:04:22+00:00
pushed_at: 2022-12-08 10:14:06+00:00
url:
https://github.com/burpheart/Gepetto-ChatGPT/blob/1ab4ee0b0e3b0cdfaad47141b6dc6a8a4216d260/Gepetto-ChatGPT.py

content:
import functools
import json
import idaapi
import ida_hexrays
import ida_kernwin
import idc
# import openai
import os
import re
import textwrap
import threading
from revChatGPT.revChatGPT import Chatbot

config = {
     "email": "<YOUR_EMAIL>",
     "password": "<YOUR_PASSWORD>"#,
    #"session_token": "",
     #Use session_token or email/password. But the session_token has a very short validity
    #"proxy": "127.0.0.1:7890"
}
ZH_CN = True  # 是否使用中文代码解释 # Use Chinese explain


# =============================================================================
# Setup the context menu and hotkey in IDA
# =============================================================================

class Gepetto_CHATPlugin(idaapi.plugin_t):
    flags = 0
    explain_action_name = "Gepetto_CHAT:explain_function_CHAT"
    explain_menu_path = "Edit/Gepetto_CHAT/Explain function_CHAT"
    rename_action_name = "Gepetto_CHAT:rename_function_CHAT"
    rename_menu_path = "Edit/Gepetto_CHAT/Rename variables_CHAT"
    wanted_name = 'Gepetto_CHAT'
    wanted_hotkey = ''
    comment = "Uses ChatGPT to enrich the decompiler's output"
    help = "See usage instructions on GitHub"
    menu = None

    def init(self):
        # Check whether the decompiler is available
        if not ida_hexrays.init_hexrays_plugin():
            return idaapi.PLUGIN_SKIP

        # Function explaining action
        explain_action = idaapi.action_desc_t(self.explain_action_name,
                                              'Explain function_CHAT',
                                              ExplainHandler(),
                                              "Ctrl+Alt+G",
                                              'Use ChatGPT to explain the currently selected function',
                                              199)
        idaapi.register_action(explain_action)
        idaapi.attach_action_to_menu(self.explain_menu_path, self.explain_action_name, idaapi.SETMENU_APP)

        # Variable renaming action
        rename_action = idaapi.action_desc_t(self.rename_action_name,
                                             'Rename variables_CHAT',
                                             RenameHandler(),
                                             "Ctrl+Alt+R",
                                             "Use ChatGPT to rename this function's variables",
                                             199)
        idaapi.register_action(rename_action)
        idaapi.attach_action_to_menu(self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)

        # Register context menu actions
        self.menu = ContextMenuHooks()
        self.menu.hook()

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        idaapi.detach_action_from_menu(self.explain_menu_path, self.explain_action_name)
        idaapi.detach_action_from_menu(self.rename_menu_path, self.rename_action_name)
        if self.menu:
            self.menu.unhook()
        return


# -----------------------------------------------------------------------------

class ContextMenuHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        # Add actions to the context menu of the Pseudocode view
        if idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE:
            idaapi.attach_action_to_popup(form, popup, Gepetto_CHATPlugin.explain_action_name, "Gepetto_CHAT/")
            idaapi.attach_action_to_popup(form, popup, Gepetto_CHATPlugin.rename_action_name, "Gepetto_CHAT/")


# -----------------------------------------------------------------------------

def comment_callback(address, view, response):
    """
    Callback that sets a comment at the given address.
    :param address: The address of the function to comment
    :param view: A handle to the decompiler window
    :param response: The comment to add
    """
    # Add newlines at the end of each sentence.
    response = "\n".join(textwrap.wrap(response, 80, replace_whitespace=False))

    # Add the response as a comment in IDA.
    idc.set_func_cmt(address, response, 0)
    # Refresh the window so the comment is displayed properly
    if view:
        view.refresh_view(False)
    print("ChatGPT query finished!")


# -----------------------------------------------------------------------------

class ExplainHandler(idaapi.action_handler_t):
    """
    This handler is tasked with querying ChatGPT for an explanation of the
    given function. Once the reply is received, it is added as a function
    comment.
    """

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

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        if ZH_CN:
            query_model_async(
                "对下面的C语言伪代码函数进行分析 推测关于该函数的使用环境和预期目的详细的函数功能等信息 并为这个函数取一个新的名字 不要返回其他的内容 (开始前加上GPTSTART 结束后加上GPTEND字符串)\n"
                + str(decompiler_output),
                functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
        else:
            query_model_async(
                "Can you explain what the following C function does and suggest a better name for it?(Add GPTSTART before the beginning of the conversation and GPTEND after the end.)\n"
                + str(decompiler_output),
                functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


# -----------------------------------------------------------------------------

def rename_callback(address, view, response):
    """
    Callback that extracts a JSON array of old names and new names from the
    response and sets them in the pseudocode.
    :param address: The address of the function to work on
    :param view: A handle to the decompiler window
    :param response: The response from ChatGPT
    """
    j = re.search(r"\{[^}]*?\}", response)
    if not j:
        print(f"Error: couldn't extract a response from ChatGPT's output:\n{response}")
        return
    try:
        names = json.loads(j.group(0))
    except json.decoder.JSONDecodeError:
        print(f"The data returned by the model cannot be parsed. Asking the model to fix it...")
        query_model_async("Please fix the following JSON document:\n" + j.group(0),
                          functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=view))
        return

    # The rename function needs the start address of the function
    function_addr = idaapi.get_func(address).start_ea

    replaced = []
    for n in names:
        if ida_hexrays.rename_lvar(function_addr, n, names[n]):
            replaced.append(n)

    # Update possible names left in the function comment
    comment = idc.get_func_cmt(address, 0)
    if comment and len(replaced) > 0:
        for n in replaced:
            comment = re.sub(r'\b%s\b' % n, names[n], comment)
        idc.set_func_cmt(address, comment, 0)

    # Refresh the window to show the new names
    if view:
        view.refresh_view(True)
    print(f"ChatGPT query finished! {len(replaced)} variable(s) renamed.")


# -----------------------------------------------------------------------------

class RenameHandler(idaapi.action_handler_t):
    """
    This handler requests new variable names from ChatGPT and updates the
    decompiler's output.
    """

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

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async(
            "Analyze the following C function. Suggest better variable names, reply with a JSON array where keys are the original names and values are the proposed names. Do not explain anything, only print the JSON dictionary(Add GPTSTART before the beginning of the reply and GPTEND after the end.):\n" + str(
                decompiler_output),
            functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


# =============================================================================
# ChatGPT interaction
# =============================================================================

def query_model(query, cb):
    """
    Function which sends a query to ChatGPT and calls a callback when the response is available.
    Blocks until the response is received
    :param query: The request to send to ChatGPT
    :param cb: Tu function to which the response will be passed to.
    """
    try:
        chatbot = Chatbot(config, conversation_id=None)
        response = chatbot.get_chat_response(query)['message']
        if response.find("GPTSTART") == -1:
            raise Exception("Unexpected response: " + response)
        times = 1
        retry = 0
        data = response
        print(f"response[" + str(times) + "]: " + response)
        while response.find("GPTEND") == -1:
            try:
                times += 1
                response = chatbot.get_chat_response("next")['message']
                if response.find("GPTSTART") != -1:
                    times = 99
                    raise Exception("Duplicate responses appear: " + response)
                print(f"response[" + str(times) + "]: " + response)
                data += response
                # print(message)
                times = times - retry
                retry = 0
            except Exception as e:
                if times > 5:
                    raise Exception("Request 5 times and still not return full results: " + response)
                if retry > 3:
                    raise Exception("Retry 3 times and the request still fails: " + response)
                retry += 1
        ida_kernwin.execute_sync(functools.partial(cb, response=data.replace('GPTEND', '').replace('GPTSTART', '')),
                                 ida_kernwin.MFF_WRITE)
    except Exception as e:
        print(f"General exception encountered while running the query: {str(e)}")


# -----------------------------------------------------------------------------

def query_model_async(query, cb):
    """
    Function which sends a query to ChatGPT and calls a callback when the response is available.
    :param query: The request to send to ChatGPT
    :param cb: Tu function to which the response will be passed to.
    """
    print("Request to ChatGPT sent...")
    t = threading.Thread(target=query_model, args=[query, cb])
    t.start()


# =============================================================================
# Main
# =============================================================================

def PLUGIN_ENTRY():
    return Gepetto_CHATPlugin()


repository: 19h/ida-plugins
file: xrefdump/xrefdump.py
wanted_name: CallGraph Exporter
comment: Exports a Graphviz DOT file representing all direct function calls.
created_at: 2025-04-08 22:54:27+00:00
pushed_at: 2025-04-08 23:07:17+00:00
url:
https://github.com/19h/ida-plugins/blob/b045b6f6bc9ae2a53104986fb74f3fe1198e5d2e/xrefdump/xrefdump.py

content:
import idaapi
import idautils
import idc
import ida_kernwin

def export_call_graph():
    """
    Enumerates all functions in the current IDB, inspects each instruction for direct call
    opcodes, and for each call with an immediate operand referencing a valid function,
    records an edge from the caller to the callee. The complete set of edges is then written
    to a Graphviz DOT file.
    """
    outfile = ida_kernwin.ask_file(False, "*.dot", "Select output DOT file")
    if not outfile:
        print("No output file selected.")
        return

    edges = set()
    # Iterate over all function entry points in the database.
    for func_ea in idautils.Functions():
        src_name = idc.get_func_name(func_ea)
        if not src_name:
            continue
        # Iterate over all instruction addresses in the function.
        for ea in idautils.FuncItems(func_ea):
            mnem = idc.print_insn_mnem(ea)
            if mnem.lower() != "call":
                continue
            # Ensure operand type is immediate (o_near); indirect calls are skipped.
            if idc.get_operand_type(ea, 0) != idc.o_near:
                continue
            target = idc.get_operand_value(ea, 0)
            # Use idaapi.get_func() instead of idc.get_func() (not available in IDA Pro 9.0)
            if idaapi.get_func(target) is None:
                continue
            tgt_name = idc.get_func_name(target)
            if not tgt_name:
                continue
            edges.add((src_name, tgt_name))

    try:
        with open(outfile, "w") as f:
            f.write("digraph CallGraph {\n")
            for (src, tgt) in sorted(edges):
                f.write('    "{}" -> "{}";\n'.format(src, tgt))
            f.write("}\n")
        print("[*] Call graph exported to: {}".format(outfile))
    except Exception as e:
        print("Error writing DOT file: {}".format(e))


class CallGraphExportPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "Exports a Graphviz DOT file representing all direct function calls."
    help = ("This plugin iterates over all functions, inspects call instructions with immediate operands, "
            "and exports the call relationships in Graphviz DOT format.")
    wanted_name = "CallGraph Exporter"
    wanted_hotkey = ""

    def init(self):
        print("[*] CallGraph Exporter plugin initialized.")
        return idaapi.PLUGIN_OK

    def run(self, arg):
        export_call_graph()

    def term(self):
        print("[*] CallGraph Exporter plugin terminated.")


def PLUGIN_ENTRY():
    return CallGraphExportPlugin()


repository: poohzaza166/pooh-plugin-assistant
file: plugins/zetime/main.py
description: This is a showcase of pytohn plugin system for creating voice assistant
created_at: 2025-03-01 22:01:41+00:00
pushed_at: 2025-03-23 16:19:46+00:00
url:
https://github.com/poohzaza166/pooh-plugin-assistant/blob/21decaa3651a0191e1006927a31c097dd3c8bde3/plugins/zetime/main.py

content:
import datetime
import time

import pytz

from .pooh_lib import (PluginBase, loop_method, provide_data, push_event,
                       register_command, subscribe)


class TimeManagement(PluginBase):
    def __init__(self, message_bus, parser):
        super().__init__(message_bus, parser)
        self.timers = {}
        self.timezone = "UTC"
        
    @register_command(
        patterns=[r"what(?:'s|s| is) the (?:current )?time", r"time now", r"current time"],
        example_phrase={"time_query": "What's the current time?"}
    )
    def get_current_time(self, format: str = "%H:%M:%S") -> str:
        """Get the current time in the specified format"""
        current_time = datetime.datetime.now().strftime(format)
        return f"The current time is {current_time}"
    
    @register_command(
        patterns=[r"set (?:a|the) timer for (\d+) (second|minute|hour|day)s?", 
                  r"remind me in (\d+) (second|minute|hour|day)s?"],
        example_phrase={"timer": "Set a timer for 5 minutes", 
                       "reminder": "Remind me in 30 minutes"}
    )
    def set_timer(self, duration: int, unit: str, label: str = "Timer") -> str:
        """Set a timer or reminder for the specified duration"""
        # Convert everything to seconds
        unit_multipliers = {
            "second": 1,
            "minute": 60,
            "hour": 3600,
            "day": 86400
        }
        
        seconds = duration * unit_multipliers.get(unit.lower(), 1)
        timer_id = str(time.time())
        end_time = time.time() + seconds
        
        self.timers[timer_id] = {
            "end_time": end_time,
            "label": label,
            "duration": f"{duration} {unit}{'s' if duration > 1 else ''}"
        }
        
        return f"{label} set for {duration} {unit}{'s' if duration > 1 else ''}"
        
    @loop_method(delay=1.0)
    def check_timers(self):
        """Check for completed timers and notify user"""
        current_time = time.time()
        completed_timers = []
        
        for timer_id, timer_info in self.timers.items():
            if current_time >= timer_info["end_time"]:
                completed_timers.append(timer_id)
                self.message_bus.publish("notification", 
                                        f"{timer_info['label']} for {timer_info['duration']} is complete!")
        
        # Remove completed timers
        for timer_id in completed_timers:
            del self.timers[timer_id]
    
    @register_command(
        patterns=[r"convert (\d+:\d+) from (\w+) to (\w+)", 
                  r"what time is (\d+:\d+) (\w+) in (\w+)"]
    )
    def convert_timezone(self, time_str: str, from_zone: str, to_zone: str) -> str:
        """Convert time between different timezones"""
        try:
            # Parse the time string
            time_format = "%H:%M"
            time_obj = datetime.datetime.strptime(time_str, time_format)
            
            # Get timezone objects
            from_tz = pytz.timezone(from_zone)
            to_tz = pytz.timezone(to_zone)
            
            # Set the date to today
            today = datetime.datetime.now().date()
            time_obj = datetime.datetime.combine(today, time_obj.time())
            
            # Localize and convert
            localized_time = from_tz.localize(time_obj)
            converted_time = localized_time.astimezone(to_tz)
            
            return f"{time_str} in {from_zone} is {converted_time.strftime(time_format)} in {to_zone}"
        except (ValueError, pytz.exceptions.UnknownTimeZoneError):
            return f"Sorry, I couldn't convert between {from_zone} and {to_zone}. Please check timezone names."
    
    @register_command(
        patterns=[r"set (?:default )?timezone to (\w+/\w+)"]
    )
    def set_timezone(self, timezone: str) -> str:
        """Set the default timezone for time-related commands"""
        try:
            # Validate timezone
            pytz.timezone(timezone)
            self.timezone = timezone
            return f"Default timezone set to {timezone}"
        except pytz.exceptions.UnknownTimeZoneError:
            return f"Unknown timezone: {timezone}. Please use a valid timezone like 'America/New_York'."
    
    @subscribe("text_input")
    def check_for_time_queries(self, text):
        """Listen for simple time queries that might not match exact commands"""
        time_keywords = ["time", "clock", "hour", "minute", "second"]
        
        if any(keyword in text.lower() for keyword in time_keywords):
            # This is a time-related query, but might not match our patterns exactly
            # We could use a more advanced NLP approach here or just default to current time
            current_time = datetime.datetime.now().strftime("%H:%M:%S")
            self.message_bus.publish("response", f"It's currently {current_time}")


def plugin_entry(message_bus, parser):
    print("loading plugin")
    print("checking parser object")
    print(parser)
    plugin = TimeManagement(message_bus, parser)
    return  plugin


repository: Skwteinopteros/ida-venv
file: ida_venv.py
description: An IDA Pro plugin to allow use of Python venvs
wanted_name: IDAVenv
comment: Plugin to help with using venvs in IDA Pro.
created_at: 2024-01-25 05:18:23+00:00
pushed_at: 2024-10-04 05:54:06+00:00
url:
https://github.com/Skwteinopteros/ida-venv/blob/27c39f2b523cb19968b34aaf527e8f5846e569be/ida_venv.py

content:
import argparse
import json
import os
import runpy
import site
import subprocess
import sys
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import NotRequired, TypedDict
from venv import create

import ida_idaapi
import ida_kernwin
import ida_loader
import pkg_resources

PLUGIN_NAME = "IDAVenv"

PYTHON_VERSION = f"python{sys.version_info.major}.{sys.version_info.minor}"

BIN_PATH = {
    "nt": "Scripts",
    "posix": "bin",
}

LIB_PATH = {
    "nt": "Lib/site-packages",
    "posix": f"lib/{PYTHON_VERSION}/site-packages",
}

PYTHON_EXECUTABLE_PATH = {
    "nt": "Scripts/python.exe",
    "posix": "bin/python",
}


def _activate_venv(venv_path: Path) -> None:
    base_path = str(venv_path.resolve())
    venv_dict = {
        "VIRTUAL_ENV": base_path,
        "_OLD_VIRTUAL_PATH": os.environ.pop("PATH", ""),
        "_OLD_SYS_PATH": sys.path,
        "_OLD_PREFIX": sys.prefix,
    }

    try:
        bin_dir = f"{base_path}/{BIN_PATH[os.name]}"
    except KeyError:
        # java/jython
        raise RuntimeError("Jython not supported (yet?)") from None

    # add venv to list
    venvs = json.loads(os.environ.get("IDAVenvs", "[]"))
    venvs.append(venv_dict)
    os.environ["IDAVenvs"] = json.dumps(venvs)

    # save old path
    os.environ["_OLD_VIRTUAL_PATH"] = venv_dict["_OLD_VIRTUAL_PATH"]

    # prepend bin to PATH (this file is inside the bin directory)
    os.environ["PATH"] = os.pathsep.join(
        [bin_dir] + os.environ.get("PATH", "").split(os.pathsep)
    )
    os.environ["VIRTUAL_ENV"] = venv_dict["VIRTUAL_ENV"]

    # add the virtual environments libraries to the host python import mechanism
    prev_length = len(sys.path)

    lib = f"{base_path}/{LIB_PATH[os.name]}"
    path = os.path.realpath(os.path.join(bin_dir, lib))
    site.addsitedir(path)

    sys.path[:] = sys.path[prev_length:] + sys.path[0:prev_length]
    sys.prefix = sys.exec_prefix = base_path


def create_venv(venv_path: Path) -> None:
    with get_context():
        create(venv_path, system_site_packages=True, symlinks=False, with_pip=True)


def activate_venv(venv_path: Path, dependencies: list[str] | None = None) -> None:
    print("[IDAVenv] activate_venv")
    python_path = venv_path / PYTHON_EXECUTABLE_PATH[os.name]
    if not python_path.is_file():
        create_venv(venv_path)

    _activate_venv(venv_path)

    if dependencies:
        install_dependencies(venv_path=venv_path, dependencies=dependencies)


def deactivate_venv() -> None:
    print("[IDAVenv] deactivate_venv")
    venv_dicts = json.loads(os.environ.get("IDAVenvs", "[]"))
    if not venv_dicts:
        return

    venv_dict = venv_dicts.pop()
    if not venv_dict:
        return

    # restore old path
    os.environ["PATH"] = os.environ.pop("_OLD_VIRTUAL_PATH", "")
    sys.path = venv_dict["_OLD_SYS_PATH"]
    # restore prefix
    sys.prefix = sys.exec_prefix = venv_dict["_OLD_PREFIX"]
    os.environ["IDAVenvs"] = json.dumps(venv_dicts)

    # TODO! improve this
    to_remove = []
    for name, module in sys.modules.items():
        if not module:
            continue
        try:
            if module.__file__.startswith(venv_dict["VIRTUAL_ENV"]):
                to_remove.append(name)
        except AttributeError:
            pass

    for name in to_remove:
        del sys.modules[name]


def install_dependencies(
    venv_path: Path, dependencies: list[str], *, pip_args: list[str] | None = None
) -> None:
    if not dependencies:
        return

    python_path = f"{venv_path}/{PYTHON_EXECUTABLE_PATH[os.name]}"

    cmd = [python_path, "-m", "pip", "install"]
    cmd.extend(dependencies)
    if pip_args:
        cmd.extend(pip_args)
    subprocess.run(cmd, check=True)


@contextmanager
def change_executable(executable: str):
    _executable = sys.executable

    sys.executable = executable
    sys._base_executable = sys.executable
    yield
    sys._base_executable = sys.executable = _executable


def get_context():
    if os.name == "nt":
        executable = Path(sys.base_prefix, "python.exe")
    elif os.name == "posix":
        # On Linux we use system python
        executable = Path(sys.base_prefix, "bin", "python3")
    else:
        # java/jython
        raise RuntimeError("Jython not supported (yet?)")
    return change_executable(str(executable))


@contextmanager
def venv_context(venv_path: Path, dependencies: list[str] | None = None):
    activate_venv(venv_path=venv_path, dependencies=dependencies)
    yield
    deactivate_venv()


def run_script_in_env(
    script_path: str | None = None,
    venv_path: Path | None = None,
    dependencies: list[str] | None = None,
) -> None:
    if not script_path:
        script_path = ida_kernwin.ask_file(
            0,
            "*.py|*.*",
            "Select script",
        )
    if not script_path:
        return

    script = Path(script_path).resolve()
    if not script.is_file():
        return

    parent_dir = os.environ.get("IDAVENV_VENV_DIR", str(script.parent))

    if not venv_path:
        venv_path = Path(parent_dir, f".venvs/{script.stem}")
    with venv_context(venv_path=venv_path, dependencies=dependencies):
        sys.path.append(parent_dir)
        runpy.run_path(str(script), run_name="__main__")


def _validate_file(filename):
    path = Path(filename).resolve()
    if path.is_file():
        return path

    raise argparse.ArgumentTypeError(f"{filename} is not a valid file")


@dataclass
class ModuleArgsNamespace:
    venv: Path | None = None
    dependencies: list[str] | None = None
    requirements_file: Path | None = None


def _parse_module_args(args: list[str]) -> ModuleArgsNamespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--venv", type=Path)

    dependencies_group = parser.add_mutually_exclusive_group()
    dependencies_group.add_argument("-r", "--requirements-file", type=_validate_file)
    dependencies_group.add_argument("-d", "--dependencies", nargs="*", type=str)

    namespace = ModuleArgsNamespace()
    parsed_args, _ = parser.parse_known_args(args, namespace=namespace)

    return parsed_args


def _process_module_args(args: ModuleArgsNamespace):
    if not args.venv:
        return

    dependencies = []
    if args.dependencies:
        dependencies = args.dependencies
    elif args.requirements_file:
        dependencies = [
            str(dep)
            for dep in pkg_resources.parse_requirements(
                args.requirements_file.read_text()
            )
        ]

    activate_venv(venv_path=args.venv, dependencies=dependencies)


class ActionDeclaration(TypedDict):
    name: str
    label: str
    handler: ida_kernwin.action_handler_t
    shortcut: NotRequired[str]
    tooltip: NotRequired[str]
    icon: NotRequired[int]
    flags: NotRequired[int]
    menu_location: str


class ActionHandler(ida_kernwin.action_handler_t):
    def __init__(self, name: str, callback) -> None:
        ida_kernwin.action_handler_t.__init__(self)
        self._name = name
        self._callback = callback

    def activate(self, ctx) -> None:
        self._callback()

    def update(self, ctx) -> bool:
        return ida_kernwin.AST_ENABLE_ALWAYS


def register_action(action: ActionDeclaration):
    ida_action = ida_kernwin.action_desc_t(
        action["name"],
        action["label"],
        action["handler"],
        action.get("shortcut", ""),
        action.get("tooltip", ""),
        action.get("icon", 0),
        action.get("flags", 0),
    )
    if not ida_kernwin.register_action(ida_action):
        print("Failed to register {action['id']}")

    if not ida_kernwin.attach_action_to_menu(
        action["menu_location"],  # The menu location
        action["name"],  # The unique function ID
        ida_kernwin.SETMENU_APP,  # Flags
    ):
        print(f"Failed to attach to menu { action['name']}")


def register_actions(actions: list[ActionDeclaration]):
    for action in actions:
        register_action(action=action)


class IDAVenvPlugin(ida_idaapi.plugin_t):
    # Use the HIDE flag to avoid the entry in
    # Edit/Plugins since this plugin's run()
    # method has no functionality...it's all
    # in the actions.

    flags = ida_idaapi.PLUGIN_HIDE
    comment = "Plugin to help with using venvs in IDA Pro."
    help = " See https://github.com/Skwteinopteros/ida-venv?tab=readme-ov-file#usage"
    wanted_name = "IDAVenv"
    wanted_hotkey = ""
    actions = [
        ActionDeclaration(
            name=f"{PLUGIN_NAME}:run",
            label="Script file in venv",
            handler=ActionHandler(
                name=f"{PLUGIN_NAME}:run", callback=run_script_in_env
            ),
            shortcut="Ctrl+Alt+F7",
            tooltip="Runs script in venv",
            menu_location="File/Script file...",
        )
    ]

    def init(self):
        print(f"[{PLUGIN_NAME}] initialization start")

        register_actions(actions=self.actions)
        args = ida_loader.get_plugin_options(PLUGIN_NAME)
        if args:
            parsed_args = _parse_module_args(args=args.split(":"))
            _process_module_args(args=parsed_args)

        print(f"[{PLUGIN_NAME}] initialization complete")
        # Return KEEP instead of OK to keep the
        # plugin loaded since it registers
        # callback actions and hotkeys
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        print(f"[{PLUGIN_NAME}] run (no effect)")

    def term(self):
        print(f"[{PLUGIN_NAME}] terminated")
        deactivate_venv()


def PLUGIN_ENTRY():
    plugin = IDAVenvPlugin()

    return plugin


repository: xthebat/optolithium
file: OptolithiumGui/pcpi.py
description: Optical lithography simulation software
created_at: 2020-03-20 07:43:25+00:00
pushed_at: 2023-04-09 10:37:01+00:00
url:
https://github.com/xthebat/optolithium/blob/b3b38a93fdcf5aa84aba306b83b346077de716c7/OptolithiumGui/pcpi.py

content:
# This file is part of Optolithium lithography modelling software.
#
# Copyright (C) 2015 Alexei Gladkikh
#
# This software is dual-licensed: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version only for NON-COMMERCIAL usage.
#
# 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
#
# If you are interested in other licensing models, including a commercial-
# license, please contact the author at gladkikhalexei@gmail.com


from ctypes import POINTER, CFUNCTYPE, Structure, cast
from ctypes import c_void_p, c_char_p, c_double, c_int
import inspect


__author__ = 'Alexei Gladkikh'


class PluginCommonError(Exception):
    pass


class PluginNotFoundError(PluginCommonError):
    pass


class PluginRegistry(object):

    def __init__(self):
        self.__container_id = dict()
        self.__container_name = dict()

    def get_by_id(self, plugin_id):
        """:rtype: plugins.Plugin"""
        try:
            return self.__container_id[plugin_id]
        except KeyError:
            raise PluginNotFoundError("Plugin with id %s not found" % plugin_id)

    def get_by_name(self, plugin_name):
        """:rtype: plugins.Plugin"""
        try:
            return self.__container_name[plugin_name]
        except KeyError:
            raise PluginNotFoundError("Plugin with name %s not found" % plugin_name)

    def add_plugin(self, plugin):
        """:type plugin: plugins.Plugin"""
        self.__container_id[plugin.record.id] = plugin
        self.__container_name[plugin.record.name] = plugin


# Plugin registry must be in pcpi module not in plugins module otherwise it result in circular imports
PLUGIN_REGISTRY = PluginRegistry()


class CPluginInterface(Structure):
    plugin_id = None
    """:type: int"""
    name = None
    """:type: str"""
    desc = None
    """:type: str"""


def deref(p):
    return p.contents.value if p else None


def ro_property(attribute):
    return property(fget=lambda self: getattr(self, attribute))


def rw_property(attribute):
    return property(fget=lambda self: getattr(self, attribute),
                    fset=lambda self, value: setattr(self, attribute, value))


# noinspection PyPep8Naming
class _variable_part_t(Structure):
    _fields_ = [
        ("_name", c_char_p),
        ("_defv", c_double),
        ("_min", POINTER(c_double)),
        ("_max", POINTER(c_double))
    ]

    name = property(lambda self: self._name)
    """:type: str"""
    defv = property(lambda self: self._defv)
    """:type: float"""
    min = property(lambda self: deref(self._min))
    """:type: float or None"""
    max = property(lambda self: deref(self._max))
    """:type: float or None"""

    def __str__(self):
        vmin = "Min %.2f, " % self.min if self.min is not None else str()
        vmax = "Max %.2f" % self.max if self.max is not None else str()
        return "Name: \"%s\", %s%s Default: %.2f" % (self.name, vmin, vmax, self.defv)


# ------------------------------------------ Development model Plugin --------------------------------------------------


# noinspection PyPep8Naming
class dev_model_arg_t(_variable_part_t):
    pass


# Return value: double
# Parameters: pac, depth, values
dev_model_expr_t = CFUNCTYPE(c_double, c_double, c_double, POINTER(c_double))


# noinspection PyPep8Naming
class dev_model_t(CPluginInterface):
    plugin_id = 1

    _fields_ = [
        ("_prolith_id", POINTER(c_int)),
        ("_name", c_char_p),
        ("_desc", c_char_p),
        ("_expr", dev_model_expr_t),
        ("_args_count", c_int),
        ("_args", POINTER(dev_model_arg_t))
    ]

    prolith_id = property(lambda self: deref(self._prolith_id))
    """:type: int or None"""
    name = property(lambda self: self._name)
    """:type: str"""
    desc = property(lambda self: self._desc)
    """:type: str"""
    expr = property(lambda self: self._expr)
    """:type: dev_model_expr_t"""
    args_count = property(lambda self: self._args_count)
    """:type: int"""
    args = property(lambda self: self._args)
    """:type: list of dev_model_arg_t"""

    def __str__(self):
        args = "\n".join(["\t%d. %s" % (k, self.args[k]) for k in xrange(self.args_count)])
        return ("Name: \"%s\"\nDesc: %s\nExpr: %s\nCount: %d\nProlith Id: %s\nArgs:\n%s" %
                (self.name, self.desc, self.expr, self.args_count, self.prolith_id, args))


# ------------------------------------------------ Mask Plugin ---------------------------------------------------------


# noinspection PyPep8Naming
class point_t(Structure):

    _fields_ = [
        ("_x", c_double),
        ("_y", c_double)
    ]

    x = rw_property("_x")
    """:type: float"""
    y = rw_property("_y")
    """:type: float"""


# noinspection PyPep8Naming
class mask_region_t(Structure):

    _fields_ = [
        ("_transmittance", c_double),
        ("_phase", c_double),
        ("_length", c_int),
        ("_points", POINTER(point_t))
    ]

    transmittance = rw_property("_transmittance")
    """:type: float"""
    phase = rw_property("_phase")
    """:type: float"""
    length = ro_property("_length")
    """:type: int"""
    points = ro_property("_points")
    """:type: list of point_t"""


# noinspection PyPep8Naming
class mask_parameter_t(_variable_part_t):
    pass


# noinspection PyPep8Naming
class mask_t(Structure):

    _fields_ = [
        ("_boundary", mask_region_t),
        ("_regions_count", c_int),
        ("_regions", POINTER(mask_region_t)),
    ]

    boundary = ro_property("_boundary")
    """:type: mask_region_t"""
    regions_count = ro_property("_regions_count")
    """:type: int"""
    regions = ro_property("_regions")
    """:type: list of mask_region_t"""


# Return value: int - status (0 - Ok)
# Parameters: mask, parameters
mask_create_t = CFUNCTYPE(c_int, POINTER(mask_t), POINTER(c_double))


# noinspection PyPep8Naming
class mask_plugin_t(CPluginInterface):
    plugin_id = 0

    _fields_ = [
        ("_name", c_char_p),
        ("_desc", c_char_p),
        ("_dimensions", c_int),
        ("_create", mask_create_t),
        ("_parameters_count", c_int),
        ("_parameters", POINTER(mask_parameter_t))
    ]

    name = ro_property("_name")
    """:type: str"""
    desc = ro_property("_desc")
    """:type: str"""
    dimensions = ro_property("_dimensions")
    """:type: int"""
    create = ro_property("_create")
    """:type: mask_create_t"""
    parameters_count = ro_property("_parameters_count")
    """:type: int"""
    parameters = ro_property("_parameters")
    """:type: mask_parameter_t"""


# --------------------------------------------- Source Shape Plugin ----------------------------------------------------

# noinspection PyPep8Naming
class source_shape_parameter_t(_variable_part_t):
    pass


# Return value: double - intensity value
# Parameters: direction cosine x, y, and parameters
source_shape_expr_t = CFUNCTYPE(c_double, c_double, c_double, POINTER(c_double))


# noinspection PyPep8Naming
class source_shape_plugin_t(CPluginInterface):
    plugin_id = 2

    _fields_ = [
        ("_name", c_char_p),
        ("_desc", c_char_p),
        ("_expression", source_shape_expr_t),
        ("_parameters_count", c_int),
        ("_parameters", POINTER(source_shape_parameter_t))
    ]

    name = ro_property("_name")
    """:type: str"""
    desc = ro_property("_desc")
    """:type: str"""
    expr = ro_property("_expression")
    """:type: source_shape_expr_t"""
    parameters_count = ro_property("_parameters_count")
    """:type: int"""
    parameters = ro_property("_parameters")
    """:type: source_shape_parameter_t"""


# --------------------------------------------- Pupil Filter Plugin ----------------------------------------------------

# noinspection PyPep8Naming
class pupil_filter_parameter_t(_variable_part_t):
    pass


# noinspection PyPep8Naming
class c_complex(Structure):
    _fields_ = [("_real", c_double), ("_imag", c_double)]
    real = rw_property("_real")
    imag = rw_property("_imag")


# Return value: double - diffraction term coefficient value
# Parameters: direction cosine x, y, and parameters
pupil_filter_expr_t = CFUNCTYPE(c_complex, c_double, c_double, POINTER(c_double))


# noinspection PyPep8Naming
class pupil_filter_plugin_t(CPluginInterface):
    plugin_id = 5

    _fields_ = [
        ("_name", c_char_p),
        ("_desc", c_char_p),
        ("_expression", pupil_filter_expr_t),
        ("_parameters_count", c_int),
        ("_parameters", POINTER(pupil_filter_parameter_t))
    ]

    name = ro_property("_name")
    """:type: str"""
    desc = ro_property("_desc")
    """:type: str"""
    expr = ro_property("_expression")
    """:type: pupil_filter_expr_t"""
    parameters_count = ro_property("_parameters_count")
    """:type: int"""
    parameters = ro_property("_parameters")
    """:type: pupil_filter_parameter_t"""


# ----------------------------------------------------------------------------------------------------------------------


plugin_types = {cls.plugin_id: cls for cls in globals().values()
                if inspect.isclass(cls) and issubclass(cls, CPluginInterface) and cls.plugin_id is not None}


ENTRY_POINT = "PluginDescriptor"


# noinspection PyPep8Naming
class plugin_descriptor_t(Structure):
    _fields_ = [
        ("_plugin_type", c_int),
        ("_plugin_entry", c_void_p)
    ]

    plugin_type = property(lambda self: self._plugin_type)
    """:type: int"""

    @property
    def plugin_entry(self):
        """:rtype: CPluginInterface"""
        return cast(self._plugin_entry, POINTER(plugin_types[self.plugin_type])).contents


repository: grayhatacademy/ida
file: plugins/fluorescence/fluorescence.py
wanted_name: fluorescence
comment: Highlights function calls
created_at: 2013-10-16 18:52:21+00:00
pushed_at: 2023-10-23 10:38:19+00:00
url:
https://github.com/grayhatacademy/ida/blob/faf13a13c4f5070fe8f0666c1508f928ffe40e23/plugins/fluorescence/fluorescence.py

content:
# Plugin to highlight and un-highlight call instructions.
import idc
import idaapi
import idautils

from shims import ida_shims


class CallHighlighter(object):
    COLOR = 0xFF99FF  # BBGGRR

    def highlight(self):
        for ea in idautils.Heads():
            flags = ida_shims.get_full_flags(ea)
            if ida_shims.is_code(flags) and idaapi.is_call_insn(ea):
                current_color = idaapi.get_item_color(ea)
                if current_color == self.COLOR:
                    idaapi.set_item_color(ea, idc.DEFCOLOR)
                elif current_color == idc.DEFCOLOR:
                    idaapi.set_item_color(ea, self.COLOR)


try:
    class FluorescenceActionHandler(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

        def activate(self, ctx):
            CallHighlighter().highlight()
            return 1

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


class fluorescence_blower_t(idaapi.plugin_t):
    flags = 0
    comment = "Highlights function calls"
    help = ''
    wanted_name = 'fluorescence'
    wanted_hotkey = ''
    fluorescence_action_name = 'fluorescence:action'
    menu_tab = 'Options/'
    menu_name = 'Un/highlight call instructions'
    context_menu = None

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            fluorescence_desc = idaapi.action_desc_t(
                self.fluorescence_action_name, self.menu_name,
                FluorescenceActionHandler(), self.wanted_hotkey,
                'Highlights function calls.', 199)

            idaapi.register_action(fluorescence_desc)
            idaapi.attach_action_to_menu(self.menu_tab,
                                         self.fluorescence_action_name,
                                         idaapi.SETMENU_APP)
        else:
            self.context_menu = idaapi.add_menu_item(self.menu_tab,
                                                     self.menu_name,
                                                     "",
                                                     0,
                                                     self.highlight,
                                                     (None,))
        return idaapi.PLUGIN_KEEP

    def term(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.detach_action_from_menu(self.menu_tab,
                                           self.fluorescence_action_name)
        else:
            if self.context_menu is not None:
                idaapi.del_menu_item(self.context_menu)
        return None

    def run(self, arg):
        pass

    def highlight(self, arg):
        CallHighlighter().highlight()


def PLUGIN_ENTRY():
    return fluorescence_blower_t()


repository: tbxy09/mdnav
file: ftplugin/markdown/mdnav.py
description: Vim plugin for markdown markup suggestion
created_at: 2018-09-04 06:04:19+00:00
pushed_at: 2018-09-10 09:05:43+00:00
url:
https://github.com/tbxy09/mdnav/blob/4e34a1c09fae873d04c75e7c7d8d77bd12636c40/ftplugin/markdown/mdnav.py

content:
from __future__ import print_function

import collections
import json
import os.path
import re
import sys
import subprocess
import webbrowser

try:
    from urllib.parse import urlparse
except ImportError:
    from urlparse import urlparse


class FakeLogger(object):
    def __init__(self, active=False):
        self.active = active

    def info(self, fmt, *args):
        if not self.active:
            return

        print(fmt % args)


_logger = FakeLogger()


def plugin_entry_point():
    import vim

    if int(vim.eval("exists('g:mdnav#Extensions')")):
        extensions = vim.eval('g:mdnav#Extensions')
        extensions = [ext.strip() for ext in extensions.split(',')]

    else:
        extensions = []

    if int(vim.eval("exists('g:mdnav#DebugMode')")):
        _logger.active = vim.eval('g:mdnav#DebugMode') == 'true'

    row, col = vim.current.window.cursor
    cursor = (row - 1, col)
    lines = vim.current.buffer

    target = parse_link(cursor, lines)
    _logger.info('open %s', target)
    action = open_link(
        target,
        current_file=vim.eval("expand('%:p')"),
        open_in_vim_extensions=extensions,
    )
    action()


def open_link(target, current_file, open_in_vim_extensions=set()):
    """
    :returns: a callable that encapsulates the action to perform
    """
    if target is not None:
        target = target.strip()

    if not target:
        _logger.info('no target')
        return NoOp(target)

    if target.startswith('#'):
        return JumpToAnchor(target)

    if has_scheme(target):
        _logger.info('has scheme -> open in browser')
        return BrowserOpen(target)

    if not has_extension(target, open_in_vim_extensions):
        _logger.info('has no extension for opening in vim')
        return OSOpen(anchor_path(target, current_file))

    if target.startswith('|filename|'):
        target = target[len('|filename|'):]

    return VimOpen(anchor_path(target, current_file))


def anchor_path(target, current_file):
    if os.path.isabs(target):
        return target

    _logger.info('anchor path relative to %s', current_file)
    return os.path.join(os.path.dirname(current_file), target)


def has_extension(path, extensions):
    if not extensions:
        return True

    path = parse_path(path)
    _, ext = os.path.splitext(path.path)
    return ext in extensions


def has_scheme(target):
    return bool(urlparse(target).scheme)


class Action(object):
    def __init__(self, target):
        self.target = target

    def __eq__(self, other):
        return type(self) == type(other) and self.target == other.target

    def __repr__(self):
        return '{}({!r})'.format(type(self).__name__, self.target)


class NoOp(Action):
    def __call__(self):
        print('<mdnav: no link>')


class BrowserOpen(Action):
    def __call__(self):
        print('<mdnav: open browser tab>')
        webbrowser.open_new_tab(self.target)


class OSOpen(Action):
    def __call__(self):
        if sys.platform.startswith('linux'):
            call(['xdg-open', self.target])

        elif sys.platform.startswith('darwin'):
            call(['open', self.target])

        else:
            os.startfile(self.target)


class VimOpen(Action):
    def __call__(self):
        import vim

        path = parse_path(self.target)

        # TODO: make space handling more robust?
        vim.command('e {}'.format(path.path.replace(' ', '\\ ')))
        if path.line is not None:
            try:
                line = int(path.line)

            except:
                print('invalid line number')
                return

            else:
                vim.current.window.cursor = (line, 0)

        if path.anchor is not None:
            JumpToAnchor(path.anchor)()


class JumpToAnchor(Action):
    heading_pattern = re.compile(r'^#+(?P<title>.*)$')
    attr_list_pattern = re.compile(r'{:\s+#(?P<id>\S+)\s')

    def __call__(self):
        import vim
        line = self.find_anchor(self.target, vim.current.buffer)

        if line is None:
            return

        vim.current.window.cursor = (line + 1, 0)

    @classmethod
    def find_anchor(cls, target, buffer):
        needle = cls.norm_target(target)

        for (idx, line) in enumerate(buffer):
            m = cls.heading_pattern.match(line)
            if (
                m is not None and
                cls.title_to_anchor(m.group('title')) == needle
            ):
                return idx

            m = cls.attr_list_pattern.search(line)
            if m is not None and needle == m.group('id'):
                return idx

    @staticmethod
    def title_to_anchor(title):
        return '-'.join(fragment.lower() for fragment in title.split())

    @staticmethod
    def norm_target(target):
        if target.startswith('#'):
            target = target[1:]

        return target.lower()


def call(args):
    """If available use vims shell mechanism to work around display issues
    """
    try:
        import vim

    except ImportError:
        subprocess.call(args)

    else:
        args = ['shellescape(' + json.dumps(arg) + ')' for arg in args]
        vim.command('execute "! " . ' + ' . " " . '.join(args))


def parse_path(path):
    """Parse a path with optional line number of anchor into its parts.

    For example::

        parse_path('foo.md:30') == ParsedPath('foo.md', line=30)
        parse_path('foo.md#anchor') == ParsedPath('foo.md', anchor='anchor')

    """
    path, ext = os.path.splitext(path)
    if '#' in ext:
        ext, anchor = ext.rsplit('#', 1)
        return ParsedPath(path + ext, anchor=anchor)

    if ':' in ext:
        ext, line = ext.rsplit(':', 1)
        return ParsedPath(path + ext, line=line)

    return ParsedPath(path + ext)


class ParsedPath(object):
    def __init__(self, path, line=None, anchor=None):
        self.path = path
        self.line = line
        self.anchor = anchor

    def __repr__(self):
        return 'ParsedPath({!r}, line={}, anchor={!r})'.format(self.path, self.line, self.anchor)


def parse_link(cursor, lines):
    row, column = cursor
    line = lines[row]

    _logger.info('handle line %s (%s, %s)', line, row, column)
    m = reference_definition_pattern.match(line)
    if m is not None:
        return m.group('link').strip()

    link_text, rel_column = select_from_start_of_link(line, column)

    if not link_text:
        _logger.info('could not find link text')
        return None

    m = link_pattern.match(link_text)

    if not m:
        _logger.info('does not match link pattern')
        return None

    if m.end('link') <= rel_column:
        _logger.info('cursor outside link')
        return None

    _logger.info('found match: %s', m.groups())
    assert (m.group('direct') is None) != (m.group('indirect') is None)

    if m.group('direct') is not None:
        _logger.info('found direct link: %s', m.group('direct'))
        return m.group('direct')

    _logger.info('follow indirect link %s', m.group('indirect'))
    indirect_ref = m.group('indirect')
    if not indirect_ref:
        indirect_ref = m.group('text')

    indrect_link_pattern = re.compile(
        r'^\[' + re.escape(indirect_ref) + r'\]:(.*)$'
    )

    for line in lines:
        m = indrect_link_pattern.match(line)

        if m:
            return m.group(1).strip()

    _logger.info('could not match for indirect link')
    return None


reference_definition_pattern = re.compile(r'''
    ^
        \[[^\]]*\]:             # reference def at start of line
        (?P<link>.*)            # interpret everything else as link text
    $
''', re.VERBOSE)

link_pattern = re.compile(r'''
    ^
    (?P<link>
        \[                      # start of link text
            (?P<text>[^\]]*)    # link text
        \]                      # end of link text
        (?:
            \(                  # start of target
                (?P<direct>
                    [^\)]*
                )
            \)                  # collect
            |
            \[
                (?P<indirect>
                    [^\]]*
                )
            \]
        )
    )
    .*                  # any non matching characters
    $
''', re.VERBOSE)


def select_from_start_of_link(line, pos):
    """Return the start of the link string and the new cursor
    """
    if pos < len(line) and line[pos] == '[':
        start = pos

    else:
        start = line[:pos].rfind('[')

    # TODO: handle escapes

    if start < 0:
        return None, pos

    # check for indirect links
    if start != 0 and line[start - 1] == ']':
        alt_start = line[:start].rfind('[')
        if alt_start >= 0:
            start = alt_start

    return line[start:], pos - start


if __name__ == "__main__":
    plugin_entry_point()


repository: N1nEmAn/VxHPub
file: ida_tools/ida_vxworks_sym_load.py
description: VxHello for public
wanted_name: Symbol Importer
comment: vxworks Symbol Importer
created_at: 2024-07-30 13:02:24+00:00
pushed_at: 2024-07-30 13:29:22+00:00
url:
https://github.com/N1nEmAn/VxHPub/blob/99f038dae5994f4017aacafe02de520476d85689/ida_tools/ida_vxworks_sym_load.py

content:
import ida_kernwin
import ida_name
import ida_funcs
import ida_idaapi
import struct
import collections


def read_symbols(file: str):
    with open(file, 'rb') as fp:
        total_size = struct.unpack(">I", fp.read(4))[0]
        entry_count = struct.unpack(">I", fp.read(4))[0]
        entry_data = fp.read(entry_count * 8)
        string_data = fp.read(total_size - entry_count * 8)
        if len(string_data) != total_size - entry_count * 8 - 8:
            raise Exception("The file is not symbol table")

        symbol_name_dict = {}
        idx = 0
        for symbol_name in string_data.split(b"\0"):
            symbol_name_dict[idx] = symbol_name.decode()
            idx += len(symbol_name) + 1  # the terminal null

        SymbolEntry = collections.namedtuple("SymbolEntry", "type address name")
        symbol_entries = []
        for i in range(entry_count):
            t, high_bit_idx, low_bit_idx, address = struct.unpack(">BBHI", entry_data[i * 8:i * 8 + 8])
            idx = (high_bit_idx << 16) + low_bit_idx
            name = symbol_name_dict[idx]
            symbol_entries.append(SymbolEntry(t, address, name))
        return symbol_entries


def process_symbols():
    sym_file = ida_kernwin.ask_file(0, "*", "symbol table file")
    if sym_file is None:
        return

    print(f"read symbol file from {sym_file}")

    entries = read_symbols(sym_file)

    for entry in entries:
        ida_name.set_name(entry.address, entry.name, ida_name.SN_CHECK)
        if entry.type == 84 or entry.type == 116:  # "t" and "T"
            ida_funcs.add_func(entry.address)


class SymbolImporterPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_UNL
    comment = "vxworks Symbol Importer"
    help = "Imports symbols from a custom symbol table file"
    wanted_name = "Symbol Importer"
    wanted_hotkey = "Ctrl-Shift-S"

    def init(self):
        print("Symbol Importer Plugin initialized")
        return ida_idaapi.PLUGIN_OK

    def run(self, arg):
        process_symbols()

    def term(self):
        print("Symbol Importer Plugin terminated")


def PLUGIN_ENTRY():
    return SymbolImporterPlugin()


repository: Thehepta/D810Learn
file: D810_demo/vds5.py
wanted_name: Hex-Rays show C graph (IDAPython)
comment: Sample plugin5 for Hex-Rays decompiler
created_at: 2024-06-02 06:36:20+00:00
pushed_at: 2025-05-11 10:47:50+00:00
url:
https://github.com/Thehepta/D810Learn/blob/52ad29245037f279411a2ddd6b8b6e3f15e9bc9a/D810_demo/vds5.py

content:
"""
summary: show ctree graph

description:
  Registers an action that can be used to show the graph of the ctree.
  The current item will be highlighted in the graph.

  The command shortcut is `Ctrl+Shift+G`, and is also added
  to the context menu.

  To display the graph, we produce a .gdl file, and
  request that ida displays that using `ida_gdl.display_gdl`.
"""

from __future__ import print_function

import ida_pro
import ida_hexrays
import ida_kernwin
import ida_gdl
import ida_lines
import ida_idaapi

ACTION_NAME = "vds5.py:displaygraph"
ACTION_SHORTCUT = "Ctrl+Shift+G"

CL_WHITE            = ((255)+  (255<<8)+  (255<<16)) #   0
CL_BLUE             = ((0  )+  (0  <<8)+  (255<<16)) #   1
CL_RED              = ((255)+  (0  <<8)+  (0  <<16)) #   2
CL_GREEN            = ((0  )+  (255<<8)+  (0  <<16)) #   3
CL_YELLOW           = ((255)+  (255<<8)+  (0  <<16)) #   4
CL_MAGENTA          = ((255)+  (0  <<8)+  (255<<16)) #   5
CL_CYAN             = ((0  )+  (255<<8)+  (255<<16)) #   6
CL_DARKGREY         = ((85 )+  (85 <<8)+  (85 <<16)) #   7
CL_DARKBLUE         = ((0  )+  (0  <<8)+  (128<<16)) #   8
CL_DARKRED          = ((128)+  (0  <<8)+  (0  <<16)) #   9
CL_DARKGREEN        = ((0  )+  (128<<8)+  (0  <<16)) #  10
CL_DARKYELLOW       = ((128)+  (128<<8)+  (0  <<16)) #  11
CL_DARKMAGENTA      = ((128)+  (0  <<8)+  (128<<16)) #  12
CL_DARKCYAN         = ((0  )+  (128<<8)+  (128<<16)) #  13
CL_GOLD             = ((255)+  (215<<8)+  (0  <<16)) #  14
CL_LIGHTGREY        = ((170)+  (170<<8)+  (170<<16)) #  15
CL_LIGHTBLUE        = ((128)+  (128<<8)+  (255<<16)) #  16
CL_LIGHTRED         = ((255)+  (128<<8)+  (128<<16)) #  17
CL_LIGHTGREEN       = ((128)+  (255<<8)+  (128<<16)) #  18
CL_LIGHTYELLOW      = ((255)+  (255<<8)+  (128<<16)) #  19
CL_LIGHTMAGENTA     = ((255)+  (128<<8)+  (255<<16)) #  20
CL_LIGHTCYAN        = ((128)+  (255<<8)+  (255<<16)) #  21
CL_LILAC            = ((238)+  (130<<8)+  (238<<16)) #  22
CL_TURQUOISE        = ((64 )+  (224<<8)+  (208<<16)) #  23
CL_AQUAMARINE       = ((127)+  (255<<8)+  (212<<16)) #  24
CL_KHAKI            = ((240)+  (230<<8)+  (140<<16)) #  25
CL_PURPLE           = ((160)+  (32 <<8)+  (240<<16)) #  26
CL_YELLOWGREEN      = ((154)+  (205<<8)+  (50 <<16)) #  27
CL_PINK             = ((255)+  (192<<8)+  (203<<16)) #  28
CL_ORANGE           = ((255)+  (165<<8)+  (0  <<16)) #  29
CL_ORCHID           = ((218)+  (112<<8)+  (214<<16)) #  30
CL_BLACK            = ((0  )+  (0  <<8)+  (0  <<16)) #  31

COLORS_LUT = {
    CL_WHITE             : "white",
    CL_BLUE              : "blue",
    CL_RED               : "red",
    CL_GREEN             : "green",
    CL_YELLOW            : "yellow",
    CL_MAGENTA           : "magenta",
    CL_CYAN              : "cyan",
    CL_DARKGREY          : "darkgrey",
    CL_DARKBLUE          : "darkblue",
    CL_DARKRED           : "darkred",
    CL_DARKGREEN         : "darkgreen",
    CL_DARKYELLOW        : "darkyellow",
    CL_DARKMAGENTA       : "darkmagenta",
    CL_DARKCYAN          : "darkcyan",
    CL_GOLD              : "gold",
    CL_LIGHTGREY         : "lightgrey",
    CL_LIGHTBLUE         : "lightblue",
    CL_LIGHTRED          : "lightred",
    CL_LIGHTGREEN        : "lightgreen",
    CL_LIGHTYELLOW       : "lightyellow",
    CL_LIGHTMAGENTA      : "lightmagenta",
    CL_LIGHTCYAN         : "lightcyan",
    CL_LILAC             : "lilac",
    CL_TURQUOISE         : "turquoise",
    CL_AQUAMARINE        : "aquamarine",
    CL_KHAKI             : "khaki",
    CL_PURPLE            : "purple",
    CL_YELLOWGREEN       : "yellowgreen",
    CL_PINK              : "pink",
    CL_ORANGE            : "orange",
    CL_ORCHID            : "orchid",
    CL_BLACK             : "black",
}

def get_color_name(c):
    return COLORS_LUT[c] if c in COLORS_LUT.keys() else "?"

class cfunc_graph_t: # alas we can't inherit gdl_graph_t
    def __init__(self, highlight):
        self.items = [] # list of citem_t
        self.highlight = highlight
        self.succs = [] # list of lists of next nodes
        self.preds = [] # list of lists of previous nodes

    def nsucc(self, n):
        return len(self.succs[n]) if self.size() else 0

    def npred(self, n):
        return len(self.preds[n]) if self.size() else 0

    def succ(self, n, i):
        return self.succs[n][i]

    def pred(self, n, i):
        return self.preds[n][i]

    def size(self):
        return len(self.preds)

    def add_node(self):
        n = self.size()

        def resize(array, new_size):
            if new_size > len(array):
                while len(array) < new_size:
                    array.append([])
            else:
                array = array[:new_size]
            return array

        self.preds = resize(self.preds, n+1)
        self.succs = resize(self.succs, n+1)
        return n

    def add_edge(self, x, y):
        self.preds[y].append(x)
        self.succs[x].append(y)

    def get_expr_name(self, expr):
        name = expr.print1(None)
        name = ida_lines.tag_remove(name)
        name = ida_pro.str2user(name)
        return name

    def get_node_label(self, n):
        item = self.items[n]
        op = item.op
        insn = item.cinsn
        expr = item.cexpr
        parts = [ida_hexrays.get_ctype_name(op)]
        if op == ida_hexrays.cot_ptr:
            parts.append(".%d" % expr.ptrsize)
        elif op == ida_hexrays.cot_memptr:
            parts.append(".%d (m=%d)" % (expr.ptrsize, expr.m))
        elif op == ida_hexrays.cot_memref:
            parts.append(" (m=%d)" % (expr.m,))
        elif op in [
                ida_hexrays.cot_obj,
                ida_hexrays.cot_var]:
            name = self.get_expr_name(expr)
            parts.append(".%d %s" % (expr.refwidth, name))
        elif op in [
                ida_hexrays.cot_num,
                ida_hexrays.cot_helper,
                ida_hexrays.cot_str]:
            name = self.get_expr_name(expr)
            parts.append(" %s" % (name,))
        elif op == ida_hexrays.cit_goto:
            parts.append(" LABEL_%d" % insn.cgoto.label_num)
        elif op == ida_hexrays.cit_asm:
            parts.append("<asm statements; unsupported ATM>")
            # parts.append(" %a.%d" % ())
        parts.append("\\n")
        parts.append("ea: %08X" % item.ea)
        if item.is_expr() and not expr.type.empty():
            parts.append("\\n")
            tstr = expr.type._print()
            parts.append(tstr if tstr else "?")
        return "".join(parts)

    def get_node_color(self, n):
        item = self.items[n]
        if self.highlight is not None and item.obj_id == self.highlight.obj_id:
            return CL_GREEN
        return None

    def gen_gdl(self, fname):
        with open(fname, "wb") as out:
            def write_out(s):
                out.write(s.encode("UTF-8"))
            write_out("graph: {\n")

            write_out("// *** nodes\n")
            for n in range(len(self.items)):
                item = self.items[n]
                node_label = self.get_node_label(n)
                node_props = [""]
                if n == 0:
                    node_props.append("vertical_order: 0")
                color = self.get_node_color(n)
                if color is not None:
                    node_props.append("color: %s" % get_color_name(color))
                write_out("""node: { title: "%d" label: "%d: %s" %s}\n""" % (
                    n,
                    n,
                    node_label,
                    " ".join(node_props)))

            write_out("// *** edges\n")
            for n in range(len(self.items)):
                item = self.items[n]
                write_out("// edges %d -> ?\n" % n)
                for i in range(self.nsucc(n)):
                    t = self.succ(n, i)
                    label = ""
                    if item.is_expr():
                        target = self.items[t]
                        if item.x is not None and item.x == target:
                            label = "x"
                        elif item.y is not None and item.y == target:
                            label = "y"
                        elif item.z is not None and item.z == target:
                            label = "z"
                        if label:
                            label = """ label: "%s" """ % label
                    write_out("""edge: { sourcename: "%s" targetname: "%s"%s}\n""" % (
                        str(n), str(t), label))

            write_out("}\n")

    def dump(self):
        print("%d items:" % len(self.items))
        for i in self.items:
            print("\t%s (%08x)" % (i, i.ea))

        print("succs:")
        for s in self.succs:
            print("\t%s" % s)

        print("preds:")
        for p in self.preds:
            print("\t%s" % p)


class graph_builder_t(ida_hexrays.ctree_parentee_t):

    def __init__(self, cg):
        ida_hexrays.ctree_parentee_t.__init__(self)
        self.cg = cg
        self.reverse = [] # (citem_t, node#) tuples

    def add_node(self, i):
        for k, _ in self.reverse:
            if i.obj_id == k.obj_id:
                ida_kernwin.warning("bad ctree - duplicate nodes! (i.ea=%x)" % i.ea)
                self.cg.dump()
                return -1

        n = self.cg.add_node()
        if n <= len(self.cg.items):
            self.cg.items.append(i)
        self.cg.items[n] = i
        self.reverse.append((i, n))
        return n

    def process(self, i):
        n = self.add_node(i)
        if n < 0:
            return n
        if len(self.parents) > 1:
            lp = self.parents.back().obj_id
            for k, v in self.reverse:
                if k.obj_id == lp:
                    p = v
                    break
            self.cg.add_edge(p, n)
        return 0

    def visit_insn(self, i):
        return self.process(i)

    def visit_expr(self, e):
        return self.process(e)


class display_graph_ah_t(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def activate(self, ctx):
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        vu.get_current_item(ida_hexrays.USE_KEYBOARD)
        highlight = vu.item.e if vu.item.is_citem() else None

        cg = cfunc_graph_t(highlight)
        gb = graph_builder_t(cg)
        gb.apply_to(vu.cfunc.body, None)

        import tempfile
        fname = tempfile.mktemp(suffix=".gdl")
        cg.gen_gdl(fname)
        ida_gdl.display_gdl(fname)
        return 1

    def update(self, ctx):
        return ida_kernwin.AST_ENABLE_FOR_WIDGET if \
            ctx.widget_type == ida_kernwin.BWN_PSEUDOCODE else \
            ida_kernwin.AST_DISABLE_FOR_WIDGET


class vds5_hooks_t(ida_hexrays.Hexrays_Hooks):
    def populating_popup(self, widget, handle, vu):
        ida_kernwin.attach_action_to_popup(vu.ct, None, ACTION_NAME)
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays show C graph (IDAPython)"
    wanted_hotkey = "Ctrl+Shift+G"
    comment = "Sample plugin5 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            ida_kernwin.register_action(
                ida_kernwin.action_desc_t(
                    ACTION_NAME,
                    "Hex-Rays show C graph (IDAPython)",
                    display_graph_ah_t(),
                    ACTION_SHORTCUT))
            self.vds5_hooks = vds5_hooks_t()
            self.vds5_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds5_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: tmr232/ida-plugin-loader
file: plugin_loader.py
wanted_name: PluginLoader
comment: Plugin Loader
created_at: 2018-07-15 15:48:24+00:00
pushed_at: 2023-12-07 14:14:46+00:00
url:
https://github.com/tmr232/ida-plugin-loader/blob/bd7f829c0c93af99bf7599b91f32e8ba701d99b4/plugin_loader.py

content:
import glob
import os
import itertools
import idaapi
import idc

PLUGINS_LIST = "plugins-{}.list".format(idaapi.get_kernel_version())

USER_PLUGIN_LIST_PATH = os.path.join(idaapi.get_user_idadir(), PLUGINS_LIST)
SYS_PLUGIN_LIST_PATH = os.path.join(idaapi.idadir(idaapi.CFG_SUBDIR), PLUGINS_LIST)
IDB_PATH = idaapi.get_path(idaapi.PATH_TYPE_IDB)
if IDB_PATH:
    PROJECT_PLUGIN_LIST_PATH = os.path.join(os.path.dirname(IDB_PATH), PLUGINS_LIST)
else:
    PROJECT_PLUGIN_LIST_PATH = None


def message(*messages):
    for msg in messages:
        for line in msg.splitlines():
            idaapi.msg("[PluginLoader] {}\n".format(line))


def iter_without_duplicates(*iterables):
    visited = set()
    chained_iterables = itertools.chain(*iterables)
    for item in chained_iterables:
        if item in visited:
            continue
        yield item
        visited.add(item)


def iter_paths(filepath):
    if not filepath:
        return
    try:
        with open(filepath) as f:
            for line in f:
                # Use `#` for comments
                if line.startswith("#"):
                    continue
                # Remove trailing spaces and newlines, then normalize to avoid duplicates.
                path = os.path.normpath(line.strip())
                if path:
                    # Allow for relative paths
                    real_path = os.path.join(os.path.dirname(path), path)
                    yield real_path
    except IOError:
        pass


def iter_plugin_paths():
    glob_patterns = iter_without_duplicates(iter_paths(SYS_PLUGIN_LIST_PATH),
                                            iter_paths(USER_PLUGIN_LIST_PATH),
                                            iter_paths(PROJECT_PLUGIN_LIST_PATH))
    return iter_without_duplicates(*(glob.iglob(pattern) for pattern in glob_patterns))


class PluginLoader(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = "Plugin Loader"
    help = "Plugin Loader"
    wanted_name = "PluginLoader"
    wanted_hotkey = ""

    def init(self):
        # Show usage message.
        usage_message = ["Loading plugins from:",
                         "  System-wide List:      {}".format(SYS_PLUGIN_LIST_PATH),
                         "  User-specific List:    {}".format(USER_PLUGIN_LIST_PATH)]
        if PROJECT_PLUGIN_LIST_PATH:
            usage_message.append("  Project-specific List: {}".format(PROJECT_PLUGIN_LIST_PATH))

        message(*usage_message)

        # Make sure the files exist. If not - create them.
        if not os.path.isfile(SYS_PLUGIN_LIST_PATH):
            try:
                with open(SYS_PLUGIN_LIST_PATH, "wb"):
                    message("Created system plugin list at {}".format(SYS_PLUGIN_LIST_PATH))
            except IOError:
                message("Failed creating system plugin list at {}".format(SYS_PLUGIN_LIST_PATH))

        if not os.path.isfile(USER_PLUGIN_LIST_PATH):
            try:
                with open(USER_PLUGIN_LIST_PATH, "wb"):
                    message("Created user plugin list at {}".format(USER_PLUGIN_LIST_PATH))
            except IOError:
                message("Failed creating user plugin list at {}".format(USER_PLUGIN_LIST_PATH))

        for path in iter_plugin_paths():
            # This check is not needed, but saves us from the dreaded error message-box
            # that pops when a python plugin is not found.
            if not os.path.isfile(path):
                message("Plugin not found: {}".format(path))
                continue
            idaapi.load_plugin(path)
        return idaapi.PLUGIN_SKIP

    def term(self):
        pass

    def run(self, arg):
        pass


def PLUGIN_ENTRY():
    return PluginLoader()


repository: 0rganizers/nmips
file: instr_info.py
description: nanoMIPS IDA plugin
url: https://github.com/0rganizers/nmips/blob/b8d2a984819abd4df7d7555a622c6e4f1aaa4f4a/instr_info.py
wanted_name: instr_info
comment: instr_info v0.0.1
created_at: 2021-06-18 23:49:09+00:00
pushed_at: 2021-09-15 18:16:56+00:00
content:
import idaapi
import ida_funcs
import idautils
import ida_kernwin
import ida_allins
import ida_ua

type_map = {
    idaapi.o_void: "void",
    idaapi.o_reg: "reg",
    idaapi.o_mem: "mem",
    idaapi.o_phrase: "phrase",
    idaapi.o_displ: "displ",
    idaapi.o_imm: "imm",
    idaapi.o_far: "far",
    idaapi.o_near: "near",
    idaapi.o_idpspec0: "idpspec0",
    idaapi.o_idpspec1: "idpspec1",
    idaapi.o_idpspec2: "idpspec2",
    idaapi.o_idpspec3: "idpspec3",
    idaapi.o_idpspec4: "idpspec4",
    idaapi.o_idpspec5: "idpspec5"
}

dt_map = {
    idaapi.dt_byte: "byte",
    idaapi.dt_word: "word",
    idaapi.dt_dword: "dword",
    idaapi.dt_float: "float",
    idaapi.dt_double: "double",
    idaapi.dt_tbyte: "tbyte",
    idaapi.dt_packreal: "packreal",
    idaapi.dt_qword: "qword",
    idaapi.dt_byte16: "byte16",
    idaapi.dt_code: "code",
    idaapi.dt_void: "void",
    idaapi.dt_string: "string"
}


class instr_info_plugin_t(idaapi.plugin_t):
    comment = "instr_info v0.0.1"
    help = "Display information about the current instruction, such as opcode and operands"
    wanted_name = "instr_info"
    wanted_hotkey = "Ctrl+Meta+Shift+I"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        return idaapi.PLUGIN_KEEP

    def cache_mnemonics(self):
        pass

    def format_type(self, t):
        return type_map[t]
    
    def format_dt(self, dt):
        return dt_map[dt]

    def get_value(self, op: idaapi.op_t):
        if op.type == idaapi.o_reg:
            return f"{idaapi.get_reg_name(op.reg, 4)} ({op.reg})"
        if op.type == idaapi.o_phrase:
            return op.phrase
        if op.type == idaapi.o_imm:
            return op.value
        if op.type == idaapi.o_mem or op.type == idaapi.o_near:
            return f"0x{op.addr:x}"
        if op.type == idaapi.o_displ:
            return f"({op.addr}, {op.phrase})"
        return op.specval

    def get_prop(self, op: idaapi.op_t):
        if op.type == idaapi.o_reg:
            return "reg"
        if op.type == idaapi.o_phrase:
            return "phrase"
        if op.type == idaapi.o_imm:
            return "value"
        if op.type == idaapi.o_mem or op.type == idaapi.o_near:
            return "addr"
        if op.type == idaapi.o_displ:
            return "(addr, phrase)"
        return "specval"

    def format_specval(self, op: idaapi.op_t):
        if op.specval == 0:
            return ""
        return f" (spec: {op.specval:x})"

    def format_flags(self, op: idaapi.op_t):
        return f"{op.specflag1}, {op.specflag2}, {op.specflag3}, {op.specflag4}"

    def format_value(self, op: idaapi.op_t):
        return f"{self.get_prop(op)}: {self.get_value(op)}"

    def format_operand(self, i: idaapi.insn_t, op: idaapi.op_t):
        return f"[{self.format_type(op.type)}] {self.format_value(op)} ({self.format_dt(op.dtype)}) ({self.format_flags(op)}){self.format_specval(op)}"

    def run(self, arg):
        ea = idaapi.get_screen_ea()
        i = idaapi.insn_t()
        if idaapi.decode_insn(i, ea):
            mnemonic = ida_ua.ua_mnem(ea)
            out_text = f"Instruction {mnemonic}, Operands:"
            op_lines = []
            for idx in range(8):
                op: idaapi.op_t = i.ops[idx]
                if op.type == idaapi.o_void:
                    break
                op_lines.append(f"Op{idx+1}: " + self.format_operand(i, op))
            if len(op_lines) > 0:
                out_text += "\n" + "\n".join(op_lines)
                ida_kernwin.ask_form(f"""BUTTON YES* Ok
{out_text}
""")
        else:
            ida_kernwin.warning(f"Instruction at 0x{ea:x} could not be decoded!")

    def term(self):
        pass

def PLUGIN_ENTRY():
    return instr_info_plugin_t()


repository: techbliss/Win-DBG-Helper
file: Win-Dbg-Helper.py
wanted_name: Windbg helper
comment: This is a comment
created_at: 2014-08-10 21:43:57+00:00
pushed_at: 2014-08-21 19:54:27+00:00
url:
https://github.com/techbliss/Win-DBG-Helper/blob/334eea64ade34c27c7bbcc4befe46d271eb2885c/Win-Dbg-Helper.py

content:
import subprocess
import idaapi
import idc
from idc import *
from idaapi import *



class Gullasch(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "This is a comment"

    help = "Windbg helper"
    wanted_name = "Windbg helper"
    wanted_hotkey = "Alt-F6"

    def init(self):
        idaapi.msg("windbg helper is found Plugin is found. \n")
        return idaapi.PLUGIN_OK

    def run(self, arg):
        idaapi.msg("run() called with %d!\n" % arg)

    def term(self):
        idaapi.msg("")

    def AddMenuElements(self):
        idaapi.add_menu_item("Help/", "WinDbg helper", "", 0, self.Chop, ())
        idaapi.add_menu_item("Help/", "OpCodes", "", 0, self.Onion, ())

    def run(self, arg = 0):
        idaapi.msg("helper found.\n")

        self.AddMenuElements()

    def Chop(self):
        import os
        import subprocess
        import sys
        subprocess.Popen(os.path.join(os.path.expanduser('~'), os.path.expandvars('%IDADIR%'),
    'debugger.chm'), shell=True)

    def Onion(self):
        subprocess.Popen(os.path.join(os.path.expanduser('~'), os.path.expandvars('%IDADIR%'),
    'opcodes.chm'), shell=True)


def PLUGIN_ENTRY():
    return Gullasch()


repository: kazabubu21/Pinokio
file: pinokio.py
url: https://github.com/kazabubu21/Pinokio/blob/5d8cf0537d73d76b5fd22a44bf5b5638074570c8/pinokio.py
wanted_name: Pinokio
comment: Uses davinci-003 to find vulnerabilities in the code
created_at: 2022-12-27 11:24:00+00:00
pushed_at: 2022-12-27 11:24:25+00:00
description:
Pinokio is a Python script which uses OpenAI's davinci-003 model to help with vulnerabilities research on functions decompiled by IDA Pro.

content:
import functools
import json
import idaapi
import ida_hexrays
import ida_kernwin
import idc
import openai
import os
import re
import textwrap
import threading

# Set your API key here, or put in in the OPENAI_API_KEY environment variable.
openai.api_key = ""

# =============================================================================
# Setup the context menu and hotkey in IDA
# =============================================================================

class PinokioVulPlugin(idaapi.plugin_t):
    flags = 0
    explain_action_name = "pinokio:find_vulnerabilities"
    explain_menu_path = "Edit/Pinokio/Find Vulnerabilities"
    wanted_name = 'Pinokio'
    wanted_hotkey = ''
    comment = "Uses davinci-003 to find vulnerabilities in the code"
    help = "See usage instructions on GitHub"
    menu = None

    def init(self):
        # Check whether the decompiler is available
        if not ida_hexrays.init_hexrays_plugin():
            return idaapi.PLUGIN_SKIP

        # Function Find Vulnerabilities action
        explain_action = idaapi.action_desc_t(self.explain_action_name,
                                              'Find Vulnerabilities',
                                              FindVulHandler(),
                                              "Ctrl+Alt+G",
                                              'Use davinci-003 to Find Vulnerabilities in the selected function',
                                              199)
        idaapi.register_action(explain_action)
        idaapi.attach_action_to_menu(self.explain_menu_path, self.explain_action_name, idaapi.SETMENU_APP)


        # Register context menu actions
        self.menu = ContextMenuHooks()
        self.menu.hook()

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        idaapi.detach_action_from_menu(self.explain_menu_path, self.explain_action_name)
        if self.menu:
            self.menu.unhook()
        return

# -----------------------------------------------------------------------------

class ContextMenuHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        # Add actions to the context menu of the Pseudocode view
        if idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE:
            idaapi.attach_action_to_popup(form, popup, PinokioVulPlugin.explain_action_name, "Pinokio/")

# -----------------------------------------------------------------------------

def comment_callback(address, view, response):
    """
    Callback that sets a comment at the given address.
    :param address: The address of the function to comment
    :param view: A handle to the decompiler window
    :param response: The comment to add
    """
    # Add newlines at the end of each sentence.
    response = "\n".join(textwrap.wrap(response, 80, replace_whitespace=False))

    # Add the response as a comment in IDA.
    idc.set_func_cmt(address, response, 0)
    # Refresh the window so the comment is displayed properly
    if view:
        view.refresh_view(False)
    print("davinci-003 query finished!")


# -----------------------------------------------------------------------------

class FindVulHandler(idaapi.action_handler_t):
    """
    This handler is tasked with querying davinci-003 for an Find Vulnerabilities in the
    given function. Once the reply is received, it is added as a function
    comment.
    """
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async("can you find potential security vulnerabilities on this function?\n"
                          + str(decompiler_output),
                          functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


# =============================================================================
# davinci-003 interaction
# =============================================================================

def query_model(query, cb, max_tokens=2500):
    """
    Function which sends a query to davinci-003 and calls a callback when the response is available.
    Blocks until the response is received
    :param query: The request to send to davinci-003
    :param cb: Tu function to which the response will be passed to.
    """
    try:
        response = openai.Completion.create(
            model="text-davinci-003",
            prompt=query,
            temperature=0.6,
            max_tokens=max_tokens,
            top_p=1,
            frequency_penalty=1,
            presence_penalty=1,
            timeout=60  # Wait 60 seconds maximum
        )
        ida_kernwin.execute_sync(functools.partial(cb, response=response.choices[0].text), ida_kernwin.MFF_WRITE)
    except openai.InvalidRequestError as e:
        # Context length exceeded. Determine the max number of tokens we can ask for and retry.
        m = re.search(r'maximum context length is (\d+) tokens, however you requested \d+ tokens \((\d+) in your '
                      r'prompt;', str(e))
        if not m:
            print(f"davinci-003 could not complete the request: {str(e)}")
            return
        (hard_limit, prompt_tokens) = (int(m.group(1)), int(m.group(2)))
        max_tokens = hard_limit - prompt_tokens
        if max_tokens >= 750:
            print(f"Context length exceeded! Reducing the completion tokens to {max_tokens}...")
            query_model(query, cb, max_tokens)
        else:
            print("Unfortunately, this function is too big to be analyzed with the model's current API limits.")

    except openai.OpenAIError as e:
        print(f"davinci-003 could not complete the request: {str(e)}")
    except Exception as e:
        print(f"General exception encountered while running the query: {str(e)}")

# -----------------------------------------------------------------------------

def query_model_async(query, cb):
    """
    Function which sends a query to davinci-003 and calls a callback when the response is available.
    :param query: The request to send to davinci-003
    :param cb: Tu function to which the response will be passed to.
    """
    print("Request to davinci-003 sent...")
    t = threading.Thread(target=query_model, args=[query, cb])
    t.start()

# =============================================================================
# Main
# =============================================================================

def PLUGIN_ENTRY():
    if not openai.api_key:
        openai.api_key = os.getenv("OPENAI_API_KEY")
        if not openai.api_key:
            print("Please edit this script to insert your OpenAI API key!")
            raise ValueError("No valid OpenAI API key found")

    return PinokioVulPlugin()


repository: Hipepper/antiVM
file: antiVM.py
url: https://github.com/Hipepper/antiVM/blob/2349ad405c517f5f5c4e4a4c39344fdb53255829/antiVM.py
comment: antiVM plugin for IDA Pro (using yara framework)
created_at: 2022-09-01 08:36:18+00:00
pushed_at: 2022-09-02 07:48:48+00:00
description:
Ida pro plugin. The antiVM aims to quickly identify anti-virtual machine and anti-sandbox behavior. This can speed up malware analysis.

content:
import idaapi
import idautils
import ida_bytes
import idc
import ida_kernwin
import yara
import string
import os

from tkinter import messagebox

PLUGIN_NAME = "antiVM"
PLUGIN_HOTKEY = "Ctrl-Alt-A"
VERSION = '1.0.0'
globalRuleFile = os.path.join(os.path.dirname(os.path.abspath(__file__)), "antiVM.rules")
# globalRuleFile = "antiVM.rules"

try:
    class Kp_Menu_Context(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)


        @classmethod
        def get_name(self):
            return self.__name__


        @classmethod
        def get_label(self):
            return self.label


        @classmethod
        def register(self, plugin, label):
            self.plugin = plugin
            self.label = label
            instance = self()
            return idaapi.register_action(idaapi.action_desc_t(
                self.get_name(),  # Name. Acts as an ID. Must be unique.
                instance.get_label(),  # Label. That's what users see.
                instance  # Handler. Called when activated, and for updating
            ))


        @classmethod
        def unregister(self):
            """Unregister the action.
            After unregistering the class cannot be used.
            """
            idaapi.unregister_action(self.get_name())


        @classmethod
        def activate(self, ctx):
            # dummy method
            return 1


        @classmethod
        def update(self, ctx):
            if ctx.form_type == idaapi.BWN_DISASM:
                return idaapi.AST_ENABLE_FOR_WIDGET
            return idaapi.AST_DISABLE_FOR_WIDGET


    class Searcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1

except:
    pass

p_initialized = False



class YaraSearchResultChooser(idaapi.Choose):
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False):
        idaapi.Choose.__init__(
            self,
            title,
            [
                ["Address", idaapi.Choose.CHCOL_HEX|10],
                ["Rule Name", idaapi.Choose.CHCOL_PLAIN|20],
                ["Match Name", idaapi.Choose.CHCOL_PLAIN|20],
                ["Match", idaapi.Choose.CHCOL_PLAIN|40],
                ["Type", idaapi.Choose.CHCOL_PLAIN|10],
            ],
            flags=flags,
            width=width,
            height=height,
            embedded=embedded)
        self.items = items
        self.selcount = 0
        self.n = len(items)


    def OnClose(self):
        return


    def OnSelectLine(self, n):
        self.selcount += 1
        ida_kernwin.jumpto(self.items[n][0])


    def OnGetLine(self, n):
        res = self.items[n]
        res = [idc.atoa(res[0]), res[1], res[2], res[3], res[4]]
        return res


    def OnGetSize(self):
        n = len(self.items)
        return n


    def show(self):
        return self.Show() >= 0

#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class antiVM_Plugin_t(idaapi.plugin_t):
    comment = "antiVM plugin for IDA Pro (using yara framework)"
    help = ""
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY
    flags = idaapi.PLUGIN_KEEP


    def init(self):
        global p_initialized

        # register popup menu handlers
        try:
            Searcher.register(self, "antiVM")
        except:
            pass

        if p_initialized is False:
            p_initialized = True
            idaapi.register_action(idaapi.action_desc_t(
                "antiVM",
                "Find antiVM rule matches!",
                self.search,
                None,
                None,
                0))
            idaapi.attach_action_to_menu("Edit/antiVM", "antiVM", idaapi.SETMENU_APP)
            print("=" * 80)
            print(r"              _   ___      ____  __ ")
            print(r"             | | (_) \    / /  \/  |")
            print(r"   __ _ _ __ | |_ _ \ \  / /| \  / |")
            print(r"  / _` | '_ \| __| | \ \/ / | |\/| |")
            print(r" | (_| | | | | |_| |  \  /  | |  | |")
            print(r"  \__,_|_| |_|\__|_|   \/   |_|  |_|")
            print("=" * 80)

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass


    def toVirtualAddress(self, offset, segments):
        va_offset = 0
        for seg in segments:
            if seg[1] <= offset < seg[2]:
                va_offset = seg[0] + (offset - seg[1])
        return va_offset


    def search(self, yara_file):
        memory, offsets = self.get_memory()
        try:
            rules = yara.compile(yara_file)
        except:
            print("ERROR: Cannot compile Yara rules from %s" % yara_file)
            return
        values = self.yarasearch(memory, offsets, rules)
        c = YaraSearchResultChooser("antiVM results", values)
        r = c.show()
        title = "antiVM result"
        message = "antiVM rules path:" + str(globalRuleFile) + "\n"
        message += "find anti nums:" + str(len(values))
        messagebox.showinfo(title,message)


    def yarasearch(self, memory, offsets, rules):
        values = list()
        matches = rules.match(data=memory)
        for rule_match in matches:
            name = rule_match.rule
            for match in rule_match.strings:
                match_string = match[2]
                match_type = 'unknown'
                if all(chr(c) in string.printable for c in match_string):
                    match_string = match_string.decode('utf-8')
                    match_type = 'ascii string'
                elif all(chr(c) in string.printable+'\x00' for c in match_string) and (b'\x00\x00' not in match_string):
                     match_string = match_string.decode('utf-16')
                     match_type = 'wide string'
                else:
                    match_string = " ".join("{:02x}".format(c) for c in match_string)
                    match_type = 'binary'

                value = [
                    self.toVirtualAddress(match[0], offsets),
                    name,
                    match[1],
                    match_string,
                    match_type
                ]
                values.append(value)
        return values


    def get_memory(self):
        result = bytearray()
        segment_starts = [ea for ea in idautils.Segments()]
        offsets = []
        start_len = 0
        for start in segment_starts:
            end = idc.get_segm_attr(start, idc.SEGATTR_END)
            result += ida_bytes.get_bytes(start, end - start)
            offsets.append((start, start_len, len(result)))
            start_len = len(result)
        return bytes(result), offsets


    def run(self, arg):
        if os.path.exists(globalRuleFile) != True:
            print("ERROR: can not find antiVM.rules in the root path!!")
            exit()
        print("antiVM INFO:",globalRuleFile)
        yara_file = globalRuleFile
        self.search(yara_file)


# register IDA plugin
def PLUGIN_ENTRY():
    return antiVM_Plugin_t()


repository: Dimonlipko/Leaf-ECU-Configurator
file: ddtplugins/ddt4all-FirmwareReceiver.py
description: Program for configuring and monitoring the parameters of the Leaf motor control unit.
created_at: 2021-06-17 22:11:57+00:00
pushed_at: 2025-01-31 13:30:21+00:00
url:
https://github.com/Dimonlipko/Leaf-ECU-Configurator/blob/55bf4718a3d7de1cd7951e42618d31bb774e395d/ddtplugins/ddt4all-FirmwareReceiver.py

content:
#!/usr/bin/python
# -*- coding: utf-8 -*-
import platform
import math

try:
    import PyQt5.QtGui as gui
    import PyQt5.QtCore as core
    import PyQt5.QtWidgets as widgets
    if platform.system() == 'Darwin':
        import PyQt5.QtWebEngine as webkit
        import PyQt5.QtWebEngineWidgets as webkitwidgets
    else:
        import PyQt5.QtWebKit as webkit
        import PyQt5.QtWebKitWidgets as webkitwidgets
    def utf8(string):
        return string
    qt5 = True
except:
    import PyQt4.QtGui as gui
    import PyQt4.QtGui as widgets
    import PyQt4.QtCore as core
    import PyQt4.QtWebKit as webkit
    import PyQt4.QtWebKit as webkitwidgets
    def utf8(string):
        return string
        # return unicode(string.toUtf8(), encoding="UTF-8")
    qt5 = False

import ecu
import options
import elm
import os
import time

_ = options.translator('ddt4all')

plugin_name = _("Firmware Update")
category = _("Firmware")
need_hw = True

class FirmwareReceiverUpdate(widgets.QDialog):

    def __init__(self):
        super(FirmwareReceiverUpdate, self).__init__()
        layout = widgets.QVBoxLayout()
        info = widgets.QLabel(_(utf8("FIRMWARE UPDATE")))
        info.setAlignment(core.Qt.AlignHCenter)
        layout.addWidget(info)

        self.open_button = widgets.QPushButton(_("Open file..."))
        self.open_button.clicked.connect(self.read_file)
        layout.addWidget(self.open_button)

        info = widgets.QLabel(_(utf8("CAN request id (hex):")))
        layout.addWidget(info)
        # TODO: filter input to hex digits
        self.canbase_input = widgets.QLineEdit("79b")
        layout.addWidget(self.canbase_input)
        
        info = widgets.QLabel(_(utf8("CAN answer id (hex):")))
        layout.addWidget(info)
        # TODO: filter input to hex digits
        self.cananswer_input = widgets.QLineEdit("7bb")
        layout.addWidget(self.cananswer_input)

        info = widgets.QLabel(_(utf8("Device token (hex):")))
        layout.addWidget(info)
        self.devtoken_input = widgets.QLineEdit("cafeface")
        layout.addWidget(self.devtoken_input)

        #self.hex_count_label = widgets.QLabel()
        self.program_size_label = widgets.QLabel()
        #self.eeprom_size_label = widgets.QLabel()
        #self.confbit_size_label = widgets.QLabel()
        #self.userid_size_label = widgets.QLabel()
        self.frames_label = widgets.QLabel()
        self.sent_frames_label = widgets.QLabel()
        #self.error_frames_label = widgets.QLabel()
        #layout.addWidget(self.hex_count_label)
        layout.addWidget(self.program_size_label)
        #layout.addWidget(self.eeprom_size_label)
        #layout.addWidget(self.confbit_size_label)
        #layout.addWidget(self.userid_size_label)
        layout.addWidget(self.frames_label)
        layout.addWidget(self.sent_frames_label)
        #layout.addWidget(self.error_frames_label)

        self.progress = widgets.QProgressBar()
        layout.addWidget(self.progress)

        self.flash_button = widgets.QPushButton(_("Flash"))
        self.flash_button.clicked.connect(self.flash_firmware)
        self.flash_button.setEnabled(False)
        layout.addWidget(self.flash_button)
        self.setLayout(layout)

        self.set_sent_frames(0)
        #self.set_error_frames(0)

    def update_firmware_size(self):
        # TODO: proper localization?
        framecount = math.ceil(float(len(self.firmware_data)) / 4)
        self.frames_label.setText(_("CAN frames: %d" % framecount))
        self.progress.setMaximum(framecount)
        self.program_size_label.setText(_("FLASH: %d bytes" % len(self.firmware_data)))

    def set_sent_frames(self, framecount):
        self.sent_frames_label.setText(_("Sent: %d" % framecount))
        self.progress.setValue(framecount)

    #def set_error_frames(self, errorcount):
        #self.error_frames_label.setText(_("Errors: %d" % errorcount))

    def read_file(self):
        fname = widgets.QFileDialog.getOpenFileName(self, "Open File", os.getcwd(), "Arduino binary (*.bin)")
        if fname:
            try:
                with open(str(fname.toUtf8()), 'rb') as fp:
                    # TODO: test with python3
                    self.firmware_data = fp.read()
                self.update_firmware_size()
                self.flash_button.setEnabled(True)
            except Exception as e:
                widgets.QMessageBox.critical(self, "Error", str(e))

    def _send_can(self, request_id, request, response_id, response):
        if options.simulation_mode:
            options.main_window.logview.append(_("Sending request: {:x} {} Expecting response: {:x} {}").format(request_id, request, response_id, response))
        else:
            options.elm.send_raw('atsh {:x}'.format(request_id))
            rsp = options.elm.send_raw(request)
            # TODO: check response
            print(rsp)

    def _invert_bytes(self, s):
        return ''.join([ s[i:i+2] for i in range(0, 8, 2) ][::-1])

    def flash_firmware(self):
        self.flash_button.setEnabled(False)
        self.open_button.setEnabled(False)

        #try:

        # TODO: do we need to restore it?
        if not options.simulation_mode:
            # Disable ISOTP
            options.elm.send_raw('atcaf0')
            # Allow 8 byte packets
            options.elm.send_raw('atal')

        # Might be excessive but we make sure it's a valid hex
        canbase = int(str(self.canbase_input.text()), 16)
        cananswer = int(str(self.cananswer_input.text()), 16)
        devtoken = int(str(self.devtoken_input.text()), 16)

        options.main_window.logview.append(_("Initiating boot mode"))

        payload = self._invert_bytes('deadbeef') + self._invert_bytes('{:04x}'.format(devtoken))
        response = self._invert_bytes('deafdead') + self._invert_bytes('{:04x}'.format(devtoken))
        self._send_can(canbase, payload, cananswer, payload)

        # TODO async?
        time_to_sleep = 5
        options.main_window.logview.append(_("Waiting {} seconds for device to reboot").format(time_to_sleep))
        time.sleep(time_to_sleep)

        for block, chunk in enumerate(self.firmware_data[i:i+4] for i in range(0, len(self.firmware_data), 4)):
            # Pad with zeroes if needed
            chunk += '\00' * (4 - len(chunk))
            addrh = block / 256
            addrl = block % 256
            self._send_can(canbase, ('{:02x}'*7).format(0xff, addrl, addrh, *(ord(i) for i in chunk)),
                            cananswer, ('{:02x}'*2).format(addrl, addrh))
            self.set_sent_frames(block +1)

        options.main_window.logview.append(_("Finalizing flashing"))
        self._send_can(canbase, self._invert_bytes('c0defade'), 0, '')

        #except:
            # TODO message boxes about errors and stuff
            #pass

        self.flash_button.setEnabled(True)
        self.open_button.setEnabled(True)

def plugin_entry():
    v = FirmwareReceiverUpdate()
    v.exec_()


repository: flatz/ida_ps5_elf_plugin
file: plugins/ps5_elf.py
description: PS5 elf loader for IDA 7.5
wanted_name: PS5 elf plugin
comment: f'{wanted_name} to extend loader functionality
created_at: 2022-05-27 16:02:53+00:00
pushed_at: 2025-03-23 14:44:56+00:00
url:
https://github.com/flatz/ida_ps5_elf_plugin/blob/04a673aef17784329e1d2e628d27b22a8499f0be/plugins/ps5_elf.py

content:
import sys, os
import struct
import hashlib
import re

import ida_allins
import ida_auto
import ida_bytes
import ida_diskio
import ida_entry
import ida_funcs
import ida_hexrays
import ida_ida
import ida_idaapi
import ida_idp
import ida_kernwin
import ida_nalt
import ida_name
import ida_netnode
import ida_search
import ida_segment
import ida_typeinf
import ida_ua
import ida_xref

import idc
import idautils

from collections import OrderedDict
from io import BytesIO
from pprint import pprint

def as_uint8(x):
	return x & 0xFF

def as_uint16(x):
	return x & 0xFFFF

def as_uint32(x):
	return x & 0xFFFFFFFF

def as_uint64(x):
	return x & 0xFFFFFFFFFFFFFFFF

def align_up(x, alignment):
	return (x + (alignment - 1)) & ~(alignment - 1)

def align_down(x, alignment):
	return x & ~(alignment - 1)

def hexdump(data, cols=16, addr=0):
	byte_filter = ''.join([(len(repr(chr(x))) == 3) and chr(x) or '.' for x in range(256)])
	lines = []

	for c in range(0, len(data), cols):
		chunk = data[c:c + cols]
		hex_data = ' '.join(['%02x' % x for x in chunk])
		printable = ''.join(['%s' % ((x <= 127 and byte_filter[x]) or '.') for x in chunk])
		lines.append('%08x %-*s %s' % (addr + c, cols * 3, hex_data, printable))

	print('\n'.join(lines))

def determine_simple_type(type_decl):
	return {
		'uint8_t': 'uint8_t',
		'__uint8': 'uint8_t',
		'unsigned int8_t': 'uint8_t',
		'unsigned __int8': 'uint8_t',
		'unsigned char': 'uint8_t',
		'int8_t': 'int8_t',
		'__int8': 'int8_t',
		'signed int8_t': 'int8_t',
		'signed __int8': 'int8_t',
		'char': 'int8_t',

		'uint16_t': 'uint16_t',
		'__uint16': 'uint16_t',
		'unsigned int16_t': 'uint16_t',
		'unsigned __int16': 'uint16_t',
		'unsigned short': 'uint16_t',
		'int16_t': 'int16_t',
		'__int16': 'int16_t',
		'signed int16_t': 'int16_t',
		'signed __int16': 'int16_t',
		'short': 'int16_t',

		'uint32_t': 'uint32_t',
		'__uint32': 'uint32_t',
		'unsigned int32_t': 'uint32_t',
		'unsigned __int32': 'uint32_t',
		'unsigned int': 'uint32_t',
		'int32_t': 'int32_t',
		'__int32': 'int32_t',
		'signed int32_t': 'int32_t',
		'signed __int32': 'int32_t',
		'int': 'int32_t',

		'uint64_t': 'uint64_t',
		'__uint64': 'uint64_t',
		'unsigned int64_t': 'uint64_t',
		'unsigned __int64': 'uint64_t',
		'unsigned long': 'uint64_t',
		'unsigned long long': 'uint64_t',
		'int64_t': 'int64_t',
		'__int64': 'int64_t',
		'signed int64_t': 'int64_t',
		'signed __int64': 'int64_t',
		'long': 'int64_t',
		'long long': 'int64_t',
	}.get(type_decl, None)

def unpack_by_type(data, size, type_info):
	assert len(data) >= size

	initial_type_name = type_info.dstr()
	initial_final_type_name = type_info.get_final_type_name()

	data = data[:size]
	endian = '>' if ida_ida.inf_is_be() else '<'

	fmt = None
	value = None

	if type_info.is_array():
		type_info = type_info.get_array_element()
		type_name = type_info.dstr()
		final_type_name = type_info.get_final_type_name()
		if final_type_name is None:
			final_type_name = type_name
		final_type_name = determine_simple_type(final_type_name)
	else:
		type_name = initial_type_name
		final_type_name = initial_final_type_name
		if final_type_name is None:
			final_type_name = type_name
		final_type_name = determine_simple_type(final_type_name)

	if final_type_name == 'uint8_t':
		fmt = 'B'
	elif final_type_name == 'int8_t':
		fmt = 'b'
	elif final_type_name == 'uint16_t':
		fmt = 'H'
	elif final_type_name == 'int16_t':
		fmt = 'h'
	elif final_type_name == 'uint32_t':
		fmt = 'I'
	elif final_type_name == 'int32_t':
		fmt = 'i'
	elif final_type_name == 'uint64_t':
		fmt = 'Q'
	elif final_type_name == 'int64_t':
		fmt = 'q'
	elif final_type_name == 'float':
		fmt = 'f'
	elif final_type_name == 'double':
		fmt = 'd'
	elif type_info.is_ptr():
		fmt = 'Q'
	else:
		# XXX: Need to handle this manually.
		value = data

	if fmt is not None:
		stride = struct.calcsize(fmt)
		count = size // stride
		if count > 1:
			value = struct.unpack('%s%d%s' % (endian, count, fmt), data)
			value = list(value)
		elif count == 1:
			value = struct.unpack('%s%s' % (endian, fmt), data)[0]
		else:
			value = []

	return value, size

def get_struct_size(name):
	assert isinstance(name, str)

	type_id = ida_typeinf.get_named_type_tid(name)
	type_info = ida_typeinf.tinfo_t()
	if type_id == ida_idaapi.BADADDR or not type_info.get_type_by_tid(type_id) or not type_info.is_udt():
		raise RuntimeError('Structure not found: %s' % name)

	return type_info.get_size()

def parse_struct(name, data):
	assert isinstance(name, str)
	assert isinstance(data, bytes)

	type_id = ida_typeinf.get_named_type_tid(name)
	type_info = ida_typeinf.tinfo_t()
	if type_id == ida_idaapi.BADADDR or not type_info.get_type_by_tid(type_id) or not type_info.is_udt():
		raise RuntimeError('Structure not found: %s' % name)

	total_size = type_info.get_size()
	assert len(data) >= total_size

	udt = ida_typeinf.udt_type_data_t()
	assert type_info.get_udt_details(udt)
	count = type_info.get_udt_nmembers()

	#if type_info.is_typedef():
	#	return None

	fields = OrderedDict()

	for idx, udm in enumerate(udt):
		if udm.is_gap():
			continue

		member_name = udm.name
		member_offset = udm.offset // 8
		member_size = udm.size // 8
		member_type_info = udm.type

		value_data = data[member_offset:member_offset + member_size]

		info = unpack_by_type(value_data, member_size, member_type_info)

		fields[member_name] = info[0]

	return fields

def check_insn_format(ea, mnem, operands=[]):
	if ida_ua.print_insn_mnem(ea).lower() != mnem.lower():
		return False

	match = True

	for i, (needed_type, needed_value) in enumerate(operands):
		if needed_type is None:
			continue
		real_type = idc.get_operand_type(ea, i)
		if real_type != needed_type:
			match = False
			break

		if needed_value is None:
			continue
		if isinstance(needed_value, str):
			# XXX: Cannot use ida_ua.print_operand because it returns garbage data.
			value = idc.print_operand(ea, i).lower().strip()
			needed_value = needed_value.lower().strip()
		else:
			value = idc.get_operand_value(ea, i)

		if value != needed_value:
			match = False
			break

	return match

def read_cstring_at(ea, encoding='utf-8'):
	length = ida_bytes.get_max_strlit_length(ea, ida_nalt.STRTYPE_C)
	data = ida_bytes.get_strlit_contents(ea, length, ida_nalt.STRTYPE_C)
	if encoding is not None:
		data = data.decode(encoding)
	return data

def refresh_views():
	ida_kernwin.refresh_idaview_anyway()

	widget = ida_kernwin.get_current_widget()
	vu = ida_hexrays.get_widget_vdui(widget)
	if vu:
		vu.refresh_ctext()

def sha1(data):
	return hashlib.sha1(data).digest()

class ObjectInfo(object):
	AUTO_EXPORT = (1 << 0)
	WEAK_EXPORT = (1 << 1)
	LOOSE_IMPORT = (1 << 3)

	CANT_STOP = (1 << 0)
	EXCLUSIVE_LOAD = (1 << 1)
	EXCLUSIVE_START = (1 << 2)
	CAN_RESTART = (1 << 3)
	CAN_RELOCATE = (1 << 4)
	CANT_SHARE = (1 << 5)

	ATTR_MASK = 0xFFFFFFFFFFFF
	EXPORT_ATTR_MASK = AUTO_EXPORT | WEAK_EXPORT | LOOSE_IMPORT
	IMPORT_ATTR_MASK = EXCLUSIVE_LOAD | CANT_STOP | CAN_RESTART | EXCLUSIVE_START | CANT_SHARE | CAN_RELOCATE

	ENCODING_CHARSET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+-'

	NID_SUFFIX = bytes.fromhex('518D64A635DED8C1E6B039B1C3E55230')

	def __init__(self):
		self.id = None
		self.version_major = None
		self.version_minor = None
		self.name_offset = None
		self.name = None
		self.attrs = None
		self.is_export = None

	def set_info(self, value, is_export=False):
		self.id = ObjectInfo.obj_id(value)
		self.version_major, self.version_minor = ObjectInfo.obj_version(value)
		self.name_offset = ObjectInfo.obj_name_offset(value)
		self.is_export = is_export

	def set_attr(self, value):
		self.id = ObjectInfo.obj_id(value)
		self.attrs = ObjectInfo.obj_attrs(value)

		if self.is_export:
			if self.attrs & (~ObjectInfo.EXPORT_ATTR_MASK):
				ida_kernwin.warning('Unsupported export attributes: %s' % ObjectInfo.stringify_attrs(self.attrs, True))
		else:
			if self.attrs & (~ObjectInfo.IMPORT_ATTR_MASK):
				ida_kernwin.warning('Unsupported import attributes: %s' % ObjectInfo.stringify_attrs(self.attrs, False))

	def update_name(self, base_ea):
		assert self.name_offset is not None
		self.name = read_cstring_at(base_ea + self.name_offset)

	def __repr__(self):
		if self.id is not None:
			return 'ObjectInfo(id=0x%x, version=%02d.%02d, name=%s, attrs=%s, is_export=%d)' % (self.id, self.version_major, self.version_minor, self.name if self.name else 'nil', ObjectInfo.stringify_attrs(self.attrs, self.is_export), self.is_export)
		else:
			return 'n/a'

	@staticmethod
	def calculate_nid(x):
		return sha1(x.encode('ascii') + ObjectInfo.NID_SUFFIX)[:8]

	@staticmethod
	def encode_nid(x):
		enc, = struct.unpack('<Q', ObjectInfo.calculate_nid(x))

		result = ObjectInfo.ENCODING_CHARSET[(enc & 0xF) << 2]
		tmp = enc >> 4
		while tmp != 0:
			result += ObjectInfo.ENCODING_CHARSET[tmp & 0x3F]
			tmp = tmp >> 6

		return result[::-1], enc

	@staticmethod
	def decode_nid(x):
		size = len(x)
		assert size == 11

		result = 0

		for i in range(size):
			idx = ord(x[i])
			assert idx < len(ObjectInfo.DECODING_CHARSET)
			value = ObjectInfo.DECODING_CHARSET[idx]
			assert value != 0xFF

			if i == size - 1:
				result = (result << 4) | (value >> 2)
			else:
				result = (result << 6) | value

		return result

	@staticmethod
	def encode_obj_id(x):
		result = ObjectInfo.ENCODING_CHARSET[(x & 0xF) << 2]
		tmp = x >> 4
		while tmp != 0:
			result += ObjectInfo.ENCODING_CHARSET[tmp & 0x3F]
			tmp = tmp >> 6

		return result[::-1]

	@staticmethod
	def decode_obj_id(x):
		size = len(x)
		assert size <= 4

		result = 0

		for i in range(size):
			idx = ord(x[i])
			assert idx < len(ObjectInfo.DECODING_CHARSET)
			value = ObjectInfo.DECODING_CHARSET[idx]
			assert value != 0xFF
			result = (result << 6) | value

		return result

	@staticmethod
	def obj_id(value):
		return (value >> 48) & 0xFFFF

	@staticmethod
	def obj_attrs(value):
		return (value & ObjectInfo.ATTR_MASK)

	@staticmethod
	def obj_name_offset(value):
		return value & ((1 << 32) - 1)

	@staticmethod
	def obj_version(value):
		major, minor = (value >> 32) & 0xFF, (value >> 40) & 0xFF
		return major, minor

	@staticmethod
	def stringify_attrs(attrs, is_export):
		if not attrs:
			return 'none'

		result = []

		if is_export:
			if attrs & ObjectInfo.AUTO_EXPORT:
				result.append('AUTO_EXPORT')
				attrs &= ~ObjectInfo.AUTO_EXPORT
			if attrs & ObjectInfo.WEAK_EXPORT:
				result.append('WEAK_EXPORT')
				attrs &= ~ObjectInfo.WEAK_EXPORT
			if attrs & ObjectInfo.LOOSE_IMPORT:
				result.append('LOOSE_IMPORT')
				attrs &= ~ObjectInfo.LOOSE_IMPORT
		else:
			if attrs & ObjectInfo.CANT_STOP:
				result.append('CANT_STOP')
				attrs &= ~ObjectInfo.CANT_STOP
			if attrs & ObjectInfo.EXCLUSIVE_LOAD:
				result.append('EXCLUSIVE_LOAD')
				attrs &= ~ObjectInfo.EXCLUSIVE_LOAD
			if attrs & ObjectInfo.EXCLUSIVE_START:
				result.append('EXCLUSIVE_START')
				attrs &= ~ObjectInfo.EXCLUSIVE_START
			if attrs & ObjectInfo.CAN_RESTART:
				result.append('CAN_RESTART')
				attrs &= ~ObjectInfo.CAN_RESTART
			if attrs & ObjectInfo.CAN_RELOCATE:
				result.append('CAN_RELOCATE')
				attrs &= ~ObjectInfo.CAN_RELOCATE
			if attrs & ObjectInfo.CANT_SHARE:
				result.append('CANT_SHARE')
				attrs &= ~ObjectInfo.CANT_SHARE
		if attrs != 0:
			result.append('0x%x' % attrs)

		return '|'.join(result)

def load_known_nids(file_name):
	nids = {}

	file_path = os.path.join(ida_diskio.idadir(ida_diskio.CFG_SUBDIR), file_name)
	if not os.path.isfile(file_path):
		ida_kernwin.warning('NID database not found: %s' % file_path)
		return nids

	with open(file_path, 'r') as f:
		for line in f.readlines():
			line = line.rstrip('\r\n').strip()
			if not line:
				continue

			name = line.split(':')
			if name:
				name = name[0]
			else:
				continue

			nid_ascii, nid = ObjectInfo.encode_nid(name)
			nids[nid] = name

			# Try to prepend underscore symbol to symbol, thus getting a new one.
			if name.startswith('_'):
				name = name[1:]
			name = '_' + name
			nid_ascii, nid = ObjectInfo.encode_nid(name)
			nids[nid] = name

	return nids

class ElfUtil(object):
	# Node and tags to find corresponding net node.
	ELF_NODE = '$ elfnode'

	ELF_PHT_TAG = 'p'
	ELF_SHT_TAG = 's'

	# File types.
	ET_LOOS = 0xFE00
	ET_HIOS = 0xFEFF
	ET_SCE_EXEC_ASLR = ET_LOOS + 0x10
	ET_SCE_DYNAMIC = ET_LOOS + 0x18

	# Segment types.
	PT_LOAD = 0x1
	PT_DYNAMIC = 0x2
	PT_NOTE = 0x4
	PT_TLS = 0x7
	PT_GNU_RELRO = 0x6474E552
	PT_GNU_EH_FRAME = 0x6474E550
	PT_SCE_PROCPARAM = 0x61000001
	PT_SCE_MODULE_PARAM = 0x61000002
	PT_SCE_COMMENT = 0x6FFFFF00
	PT_SCE_VERSION = 0x6FFFFF01

	# Segment flags.
	PF_NONE = 0x0
	PF_EXEC = 0x1
	PF_WRITE = 0x2
	PF_READ = 0x4
	PF_READ_WRITE = PF_READ | PF_WRITE

	# Supported dynamic tags.
	DT_LOOS = 0x60000000
	DT_HIOS = 0x6FFFFFFF
	DT_NULL = 0x0
	DT_NEEDED = 0x1
	DT_PLTGOT = 0x3
	DT_INIT = 0xC
	DT_FINI = 0xD
	DT_SONAME = 0xE
	DT_TEXTREL = 0x16
	DT_SCE_IDTABENTSZ = DT_LOOS + 0x5
	DT_SCE_FINGERPRINT = DT_LOOS + 0x1000007
	DT_SCE_ORIGINAL_FILENAME = DT_LOOS + 0x1000009
	DT_SCE_MODULE_INFO = DT_LOOS + 0x100000D
	DT_SCE_NEEDED_MODULE = DT_LOOS + 0x100000F
	DT_SCE_MODULE_ATTR = DT_LOOS + 0x1000011
	DT_SCE_EXPORT_LIB = DT_LOOS + 0x1000013
	DT_SCE_IMPORT_LIB = DT_LOOS + 0x1000015
	DT_SCE_EXPORT_LIB_ATTR = DT_LOOS + 0x1000017
	DT_SCE_IMPORT_LIB_ATTR = DT_LOOS + 0x1000019
	DT_61000027 = DT_LOOS + 0x1000027 # TODO: Is it related to PLTGOT?
	DT_SCE_ORIGINAL_FILENAME_PPR = DT_LOOS + 0x1000041
	DT_SCE_MODULE_INFO_PPR = DT_LOOS + 0x1000043
	DT_SCE_NEEDED_MODULE_PPR = DT_LOOS + 0x1000045
	DT_SCE_IMPORT_LIB_PPR = DT_LOOS + 0x1000047
	DT_SCE_EXPORT_LIB_PPR = DT_LOOS + 0x1000049

	# Skipped dynamic tags.
	DT_PLTRELSZ = 0x2
	DT_HASH = 0x4
	DT_STRTAB = 0x5
	DT_SYMTAB = 0x6
	DT_RELA = 0x7
	DT_RELASZ = 0x8
	DT_RELAENT = 0x9
	DT_STRSZ = 0xA
	DT_SYMENT = 0xB
	DT_PLTREL = 0x14
	DT_DEBUG = 0x15
	DT_JMPREL = 0x17
	DT_INIT_ARRAY = 0x19
	DT_FINI_ARRAY = 0x1A
	DT_INIT_ARRAYSZ = 0x1B
	DT_FINI_ARRAYSZ = 0x1C
	DT_PREINIT_ARRAY = 0x20
	DT_PREINIT_ARRAYSZ = 0x21
	DT_61000025 = DT_LOOS + 0x1000025
	DT_61000029 = DT_LOOS + 0x1000029
	DT_6100002B = DT_LOOS + 0x100002B
	DT_6100002D = DT_LOOS + 0x100002D
	DT_6100002F = DT_LOOS + 0x100002F
	DT_61000031 = DT_LOOS + 0x1000031
	DT_61000033 = DT_LOOS + 0x1000033
	DT_61000035 = DT_LOOS + 0x1000035
	DT_61000037 = DT_LOOS + 0x1000037
	DT_61000039 = DT_LOOS + 0x1000039
	DT_6100003B = DT_LOOS + 0x100003B
	DT_SCE_HASHSZ = DT_LOOS + 0x100003D
	DT_SCE_SYMTABSZ = DT_LOOS + 0x100003F
	DT_RELACOUNT = DT_LOOS + 0xFFFFFF9
	DT_FLAGS_1 = DT_LOOS + 0xFFFFFFB

	# Obsoleted dynamic tags.
	DT_SYMBOLIC = 0x10
	DT_FLAGS = 0x1E

	# Unsupported dynamic tags.
	DT_RPATH = 0xF
	DT_REL = 0x11
	DT_RELSZ = 0x12
	DT_RELENT = 0x13
	DT_BIND_NOW = 0x18
	DT_RUNPATH = 0x1D
	DT_ENCODING = 0x1F
	DT_61000008 = DT_LOOS + 0x1000008
	DT_6100000A = DT_LOOS + 0x100000A
	DT_6100000B = DT_LOOS + 0x100000B
	DT_6100000C = DT_LOOS + 0x100000C
	DT_6100000E = DT_LOOS + 0x100000E
	DT_61000010 = DT_LOOS + 0x1000010
	DT_61000012 = DT_LOOS + 0x1000012
	DT_61000014 = DT_LOOS + 0x1000014
	DT_61000016 = DT_LOOS + 0x1000016
	DT_61000018 = DT_LOOS + 0x1000018
	DT_6100001A = DT_LOOS + 0x100001A
	DT_6100001B = DT_LOOS + 0x100001B
	DT_6100001C = DT_LOOS + 0x100001C
	DT_SCE_STUB_MODULE_NAME = DT_LOOS + 0x100001D
	DT_6100001E = DT_LOOS + 0x100001E
	DT_SCE_STUB_MODULE_VERSION = DT_LOOS + 0x100001F
	DT_61000020 = DT_LOOS + 0x1000020
	DT_SCE_STUB_LIBRARY_NAME = DT_LOOS + 0x1000021
	DT_61000022 = DT_LOOS + 0x1000022
	DT_SCE_STUB_LIBRARY_VERSION = DT_LOOS + 0x1000023
	DT_61000024 = DT_LOOS + 0x1000024
	DT_61000026 = DT_LOOS + 0x1000026
	DT_61000028 = DT_LOOS + 0x1000028
	DT_6100002A = DT_LOOS + 0x100002A
	DT_6100002C = DT_LOOS + 0x100002C
	DT_6100002E = DT_LOOS + 0x100002E
	DT_61000030 = DT_LOOS + 0x1000030
	DT_61000032 = DT_LOOS + 0x1000032
	DT_61000034 = DT_LOOS + 0x1000034
	DT_61000036 = DT_LOOS + 0x1000036
	DT_61000038 = DT_LOOS + 0x1000038
	DT_6100003A = DT_LOOS + 0x100003A
	DT_6100003C = DT_LOOS + 0x100003C
	DT_6100003E = DT_LOOS + 0x100003E
	DT_61000040 = DT_LOOS + 0x1000040
	DT_61000042 = DT_LOOS + 0x1000042
	DT_61000044 = DT_LOOS + 0x1000044
	DT_61000046 = DT_LOOS + 0x1000046
	DT_61000048 = DT_LOOS + 0x1000048

	# Filler for .sce_padding.
	SCE_PADDING_SIZE = 0x10
	SCE_PADDING = b'\xCC' * SCE_PADDING_SIZE

	# Magics for sceProcessParam and sceModuleParam structures.
	SCE_PROCESS_PARAM_MAGIC = b'ORBI'
	SCE_MODULE_PARAM_MAGIC = b'\xBF\xF4\x13\x3C'

	_DT_STR_MAP = {}

	def __init__(self):
		self.elf_node = ida_netnode.netnode(ElfUtil.ELF_NODE)

		if not self.is_inited():
			return

		#print('Parsing elf header.')
		self._parse_ehdr()

		#print('Parsing program headers.')
		self._parse_phdrs()

		#print('Parsing section headers.')
		self._parse_shdrs()

	def find_phdr_by_type(self, type_id, index=0):
		phdrs = []

		for i, phdr in enumerate(self.phdrs):
			p_type = ElfUtil.phdr_type(phdr)
			if p_type != type_id:
				continue
			phdrs.append(phdr)

		return phdrs[index] if index < len(phdrs) else None

	def find_phdr_by_seg(self, seg):
		if not seg:
			return None

		idx = ida_segment.get_segm_num(seg.start_ea)

		#print('Segment #%03d (start=0x%x, end=0x%0x, perm=0x%x)' % (idx, seg.start_ea, seg.end_ea, seg.perm))

		# Try to find a program header which boundaries are exactly the same as in the segment.
		#print('Doing first pass (exact).')
		for i, phdr in enumerate(self.phdrs):
			if self._match_phdr_boundaries_with_seg(i, phdr, seg, True):
				return phdr

		# Now try to find a program header which boundaries includes the segment boundaries.
		#print('Doing second pass (inexact).')
		for i, phdr in enumerate(self.phdrs):
			if self._match_phdr_boundaries_with_seg(i, phdr, seg, False):
				return phdr

		return None

	def is_inited(self):
		if self.elf_node.valobj() is None:
			return False
		node = ida_netnode.netnode(ps5_elf_plugin_t.PROSPERO_NET_NODE)
		return ida_netnode.exist(node)

	@staticmethod
	def phdr_type(phdr):
		for name in ['p_type', '_p_type', '__p_type']:
			if name in phdr:
				return phdr[name]
		return None

	@staticmethod
	def stringify_phdr_type(type_id):
		return {
			ElfUtil.PT_LOAD: 'PT_LOAD',
			ElfUtil.PT_DYNAMIC: 'PT_DYNAMIC',
			ElfUtil.PT_NOTE: 'PT_NOTE',
			ElfUtil.PT_TLS: 'PT_TLS',
			ElfUtil.PT_GNU_RELRO: 'PT_GNU_RELRO',
			ElfUtil.PT_GNU_EH_FRAME: 'PT_GNU_EH_FRAME',
			ElfUtil.PT_SCE_PROCPARAM: 'PT_SCE_PROCPARAM',
			ElfUtil.PT_SCE_MODULE_PARAM: 'PT_SCE_MODULE_PARAM',
			ElfUtil.PT_SCE_COMMENT: 'PT_SCE_COMMENT',
			ElfUtil.PT_SCE_VERSION: 'PT_SCE_VERSION',
		}.get(type_id, '0x%08x' % type_id)

	@staticmethod
	def stringify_phdr_flags(flags):
		result = []
		if flags == ElfUtil.PF_NONE:
			result.append('PF_NONE')
		else:
			if flags & ElfUtil.PF_READ:
				result.append('PF_READ')
			if flags & ElfUtil.PF_WRITE:
				result.append('PF_WRITE')
			if flags & ElfUtil.PF_EXEC:
				result.append('PF_EXEC')
		return '|'.join(result)

	@staticmethod
	def stringify_dyn_tag(tag):
		return ElfUtil._DT_STR_MAP.get(tag, '0x%x' % tag)

	@staticmethod
	def phdr_flags_from_perm(perm):
		flags = 0
		if perm & ida_segment.SEGPERM_READ:
			flags |= ElfUtil.PF_READ
		if perm & ida_segment.SEGPERM_WRITE:
			flags |= ElfUtil.PF_WRITE
		if perm & ida_segment.SEGPERM_EXEC:
			flags |= ElfUtil.PF_EXEC
		return flags

	def _match_phdr_boundaries_with_seg(self, idx, phdr, seg, exact):
		p_type, seg_flags = ElfUtil.phdr_type(phdr), ElfUtil.phdr_flags_from_perm(seg.perm)
		va_start, va_end = phdr['p_vaddr'], phdr['p_vaddr'] + phdr['p_memsz']

		info_str = 'type=0x%x(%s), flags=0x%x(%s), offset=0x%x, va_start=0x%x, va_end=0x%x, filesz=0x%x, memsz=0x%x' % (
			p_type, ElfUtil.stringify_phdr_type(p_type),
			phdr['p_flags'], ElfUtil.stringify_phdr_flags(phdr['p_flags']),
			phdr['p_offset'],
			va_start, va_end,
			phdr['p_filesz'], phdr['p_memsz']
		)

		if p_type not in [ElfUtil.PT_LOAD, ElfUtil.PT_DYNAMIC, ElfUtil.PT_SCE_PROCPARAM, ElfUtil.PT_SCE_MODULE_PARAM, ElfUtil.PT_GNU_EH_FRAME, ElfUtil.PT_TLS]:
			#print('Skipping phdr #%03d (%s): type mismatch (0x%x)' % (idx, info_str, p_type))
			return False

		if seg_flags != phdr['p_flags']:
			#print('Skipping phdr #%03d (%s): flags mismatch (segment: 0x%x, phdr: 0x%x)' % (idx, info_str, seg_flags, phdr['p_flags']))
			return False

		result = (va_start == seg.start_ea and va_end == seg.end_ea) if exact else ((va_start <= seg.start_ea < va_end) and (va_start < seg.end_ea <= va_end))
		if not result:
			#print('Skipping phdr #%03d (%s): %s addresses mismatch' % (idx, info_str, 'exact' if exact else 'inexact'))
			return False

		#print('Matched phdr #%03d (%s).' % (idx, info_str))

		return True

	def _parse_ehdr(self):
		data = self.elf_node.valobj()
		self.ehdr = parse_struct('Elf64_Ehdr', data)

	def _parse_phdrs(self):
		self.phdrs = []

		i = 0
		while self.elf_node.supval(i, ElfUtil.ELF_PHT_TAG) is not None:
			data = self.elf_node.supval(i, ElfUtil.ELF_PHT_TAG)
			phdr = parse_struct('Elf64_Phdr', data)
			self.phdrs.append(phdr)
			i += 1

	def _parse_shdrs(self):
		self.shdrs = []

		i = 0
		while self.elf_node.supval(i, ElfUtil.ELF_SHT_TAG) is not None:
			data = self.elf_node.supval(i, ElfUtil.ELF_SHT_TAG)
			shdr = parse_struct('Elf64_Shdr', data)
			self.shdrs.append(shdr)
			i += 1

for key in dir(ElfUtil):
	if not key.startswith('DT_') or key in ['DT_LOOS', 'DT_HIOS']:
		continue

	int_value = getattr(ElfUtil, key)
	if int_value >= ElfUtil.DT_LOOS and int_value <= ElfUtil.DT_HIOS:
		str_value = '%s(DT_LOOS+0x%x)' % (key, int_value - ElfUtil.DT_LOOS)
	else:
		str_value = '%s(0x%x)' % (key, int_value)

	ElfUtil._DT_STR_MAP[int_value] = str_value

class ElfTable(object):
	def __init__(self, ea=ida_idaapi.BADADDR, size=ida_idaapi.BADADDR, entry_size=ida_idaapi.BADADDR, entry_count=ida_idaapi.BADADDR):
		self.ea = ea
		self.size = size
		self.entry_size = entry_size
		self.entry_count = entry_count

	def type_name(self):
		return self.__class__.__name__

	def struct_name(self):
		return None

	def is_loaded(self):
		return self.ea != ida_idaapi.BADADDR and self.size != ida_idaapi.BADADDR

	def is_table_loaded(self):
		return self.is_loaded() and self.entry_size != ida_idaapi.BADADDR

	def get_entry(self, idx):
		assert self.is_table_loaded()
		assert self.struct_name() is not None

		count = self.get_num_entries()
		assert idx >= 0 and idx < count

		data = ida_bytes.get_bytes(self.ea + idx * self.entry_size, self.entry_size)
		assert len(data) == self.entry_size

		entry = parse_struct(self.struct_name(), data)

		return entry

	def get_num_entries(self):
		# XXX: We do not rely on entry count because it could be not set or may a wrong value.
		return self.size // self.entry_size

	def __repr__(self):
		params = []

		if self.ea != ida_idaapi.BADADDR:
			params.append('ea=0x%x' % self.ea)
		if self.size != ida_idaapi.BADADDR:
			params.append('size=0x%x' % self.size)
		if self.entry_size != ida_idaapi.BADADDR:
			params.append('entry_size=0x%x' % self.entry_size)
			params.append('real_entry_count=%d' % self.get_num_entries())
		if self.entry_count != ida_idaapi.BADADDR:
			params.append('entry_count=%d' % self.entry_count)

		return '%s(%s)' % (self.type_name(), ', '.join(params))

class RelaRelocTable(ElfTable):
	R_AMD64_NONE = 0
	R_AMD64_64 = 1
	R_AMD64_PC32 = 2
	R_AMD64_GOT32 = 3
	R_AMD64_PLT32 = 4
	R_AMD64_COPY = 5
	R_AMD64_GLOB_DAT = 6
	R_AMD64_JUMP_SLOT = 7
	R_AMD64_RELATIVE = 8
	R_AMD64_GOTPCREL = 9
	R_AMD64_32 = 10
	R_AMD64_32S = 11
	R_AMD64_16 = 12
	R_AMD64_PC16 = 13
	R_AMD64_8 = 14
	R_AMD64_PC8 = 15
	R_AMD64_DTPMOD64 = 16
	R_AMD64_DTPOFF64 = 17
	R_AMD64_TPOFF64 = 18
	R_AMD64_TLSGD = 19
	R_AMD64_TLSLD = 20
	R_AMD64_DTPOFF32 = 21
	R_AMD64_GOTTPOFF = 22
	R_AMD64_TPOFF32 = 23
	R_AMD64_PC64 = 24
	R_AMD64_GOTOFF64 = 25
	R_AMD64_GOTPC32 = 26

	class Record(object):
		def __init__(self, entry):
			self.entry = entry

		def get_symbol_idx(self):
			return as_uint64(self.entry['r_info']) >> 32

		def get_type(self):
			return as_uint32(self.entry['r_info'])

		def __repr__(self):
			params = []

			if self.entry is not None:
				params.append('entry=%s' % repr(self.entry))

			return 'Record(%s)' % ', '.join(params)

	def type_name(self):
		return 'ELF RELA Relocation Table'

	def struct_name(self):
		return 'Elf64_Rela'

class JmpRelRelocTable(RelaRelocTable):
	def type_name(self):
		return 'ELF JMPREL Relocation Table'

class SymbolTable(ElfTable):
	STB_LOCAL = 0
	STB_GLOBAL = 1
	STB_WEAK = 2

	STT_NOTYPE = 0
	STT_OBJECT = 1
	STT_FUNC = 2
	STT_SECTION = 3
	STT_FILE = 4
	STT_COMMON = 5
	STT_TLS = 6

	SHN_UNDEF = 0

	@staticmethod
	def sanitize_name(name):
		return re.sub(r'[^a-zA-Z0-9_]', '_', name)

	class Symbol(object):
		def __init__(self, entry):
			self.entry = entry
			self.module_name = None
			self.library_name = None
			self.symbol_name = None
			self.is_export = None

		def set_descriptor(self, module_name, library_name, symbol_name, is_export):
			self.module_name = module_name
			self.library_name = library_name
			self.symbol_name = symbol_name
			self.is_export = is_export

		def has_descriptor(self):
			return self.module_name is not None and self.library_name is not None and self.symbol_name is not None

		def get_binding(self):
			return as_uint8(self.entry['st_info']) >> 4

		def get_type(self):
			return as_uint8(self.entry['st_info']) & 0xF

		def _get_name(self, kind='simple'):
			module_name = SymbolTable.sanitize_name(self.module_name)
			library_name = SymbolTable.sanitize_name(self.library_name)

			if self.is_func(): suffix = 'f'
			elif self.is_object(): suffix = 'o'
			else: suffix = 'u'

			if isinstance(self.symbol_name, str):
				if kind == 'extended':
					symbol_name = SymbolTable.sanitize_name(self.symbol_name)
					if module_name != library_name:
						symbol_name = '%s_%s_%s' % (module_name, library_name, symbol_name)
					else:
						symbol_name = '%s_%s' % (module_name, symbol_name)
				elif kind == 'comment':
					symbol_name = self.symbol_name
					demangled_name = ida_name.demangle_name(symbol_name, idc.get_inf_attr(idc.INF_LONG_DEMNAMES))
					if demangled_name is not None:
						demangled_name = demangled_name.strip()
						if demangled_name:
							symbol_name = demangled_name
				else:
					symbol_name = self.symbol_name
			else:
				if module_name != library_name:
					symbol_name = 'nid%s_%s_%s_0x%016x' % (suffix, module_name, library_name, self.symbol_name)
				else:
					symbol_name = 'nid%s_%s_0x%016x' % (suffix, module_name, self.symbol_name)

			return symbol_name

		def get_name(self):
			return self._get_name('simple')

		def get_name_ex(self):
			return self._get_name('extended')

		def get_name_comment(self):
			name = self._get_name('comment')

			if self.is_func():
				type_str = 'Function'
			elif self.is_object():
				type_str = 'Object'
			else:
				type_str = 'Unknown'

			return '\n'.join([
				'%s:' % type_str,
				'  Module: %s' % self.module_name,
				'  Library: %s' % self.library_name,
				'  Name: %s' % name
			])

		def is_local(self):
			return self.get_binding() & SymbolTable.STB_LOCAL

		def is_global(self):
			return self.get_binding() & SymbolTable.STB_GLOBAL

		def is_weak(self):
			return self.get_binding() & SymbolTable.STB_WEAK

		def is_object(self): # Variable, array, etc.
			return self.get_type() == SymbolTable.STT_OBJECT

		def is_func(self): # Method or function.
			return self.get_type() == SymbolTable.STT_FUNC

		def is_tls(self): # TLS stuff
			return self.get_type() == SymbolTable.STT_TLS

		def __repr__(self):
			params = []

			if self.module_name is not None:
				params.append('module_name=%s' % self.module_name)
			if self.library_name is not None:
				params.append('library_name=%s' % self.library_name)
			if self.symbol_name is not None:
				if isinstance(self.symbol_name, str):
					params.append('symbol_name=%s' % self.symbol_name)
				else:
					params.append('symbol_nid=%s' % hex(self.symbol_name))
			if self.is_export is not None:
				params.append('is_export=%d' % self.is_export)

			if self.entry is not None:
				params.append('entry=%s' % repr(self.entry))

			return 'Symbol(%s)' % ', '.join(params)

	def type_name(self):
		return 'ELF Symbol Table'

	def struct_name(self):
		return 'Elf64_Sym'

class StringTable(ElfTable):
	def type_name(self):
		return 'ELF String Table'

	def get_string(self, offset):
		assert self.is_loaded()

		assert offset >= 0 and offset < self.size

		return read_cstring_at(self.ea + offset)

class HashTable(ElfTable):
	def type_name(self):
		return 'ELF Hash Table'

class IdTable(ElfTable):
	def type_name(self):
		return 'ELF ID Table'

class ps5_elf_plugin_t(ida_idaapi.plugin_t):
	flags = ida_idaapi.PLUGIN_PROC
	wanted_name = 'PS5 elf plugin'
	comment = f'{wanted_name} to extend loader functionality'
	wanted_hotkey = ''
	help = ''

	# Inhibit flags for symbol names.
	DEMANGLED_FORM = 0x0EA3FFE7 # MNG_SHORT_FORM | MNG_NOBASEDT | MNG_NOCALLC | MNG_NOCSVOL

	# Inhibit flags for type info comments.
	DEMANGLED_TYPEINFO = 0x06400007 # MNG_LONG_FORM

	# ud2 instruction bytes.
	UD2_INSN_BYTES = b'\x0F\x0B'

	# Netnode to determine if ELF is for Prospero.
	PROSPERO_NET_NODE = '$ prospero'

	class UiHooks(ida_kernwin.UI_Hooks):
		def __init__(self, plugin):
			super().__init__()

			self.plugin = plugin

		def ready_to_run(self, *args):
			# UI ready to run (called multiple times).
			return 0

		def database_inited(self, is_new_database, idc_script):
			# It is called multiple times, not really useful.
			return 0

		def plugin_loaded(self, plugin_info):
			#print('Loading plugin: %s' % plugin_info.name)
			return 0

	class IdbHooks(ida_idp.IDB_Hooks):
		def __init__(self, plugin):
			super().__init__()

			self.plugin = plugin

		def loader_finished(self, *args):
			# External file loader finished its work.
			# Use this event to augment the existing loader functionality.
			return 0

		def determined_main(self, ea):
			return 0

		def segm_added(self, seg):
			return 0

		def auto_empty_finally(self, *args):
			self.plugin.post_initial_analysis()
			return 0

	class IdpHooks(ida_idp.IDP_Hooks):
		def __init__(self, plugin):
			super().__init__()

			self.plugin = plugin

		def ev_func_bounds(self, possible_ret_code, func, max_func_end_ea):
			self.plugin.fixup_func_bounds(func, max_func_end_ea)
			return 0

	def __init__(self):
		super().__init__()

		self.elf = None
		self.file_type = None
		self.lib_versions = None
		self.prodg_meta_data = None
		self.soname = None
		self.orig_file_path = None
		self.needed_modules = None
		self.modules = None
		self.libraries = None
		self.relocation_type = None
		self.rela_reloc_table = None
		self.jmprel_reloc_table = None
		self.symbol_table = None
		self.string_table = None
		self.hash_table = None
		self.id_table = None
		self.got_start_ea = ida_idaapi.BADADDR
		self.got_plt_start_ea = ida_idaapi.BADADDR
		self.init_proc_ea = ida_idaapi.BADADDR
		self.term_proc_ea = ida_idaapi.BADADDR
		self.nids = None
		self.symbols = None

		self.ui_hooks = ps5_elf_plugin_t.UiHooks(self)
		self.idb_hooks = ps5_elf_plugin_t.IdbHooks(self)
		self.idp_hooks = ps5_elf_plugin_t.IdpHooks(self)

	def init(self):
		# Cannot be used in terminal mode.
		if not ida_kernwin.is_idaq():
			return ida_idaapi.PLUGIN_SKIP

		if not ida_hexrays.init_hexrays_plugin():
			return ida_idaapi.PLUGIN_SKIP

		print(f'Initializing plugin: {self.wanted_name}')

		file_path = ida_nalt.get_input_file_path()
		file_name = ida_nalt.get_root_filename()

		# Sanity check.
		if ida_ida.inf_get_filetype() != idc.FT_ELF or ida_ida.inf_get_procname() != 'metapc' or ida_ida.inf_is_be() or not ida_ida.inf_is_64bit():
			return ida_idaapi.PLUGIN_SKIP

		# Load needed type info libraries and register standard types.
		idc.add_default_til('gnuunx64')

		standard_types = ['Elf64_Ehdr', 'Elf64_Phdr', 'Elf64_Shdr', 'Elf64_Nhdr', 'Elf64_Rel', 'Elf64_Rela', 'Elf64_Dyn', 'Elf64_Sym']
		for type_name in standard_types:
			idc.import_type(-1, type_name)

		# Read and parse ELF header.
		elf = ElfUtil()
		if elf.is_inited():
			ehdr = elf.ehdr
			is_just_loaded = False
		else:
			ehdr_struct_name = 'Elf64_Ehdr'
			ehdr_size = get_struct_size(ehdr_struct_name)

			phdr_struct_name = 'Elf64_Phdr'
			phdr_size = get_struct_size(phdr_struct_name)

			is_prospero_elf = False

			try:
				with open(file_path, 'rb') as f:
					data = f.read(ehdr_size)

					while True:
						if len(data) != ehdr_size:
							break
						ehdr = parse_struct(ehdr_struct_name, data)

						phdr_offset = ehdr['e_phoff']
						if phdr_offset <= 0:
							break
						f.seek(phdr_offset)

						data = f.read(phdr_size)
						if len(data) != phdr_size:
							break
						phdr = parse_struct(phdr_struct_name, data)
						phdr_type, phdr_flags = ElfUtil.phdr_type(phdr), phdr['p_flags']

						if phdr_type != ElfUtil.PT_LOAD or phdr_flags != ElfUtil.PF_EXEC:
							break

						is_prospero_elf = True
						break
			except Exception as e:
				#print('Got exception during header parsing attempt:', e)
				pass

			if not is_prospero_elf:
				return ida_idaapi.PLUGIN_SKIP
			else:
				node = ida_netnode.netnode()
				node.create(ps5_elf_plugin_t.PROSPERO_NET_NODE)

			is_just_loaded = True

		# Determine file type.
		file_type_str = {
			ElfUtil.ET_SCE_EXEC_ASLR: 'Executable',
			ElfUtil.ET_SCE_DYNAMIC: 'PRX',
		}.get(ehdr['e_type'], None)

		if file_type_str is None:
			return ida_idaapi.PLUGIN_SKIP

		self.file_type = ehdr['e_type']

		print('File type: %s' % file_type_str)

		# Reset members.
		self.lib_versions = {}
		self.prodg_meta_data = {}
		self.soname = None
		self.orig_file_path = None
		self.needed_modules = []
		self.modules = {}
		self.libraries = {}
		self.relocation_type = None
		self.rela_reloc_table = None
		self.jmprel_reloc_table = None
		self.symbol_table = None
		self.string_table = None
		self.hash_table = None
		self.id_table = None
		self.got_start_ea = ida_idaapi.BADADDR
		self.got_plt_start_ea = ida_idaapi.BADADDR
		self.init_proc_ea = ida_idaapi.BADADDR
		self.term_proc_ea = ida_idaapi.BADADDR

		# Load additional type info libraries.
		for name in ['prospero']:
			idc.add_default_til(name)

		# Load known NIDS.
		self.nids = load_known_nids('ps5_symbols.txt')

		self.symbols = []

		# Load additional modules.
		ida_idaapi.require('gcc_extab')

		# Set up analyzer on the first load.
		if is_just_loaded:
			self.setup_analysis()
		else:
			self.elf = elf

		self.ui_hooks.hook()
		self.idb_hooks.hook()
		self.idp_hooks.hook()

		return ida_idaapi.PLUGIN_KEEP

	def term(self):
		#print(f'Terminating plugin: {self.wanted_name}')

		self.idp_hooks.unhook()
		self.idb_hooks.unhook()
		self.ui_hooks.unhook()

	def setup_analysis(self):
		# Set up common parameters.
		ida_ida.inf_set_ostype(0x6) # BSD OS
		ida_ida.inf_set_demnames(ida_ida.DEMNAM_NAME | ida_ida.DEMNAM_GCC3) # use GCC mangling names

		# Set up compiler parameters.
		ida_ida.inf_set_cc_id(ida_typeinf.COMP_GNU)
		ida_ida.inf_set_cc_cm(ida_typeinf.CM_N64 | ida_typeinf.CM_M_NN | ida_typeinf.CM_CC_CDECL)
		ida_ida.inf_set_cc_size_b(1)
		ida_ida.inf_set_cc_size_s(2)
		ida_ida.inf_set_cc_size_i(4)
		ida_ida.inf_set_cc_size_e(4)
		ida_ida.inf_set_cc_size_l(8)
		ida_ida.inf_set_cc_size_l(8)
		ida_ida.inf_set_cc_size_ldbl(8)
		ida_ida.inf_set_cc_defalign(0)

		# Set up analysis parameters.
		ida_ida.inf_set_mark_code(False) # Do not find functions inside .data segments.
		ida_ida.inf_set_create_func_tails(False) # Don not create function tails.
		ida_ida.inf_set_noflow_to_data(True) # Control flow to data segment is ignored.
		ida_ida.inf_set_rename_jumpfunc(False) # Rename jump functions as J_.

	def _fixup_segment(self, seg):
		image_base = ida_nalt.get_imagebase()

		name = ida_segment.get_segm_name(seg)
		type_id = ida_segment.segtype(seg.start_ea)

		print('Fixing up segment at 0x%x (type: %d, perm: 0x%x).' % (seg.start_ea, type_id, seg.perm))

		if type_id == ida_segment.SEG_CODE:
			other_seg = ida_segment.get_segm_by_name('.text')
			if seg.start_ea == image_base and seg.perm == ida_segment.SEGPERM_EXEC and not other_seg:
				ida_segment.set_segm_name(seg, '.text')
				print('Found .text segment.')
				return True
		elif type_id == ida_segment.SEG_DATA:
			other_seg = ida_segment.get_segm_by_name('.rodata')
			if seg.perm == ida_segment.SEGPERM_READ:
				if not other_seg:
					ida_segment.set_segm_name(seg, '.rodata')
					print('Found .rodata segment.')
					return True
				else:
					if name.lower().strip() == 'note':
						print('Deleting note segment.')
						ida_segment.del_segm(seg.start_ea, ida_segment.SEGMOD_KILL | ida_segment.SEGMOD_SILENT)
						return True
			elif seg.perm == ida_segment.SEGPERM_READ | ida_segment.SEGPERM_WRITE:
				# There are multiple R/W segments and we need more info to recognize them, so skip now and process them later.
				return False
			elif seg.perm == 0:
				other_seg = ida_segment.get_segm_by_name('.dynsym')
				if not other_seg:
					ida_segment.set_segm_name(seg, '.dynsym')
					print('Found .dynsym segment.')
					return True
		elif type_id == ida_segment.SEG_XTRN:
			other_seg = ida_segment.get_segm_by_name('extern')
			if seg.perm == 0 and not other_seg:
				ida_segment.set_segm_name(seg, 'extern')
				print('Found extern segment.')
				return True

		return False

	def _parse_extra_segments(self):
		assert self.elf.is_inited()

		file_path = ida_nalt.get_input_file_path()

		result = False

		dynamic_phdr = self.elf.find_phdr_by_type(ElfUtil.PT_DYNAMIC)

		if dynamic_phdr is not None:
			result |= self._parse_dynamic_segment(dynamic_phdr)

		comment_phdr = self.elf.find_phdr_by_type(ElfUtil.PT_SCE_COMMENT)
		version_phdr = self.elf.find_phdr_by_type(ElfUtil.PT_SCE_VERSION)

		if not comment_phdr and not version_phdr:
			return False

		with open(file_path, 'rb') as f:
			if comment_phdr is not None:
				f.seek(comment_phdr['p_offset'])
				comment_data = f.read(comment_phdr['p_filesz'])
				if len(comment_data) != comment_phdr['p_filesz']:
					comment_data = None
			else:
				comment_data = None

			if version_phdr is not None:
				f.seek(version_phdr['p_offset'])
				version_data = f.read(version_phdr['p_filesz'])
				if len(version_data) != version_phdr['p_filesz']:
					version_data = None
			else:
				version_data = None

		if comment_data:
			result |= self._parse_comment_segment(comment_data)
		if version_data:
			result |= self._parse_version_segment(version_data)

		return result

	def _parse_dynamic_segment(self, dynamic_phdr):
		print('Processing dynamic segment.')

		struct_name = 'Elf64_Dyn'
		struct_size = get_struct_size(struct_name)

		seg = ida_segment.get_segm_by_name('.dynsym')
		if not seg:
			ida_kernwin.warning('Unable to find .dynsym segment, cannot parse dynamic segment.')
			return False
		dynsym_base_ea = seg.start_ea

		ea = dynamic_phdr['p_vaddr']
		end_ea = dynamic_phdr['p_vaddr'] + dynamic_phdr['p_memsz']

		dyns = []
		while ea < end_ea:
			data = ida_bytes.get_bytes(ea, struct_size)
			if len(data) != struct_size:
				raise RuntimeError('Insufficient data of %s structure: 0x%x (expected: 0x%x)' % (struct_name, len(data), struct_size))

			dyn = parse_struct(struct_name, data)
			if dyn['d_tag'] == ElfUtil.DT_NULL:
				break
			dyns.append(dyn)

			ea += struct_size

		if not dyns:
			print('No dynamic tags found.')
			return True

		self.rela_reloc_table = RelaRelocTable()
		self.jmprel_reloc_table = JmpRelRelocTable()
		self.symbol_table = SymbolTable()
		self.string_table = StringTable()
		self.hash_table = HashTable()
		self.id_table = IdTable()

		print('Dynamic tags:')
		for dyn in dyns:
			tag, value = dyn['d_tag'], dyn['d_un']
			print('  %s: 0x%x' % (ElfUtil.stringify_dyn_tag(tag), value))

			if tag == ElfUtil.DT_NEEDED:
				name = read_cstring_at(dynsym_base_ea + value)
				self.needed_modules.append(name)
			elif tag == ElfUtil.DT_SONAME:
				self.soname = read_cstring_at(dynsym_base_ea + value)
			elif tag in [ElfUtil.DT_SCE_NEEDED_MODULE, ElfUtil.DT_SCE_NEEDED_MODULE_PPR]:
				module_id = ObjectInfo.obj_id(value)
				if module_id not in self.modules:
					self.modules[module_id] = ObjectInfo()
				self.modules[module_id].set_info(value)
				self.modules[module_id].update_name(dynsym_base_ea)
			elif tag in [ElfUtil.DT_SCE_EXPORT_LIB, ElfUtil.DT_SCE_EXPORT_LIB_PPR]:
				library_id = ObjectInfo.obj_id(value)
				if library_id not in self.libraries:
					self.libraries[library_id] = ObjectInfo()
				self.libraries[library_id].set_info(value, True)
				self.libraries[library_id].update_name(dynsym_base_ea)
			elif tag == ElfUtil.DT_SCE_IMPORT_LIB_ATTR:
				library_id = ObjectInfo.obj_id(value)
				if library_id not in self.libraries:
					self.libraries[library_id] = ObjectInfo()
				self.libraries[library_id].set_attr(value)
			elif tag in [ElfUtil.DT_SCE_MODULE_INFO, ElfUtil.DT_SCE_MODULE_INFO_PPR]:
				module_id = ObjectInfo.obj_id(value)
				if module_id not in self.modules:
					self.modules[module_id] = ObjectInfo()
				self.modules[module_id].set_info(value, True)
				self.modules[module_id].update_name(dynsym_base_ea)
			elif tag == ElfUtil.DT_SCE_MODULE_ATTR:
				module_id = ObjectInfo.obj_id(value)
				if module_id not in self.modules:
					self.modules[module_id] = ObjectInfo()
				self.modules[module_id].set_attr(value)
			elif tag in [ElfUtil.DT_SCE_ORIGINAL_FILENAME, ElfUtil.DT_SCE_ORIGINAL_FILENAME_PPR]:
				self.orig_file_path = read_cstring_at(dynsym_base_ea + value)
			elif tag in [ElfUtil.DT_SCE_IMPORT_LIB, ElfUtil.DT_SCE_IMPORT_LIB_PPR]:
				library_id = ObjectInfo.obj_id(value)
				if library_id not in self.libraries:
					self.libraries[library_id] = ObjectInfo()
				self.libraries[library_id].set_info(value)
				self.libraries[library_id].update_name(dynsym_base_ea)
			elif tag == ElfUtil.DT_SCE_EXPORT_LIB_ATTR:
				library_id = ObjectInfo.obj_id(value)
				if library_id not in self.libraries:
					self.libraries[library_id] = ObjectInfo()
				self.libraries[library_id].set_attr(value)
			elif tag == ElfUtil.DT_RELA: # ELF RELA Relocation Table
				ea = as_uint64(value)
				if ea != 0 and ea != ida_idaapi.BADADDR:
					self.rela_reloc_table.ea = ea
			elif tag == ElfUtil.DT_RELASZ: # ELF RELA Relocation Table
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					self.rela_reloc_table.size = size
			elif tag == ElfUtil.DT_RELAENT: # ELF RELA Relocation Table
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					assert size == get_struct_size(self.rela_reloc_table.struct_name())
					self.rela_reloc_table.entry_size = size
			elif tag == ElfUtil.DT_RELACOUNT: # ELF RELA Relocation Table
				count = as_uint64(value)
				if count != ida_idaapi.BADADDR:
					# TODO: Why is it smaller than actual count?
					self.rela_reloc_table.entry_count = count
			elif tag == ElfUtil.DT_JMPREL: # ELF JMPREL Relocation Table
				ea = as_uint64(value)
				if ea != 0 and ea != ida_idaapi.BADADDR:
					self.jmprel_reloc_table.ea = ea
			elif tag == ElfUtil.DT_PLTRELSZ: # ELF JMPREL Relocation Table
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					self.jmprel_reloc_table.size = size
			elif tag == ElfUtil.DT_PLTGOT:
				ea = as_uint64(value)
				if ea != 0 and ea != ida_idaapi.BADADDR:
					self.got_plt_start_ea = ea
			elif tag == ElfUtil.DT_PLTREL:
				self.relocation_type = as_uint32(value)
				if self.relocation_type != ElfUtil.DT_REL and self.relocation_type != ElfUtil.DT_RELA:
					ida_kernwin.warning('Unsupported PLT relocation type: 0x%x' % self.relocation_type)
			elif tag == ElfUtil.DT_SYMTAB: # ELF Symbol Table
				ea = as_uint64(value)
				if ea != 0 and ea != ida_idaapi.BADADDR:
					self.symbol_table.ea = ea
			elif tag == ElfUtil.DT_SCE_SYMTABSZ: # ELF Symbol Table
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					self.symbol_table.size = size
			elif tag == ElfUtil.DT_SYMENT: # ELF Symbol Table
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					assert size == get_struct_size(self.symbol_table.struct_name())
					self.symbol_table.entry_size = size
			elif tag == ElfUtil.DT_STRTAB: # ELF String Table
				ea = as_uint64(value)
				if ea != 0 and ea != ida_idaapi.BADADDR:
					self.string_table.ea = ea
			elif tag == ElfUtil.DT_STRSZ: # ELF String Table
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					self.string_table.size = size
			elif tag == ElfUtil.DT_HASH: # ELF Hash Table
				ea = as_uint64(value)
				if ea != 0 and ea != ida_idaapi.BADADDR:
					self.hash_table.ea = ea
			elif tag == ElfUtil.DT_SCE_HASHSZ: # ELF Hash Table
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					self.hash_table.size = size
			elif tag == ElfUtil.DT_SCE_IDTABENTSZ: # ELF ID Table
				# TODO: Where are ea/size tags?
				size = as_uint64(value)
				if size != ida_idaapi.BADADDR:
					# TODO: assert size == get_struct_size(self.id_table.struct_name())
					self.id_table.entry_size = size
			elif tag == ElfUtil.DT_INIT:
				self.init_proc_ea = value
			elif tag == ElfUtil.DT_FINI:
				self.term_proc_ea = value
			elif tag == ElfUtil.DT_PREINIT_ARRAY:
				# TODO
				continue
			elif tag == ElfUtil.DT_PREINIT_ARRAYSZ:
				# TODO
				continue
			elif tag == ElfUtil.DT_INIT_ARRAY:
				# TODO
				continue
			elif tag == ElfUtil.DT_INIT_ARRAYSZ:
				# TODO
				continue
			elif tag == ElfUtil.DT_FINI_ARRAY:
				# TODO
				continue
			elif tag == ElfUtil.DT_FINI_ARRAYSZ:
				# TODO
				continue

		if self.jmprel_reloc_table.entry_size == ida_idaapi.BADADDR:
			self.jmprel_reloc_table.entry_size = get_struct_size(self.jmprel_reloc_table.struct_name())

		if self.rela_reloc_table.entry_size == ida_idaapi.BADADDR:
			self.rela_reloc_table.entry_size = get_struct_size(self.rela_reloc_table.struct_name())

		if self.id_table.entry_size != ida_idaapi.BADADDR and self.id_table.entry_size != 0x8:
			ida_kernwin.warning('Unsupported ID table entry size: 0x%x' % self.id_table.entry_size)

		return True

	def _parse_comment_segment(self, data):
		print('Processing comment segment.')

		f = BytesIO(data)

		while True:
			key = f.read(struct.calcsize('4s'))
			if len(key) < struct.calcsize('4s'):
				# Reached end of file.
				break
			key = key.rstrip(b'\0').decode('ascii')

			data = f.read(struct.calcsize('2I'))
			if len(data) != struct.calcsize('2I'):
				ida_kernwin.warning('Truncated data at comment segment.')
				return False

			max_length, length = struct.unpack('<2I', data)
			value = f.read(length)

			if len(value) != length:
				ida_kernwin.warning('Truncated data at comment segment.')
				return False

			# Try to decode value as UTF-8 string.
			try:
				value = value.decode('utf-8').rstrip('\0')
			except:
				pass

			self.prodg_meta_data[key] = value

		params = {
			'PATH': 'Original path',
		}

		for key, desc in params.items():
			if key not in self.prodg_meta_data:
				continue
			print('%s: %s' % (desc, self.prodg_meta_data[key]))

		return True

	def _parse_version_segment(self, data):
		print('Processing version segment.')

		f = BytesIO(data)

		while True:
			data = f.read(struct.calcsize('2H'))
			if data == b'':
				# Reached end of file.
				break
			elif len(data) != struct.calcsize('2H'):
				ida_kernwin.warning('Truncated data at version segment.')
				return False

			reserved, length = struct.unpack('<2H', data)
			assert reserved == 0
			if length == 0:
				continue

			data = f.read(length)
			if len(data) != length:
				ida_kernwin.warning('Truncated data at version segment.')
				return False

			type_id, data = data[0], data[1:]
			if type_id == 0x8:
				name, version = data.split(b':')
				name = name.decode('ascii')
				version = version.hex().upper() # TODO: Need to parse version properly.
				print('Library %s version: %s' % (name, version))
				self.lib_versions[name] = version
			else:
				ida_kernwin.warning('Unknown type id 0x%x in version info.' % type_id)
				continue

		return True

	def _fixup_padding_segment(self):
		seg = ida_segment.get_segm_by_name('.sce_padding')
		if not seg:
			image_base = ida_nalt.get_imagebase()

			has_padding = ida_bytes.get_bytes(image_base, ElfUtil.SCE_PADDING_SIZE) == ElfUtil.SCE_PADDING
			if not has_padding:
				return False

			text_seg = ida_segment.get_segm_by_name('.text')
			if not text_seg:
				ida_kernwin.warning('Unable to find .text segment, cannot fixup .sce_padding segment.')
				return False

			if text_seg.start_ea == image_base:
				print('Moving start of .text segment from 0x%x to 0x%x.' % (text_seg.start_ea, text_seg.start_ea + ElfUtil.SCE_PADDING_SIZE))
				ida_segment.set_segm_start(text_seg.start_ea, text_seg.start_ea + ElfUtil.SCE_PADDING_SIZE, ida_segment.SEGMOD_KILL | ida_segment.SEGMOD_SILENT)

			print('Creating .sce_padding segment.')
			seg = ida_segment.segment_t()
			seg.start_ea = image_base
			seg.end_ea = image_base + ElfUtil.SCE_PADDING_SIZE
			seg.bitness = text_seg.bitness
			seg.type = ida_segment.SEG_UNDF
			seg.perm = 0
			ida_segment.add_segm_ex(seg, '.sce_padding', None, ida_segment.ADDSEG_NOAA)

		seg = ida_segment.get_segm_by_name('.sce_padding')
		if not seg:
			return False

		ida_auto.auto_mark_range(seg.start_ea, seg.end_ea, ida_auto.AU_UNK)
		ida_bytes.del_items(seg.start_ea, ida_bytes.DELIT_SIMPLE, ElfUtil.SCE_PADDING_SIZE)

		print('Found .sce_padding segment.')

		return True

	def _link_segments_with_phdrs(self):
		num_segments = ida_segment.get_segm_qty()
		print('Number of segments: %d' % num_segments)

		for i in range(num_segments):
			seg = ida_segment.getnseg(i)
			if not seg:
				continue
			idx = ida_segment.get_segm_num(seg.start_ea)

			phdr = self.elf.find_phdr_by_seg(seg)
			if not phdr:
				continue

	def _fixup_segment_perms(self):
		print('Fixing up segments permissions.')

		seg, last_seg = ida_segment.get_first_seg(), ida_segment.get_last_seg()

		while seg:
			name = ida_segment.get_segm_name(seg)

			if name in ['.text', '.init', '.fini', '.plt'] or name.startswith('.text.'):
				seg.perm = ida_segment.SEGPERM_EXEC
				need_update = True
			else:
				need_update = False

			if not need_update:
				print('Updating %s segment permissions.' % name)
				ida_segment.update_segm(seg)

			seg = ida_segment.get_next_seg(seg.start_ea)
			if seg == last_seg:
				break

		return True

	def _fixup_init_fini_segments(self):
		print('Fixing up .init and .fini segments.')

		info = { '.init_proc': '.init', '.term_proc': '.fini' }
		segments = {}

		for func_name, segment_name in info.items():
			seg = ida_segment.get_segm_by_name(segment_name)
			if seg:
				continue

			ea = ida_name.get_name_ea(ida_idaapi.BADADDR, func_name)
			if ea == ida_idaapi.BADADDR:
				ida_kernwin.warning('Unable to find %s function address, cannot fixup %s segment.' % (func_name, segment_name))
				continue

			func = ida_funcs.get_func(ea)
			if not func:
				ida_kernwin.warning('Unable to find %s function, cannot fixup %s segment.' % (func_name, segment_name))
				continue
			start_ea, end_ea = func.start_ea, func.end_ea

			text_seg = ida_segment.get_segm_by_name('.text')
			if not text_seg:
				ida_kernwin.warning('Unable to find .text segment, cannot fixup %s segment.' % segment_name)
				continue

			if segment_name == '.init':
				end_ea = align_up(end_ea, 0x10)
				print('Moving start of .text segment from 0x%x to 0x%x.' % (text_seg.start_ea, end_ea))
				ida_segment.set_segm_start(text_seg.start_ea, end_ea, ida_segment.SEGMOD_KEEP | ida_segment.SEGMOD_SILENT)
			elif segment_name == '.fini':
				start_ea = align_up(start_ea, 0x10)
				print('Moving end of .text segment from 0x%x to 0x%x.' % (text_seg.end_ea, start_ea))
				ida_segment.set_segm_end(text_seg.start_ea, start_ea, ida_segment.SEGMOD_KEEP | ida_segment.SEGMOD_SILENT)

			seg = ida_segment.segment_t()
			seg.start_ea = start_ea
			seg.end_ea = end_ea
			seg.bitness = text_seg.bitness
			seg.type = text_seg.type
			seg.perm = text_seg.perm
			segments[segment_name] = seg

		text_seg = ida_segment.get_segm_by_name('.text')
		if not text_seg:
			ida_kernwin.warning('Unable to find .text segment, cannot fixup .init and .proc segments.')
			return False

		for segment_name, seg in segments.items():
			print('Creating %s segment.' % segment_name)
			ida_segment.add_segm_ex(seg, segment_name, ida_segment.get_segm_class(text_seg), ida_segment.ADDSEG_NOSREG)

		return True

	def _fixup_eh_segments(self):
		assert self.elf.is_inited()

		print('Fixing up .eh_frame and .eh_frame_hdr segments.')

		seg = ida_segment.get_segm_by_name('.eh_frame')
		if seg:
			# Segment already exists, skipping it.
			return True

		seg = ida_segment.get_segm_by_name('.eh_frame_hdr')
		if not seg:
			seg = ida_segment.get_segm_by_name('.rodata')
			if not seg:
				ida_kernwin.warning('Unable to find .rodata segment, cannot fixup .eh_frame_hdr segment.')
				return False

			phdr = self.elf.find_phdr_by_type(ElfUtil.PT_GNU_EH_FRAME)
			if phdr is None:
				ida_kernwin.warning('Unable to find program header for segment .eh_frame_hdr, cannot fixup it.')
				return False

			new_seg = ida_segment.segment_t()
			new_seg.start_ea = phdr['p_vaddr']
			new_seg.end_ea = phdr['p_vaddr'] + phdr['p_memsz']
			new_seg.bitness = seg.bitness
			new_seg.type = seg.type
			new_seg.perm = seg.perm

			print('Creating .eh_frame_hdr segment.')
			ida_segment.add_segm_ex(new_seg, '.eh_frame_hdr', ida_segment.get_segm_class(seg), ida_segment.ADDSEG_NOSREG)

			seg = ida_segment.get_segm_by_name('.eh_frame_hdr')

		if not seg:
			ida_kernwin.warning('Unable to find .eh_frame_hdr segment, cannot fixup .eh_frame segment.')
			return False

		ea = seg.start_ea

		exc_data_base_ea, exc_version = ea, gcc_extab.format_byte(ea, 'version')
		ea += struct.calcsize('B')

		exc_eh_frame_ptr_enc, ea = gcc_extab.format_enc(ea, 'eh frame ptr encoding')
		exc_fde_count_enc, ea = gcc_extab.format_enc(ea, 'fde count encoding')
		exc_ent_table_enc, ea = gcc_extab.format_enc(ea, 'ent binary table encoding')
		exc_eh_frame_ptr, ea = gcc_extab.read_enc_val(ea, exc_eh_frame_ptr_enc, True, data_ea=exc_data_base_ea)

		if exc_eh_frame_ptr != ida_idaapi.BADADDR and exc_eh_frame_ptr < seg.start_ea:
			new_seg = ida_segment.segment_t()
			new_seg.start_ea = exc_eh_frame_ptr
			new_seg.end_ea = seg.start_ea
			new_seg.bitness = seg.bitness
			new_seg.type = seg.type
			new_seg.perm = seg.perm

			print('Creating .eh_frame segment.')
			ida_segment.add_segm_ex(new_seg, '.eh_frame', ida_segment.get_segm_class(seg), ida_segment.ADDSEG_NOSREG)

		return True

	def _fixup_param_segment(self):
		assert self.elf.is_inited()

		if self.file_type == ElfUtil.ET_SCE_EXEC_ASLR:
			phdr_type = ElfUtil.PT_SCE_PROCPARAM
			segment_name = '.sce_process_param'
			struct_name = 'sceProcessParam'
			handler_cb = self._fixup_process_param_segment
		else:
			phdr_type = ElfUtil.PT_SCE_MODULE_PARAM
			segment_name = '.sce_module_param'
			struct_name = 'sceModuleParam'
			handler_cb = self._fixup_module_param_segment

		phdr = self.elf.find_phdr_by_type(phdr_type)
		if phdr is None:
			ida_kernwin.warning('Unable to find program header for segment %s, cannot fixup it.' % segment_name)
			return False

		seg = ida_segment.get_segm_by_name(segment_name)
		if not seg:
			seg = ida_segment.get_segm_by_name('.rodata')
			if not seg:
				ida_kernwin.warning('Unable to find .rodata segment, cannot fixup %s segment.' % segment_name)
				return False

			new_seg = ida_segment.segment_t()
			new_seg.start_ea = phdr['p_vaddr']
			new_seg.end_ea = align_up(phdr['p_vaddr'] + phdr['p_memsz'], 0x10)

			new_seg.bitness = seg.bitness
			new_seg.type = seg.type
			new_seg.perm = seg.perm

			print('Creating %s segment.' % segment_name)
			ida_segment.add_segm_ex(new_seg, segment_name, ida_segment.get_segm_class(seg), ida_segment.ADDSEG_NOSREG)

			seg = ida_segment.get_segm_by_name(segment_name)

		if not seg:
			ida_kernwin.warning('Unable to find %s segment, cannot fixup it.' % segment_name)
			return False

		print('Processing %s segment.' % segment_name)

		size = ida_bytes.get_qword(seg.start_ea)
		if size == ida_idaapi.BADADDR or size < struct.calcsize('Q'):
			ida_kernwin.warning('Unexpected size of %s structure.' % struct_name)
			return False
		print('%s structure size: 0x%x' % (struct_name, size))

		end_ea = align_up(seg.start_ea + size, 0x10)
		if seg.end_ea != end_ea:
			print('Moving end of %s segment from 0x%x to 0x%x.' % (segment_name, seg.end_ea, end_ea))
			ida_segment.set_segm_end(seg.start_ea, end_ea, ida_segment.SEGMOD_KEEP | ida_segment.SEGMOD_SILENT)

		data = ida_bytes.get_bytes(seg.start_ea, size)
		if len(data) != size:
			raise RuntimeError('Insufficient data of %s structure: 0x%x (expected: 0x%x)' % (struct_name, len(data), size))

		return handler_cb(segment_name, struct_name, data[struct.calcsize('Q'):])

	def _fixup_process_param_segment(self, segment_name, struct_name, data):
		fmt = '<4s3I5Q'
		extra_fmt = '<3Q'
		data_size = len(data)
		expected_size = struct.calcsize(fmt)
		expected_extra_size = struct.calcsize(extra_fmt)
		if data_size < expected_size:
			raise RuntimeError('Unsupported size of %s structure: 0x%x (expected: 0x%x)' % (struct_name, data_size, expected_size))
		elif data_size > expected_size + expected_extra_size:
			ida_kernwin.warning('Size of %s structure is larger than expected: 0x%x (expected: 0x%x)' % (struct_name, data_size, expected_size + expected_extra_size))

		# TODO: Check these fields.
		magic, entry_count, sdk_version, unk1, process_name_ea, user_main_thread_name_ea, user_main_thread_priority_ea, user_main_thread_stack_size_ea, libc_param_ea = struct.unpack(fmt, data[:expected_size])
		if magic != ElfUtil.SCE_PROCESS_PARAM_MAGIC:
			raise RuntimeError('Invalid magic in %s structure: 0x%08x' % (struct_name, magic))
		offset = expected_size
		data_size -= offset

		# TODO: Check if it is really an address and not a value.
		kernel_mem_param_ea = ida_idaapi.BADADDR
		if data_size >= 0x8:
			kernel_mem_param_ea, = struct.unpack('<Q', data[offset:offset + struct.calcsize('Q')])
			offset += struct.calcsize('Q')
			data_size -= struct.calcsize('Q')

		# TODO: Check if it is really an address and not a value.
		kernel_fs_param_ea = ida_idaapi.BADADDR
		if data_size >= 0x8:
			kernel_fs_param_ea, = struct.unpack('<Q', data[offset:offset + struct.calcsize('Q')])
			offset += struct.calcsize('Q')
			data_size -= struct.calcsize('Q')

		# TODO: Check if it is really an address and not a value.
		process_preload_enabled_ea = ida_idaapi.BADADDR
		if data_size >= 0x8:
			process_preload_enabled_ea, = struct.unpack('<Q', data[offset:offset + struct.calcsize('Q')])
			offset += struct.calcsize('Q')
			data_size -= struct.calcsize('Q')

		# TODO: Create proper structure.
		print('Process info:')
		print('  Magic: 0x%s' % magic.hex())
		print('  Entry count: %d' % entry_count)
		print('  SDK version: 0x%x' % sdk_version)
		print('  Unk1: 0x%x' % unk1)
		print('  Process name ea: 0x%x' % process_name_ea)
		print('  User main thread ea: 0x%x' % user_main_thread_name_ea)
		print('  User main thread priority ea: 0x%x' % user_main_thread_priority_ea)
		print('  User main thread stack size ea: 0x%x' % user_main_thread_stack_size_ea)
		print('  Libc param ea: 0x%x' % libc_param_ea)
		if kernel_mem_param_ea != ida_idaapi.BADADDR:
			print('  Kernel mem param ea: 0x%x' % kernel_mem_param_ea)
		if kernel_fs_param_ea != ida_idaapi.BADADDR:
			print('  Kernel fs param ea: 0x%x' % kernel_fs_param_ea)
		if process_preload_enabled_ea != ida_idaapi.BADADDR:
			print('  Process preload enabled ea: 0x%x' % process_preload_enabled_ea)

		return True

	def _fixup_module_param_segment(self, segment_name, struct_name, data):
		fmt = '<4sIQ2I'
		data_size = len(data)
		expected_size = struct.calcsize(fmt)
		if data_size < expected_size:
			raise RuntimeError('Unsupported size of %s structure: 0x%x (expected: 0x%x)' % (struct_name, data_size, expected_size))
		elif data_size > expected_size:
			ida_kernwin.warning('Size of %s structure is larger than expected: 0x%x (expected: 0x%x)' % (struct_name, data_size, expected_size))

		# TODO: Check these fields.
		magic, entry_count, sdk_version, unk1, unk2 = struct.unpack(fmt, data[:expected_size])
		if magic != ElfUtil.SCE_MODULE_PARAM_MAGIC:
			raise RuntimeError('Invalid magic in %s structure: 0x%08x' % (struct_name, magic))

		# TODO: Create proper structure.
		print('Module info:')
		print('  Magic: 0x%s' % magic.hex())
		print('  Entry count: %d' % entry_count)
		print('  SDK version: 0x%x' % sdk_version)
		print('  Unk1: 0x%x' % unk1)
		print('  Unk2: 0x%x' % unk2)

		return True

	def _fixup_data_segment(self):
		seg = ida_segment.get_segm_by_name('.data')
		if seg:
			# Segment already exists, skipping it.
			return False

		seg = self._find_last_rw_seg()
		if not seg:
			ida_kernwin.warning('Unable to find R/W segment, cannot fixup .data segment.')
			return False

		seg_name = ida_segment.get_segm_name(seg)
		if seg_name.startswith('.'):
			ida_kernwin.warning('R/W segment starts with dot already, cannot fixup .data segment.')
			return False

		ida_segment.set_segm_name(seg, '.data')

		return True

	def _fixup_extra_segments(self):
		print('Fixing up extra .data segments.')

		first_seg, last_seg = ida_segment.get_first_seg(), ida_segment.get_last_seg()

		seg = first_seg
		while seg:
			name = ida_segment.get_segm_name(seg)
			sclass = ida_segment.get_segm_class(seg)
			idx = ida_segment.get_segm_num(seg.start_ea)

			if name.lower() == 'load' and not name.startswith('.') and sclass == 'DATA':
				print('Renaming extra R/W %s segment #%03d to .data.' % (name, idx))
				ida_segment.set_segm_name(seg, '.data')

			seg = ida_segment.get_next_seg(seg.start_ea)
			if seg == last_seg:
				break

		print('Merging similar neighboring segments.')

		seg1 = first_seg
		while seg1:
			name1 = ida_segment.get_segm_name(seg1)
			sclass1 = ida_segment.get_segm_class(seg1)
			idx1 = ida_segment.get_segm_num(seg1.start_ea)

			#print('Processing segment #%03d: %s' % (idx1, name1))

			finished = False
			while not finished:
				seg2 = ida_segment.get_next_seg(seg1.start_ea)
				if not seg2:
					finished = True
					break
				is_last = seg2 == last_seg

				name2 = ida_segment.get_segm_name(seg2)
				sclass2 = ida_segment.get_segm_class(seg2)
				idx2 = ida_segment.get_segm_num(seg2.start_ea)

				#print('Comparing with segment #%03d: %s' % (idx2, name2))

				if name1 != name2 or seg1.perm != seg2.perm or seg1.end_ea != seg2.start_ea:
					#print('Merging done: params mismatch')
					finished = True
					break

				print('Merging segments #%03d(%s) and #%03d(%s): [0x%x;0x%x) / [0x%x;0x%x)' % (idx1, name1, idx2, name2, seg1.start_ea, seg1.end_ea, seg2.start_ea, seg2.end_ea))

				end_ea = seg2.end_ea
				assert end_ea >= seg1.end_ea

				ida_segment.del_segm(seg2.start_ea, ida_segment.SEGMOD_KEEP)
				ida_segment.set_segm_end(seg1.start_ea, end_ea, ida_segment.SEGMOD_KEEP)
				ida_segment.update_segm(seg1)

				if is_last:
					#print('Merging done: was last segment')
					break

			seg1 = ida_segment.get_next_seg(seg1.start_ea)
			if seg1 == last_seg:
				break

		return True

	def _fixup_got_segments(self):
		print('Fixing up .got and .got.plt segments.')

		result = False

		if self.got_plt_start_ea != ida_idaapi.BADADDR:
			print('Address of .got.plt section: 0x%x' % self.got_plt_start_ea)

			seg = ida_segment.get_segm_by_name('.got.plt')
			if not seg:
				seg = ida_segment.getseg(self.got_plt_start_ea)
				if not seg:
					ida_kernwin.warning('Unable to find segment which includes .got.plt, cannot fixup .got.plt segment.')
					return False

				new_seg = ida_segment.segment_t()
				new_seg.start_ea = self.got_plt_start_ea
				new_seg.end_ea = seg.end_ea
				new_seg.bitness = seg.bitness
				new_seg.type = seg.type
				new_seg.perm = seg.perm

				print('Creating .got.plt segment.')
				ida_segment.add_segm_ex(new_seg, '.got.plt', ida_segment.get_segm_class(seg), 0)
			else:
				# Segment already exists, skipping it.
				pass

			result |= True

		if self.got_start_ea != ida_idaapi.BADADDR:
			print('Address of .got section: 0x%x' % self.got_start_ea)

			seg = ida_segment.get_segm_by_name('.got')
			if not seg:
				seg = ida_segment.getseg(self.got_start_ea)
				if not seg:
					ida_kernwin.warning('Unable to find segment which includes .got, cannot fixup .got segment.')
					return False

				new_seg = ida_segment.segment_t()
				new_seg.start_ea = self.got_start_ea
				new_seg.end_ea = seg.end_ea
				new_seg.bitness = seg.bitness
				new_seg.type = seg.type
				new_seg.perm = seg.perm

				print('Creating .got segment.')
				ida_segment.add_segm_ex(new_seg, '.got', ida_segment.get_segm_class(seg), 0)
			else:
				# Segment already exists, skipping it.
				pass

			result |= True

		return result

	def _fixup_ctors_dtors_segments(self):
		print('Fixing up .ctors and .dtors segments.')

		# The linker must build two lists of these functions - a list of initialization functions, called __CTOR_LIST__, and a list of termination functions, called __DTOR_LIST__.
		#
		# Each list always begins with an ignored function pointer (which may hold 0, -1, or a count of the function pointers after it, depending on the environment).
		# This is followed by a series of zero or more function pointers to constructors (or destructors), followed by a function pointer containing zero.

		data_seg = ida_segment.get_segm_by_name('.data')
		if not data_seg:
			ida_kernwin.warning('Unable to find .data segment, cannot fixup .ctors and .dtors segments.')
			return False
		seg_class, seg_bitness, seg_type, seg_perm = ida_segment.get_segm_class(data_seg), data_seg.bitness, data_seg.type, data_seg.perm

		ea_pair = self._fixup_dtors_segment()
		if not ea_pair:
			return False
		dtors_start_ea, dtors_end_ea = ea_pair

		if dtors_start_ea != ida_idaapi.BADADDR and dtors_end_ea != ida_idaapi.BADADDR:
			seg = ida_segment.segment_t()
			seg.start_ea = dtors_start_ea
			seg.end_ea = dtors_end_ea
			seg.bitness = seg_bitness
			seg.type = seg_type
			seg.perm = seg_perm

			print('Creating .dtors segment.')
			ida_segment.add_segm_ex(seg, '.dtors', seg_class, 0)
		else:
			ida_kernwin.warning('Unable to find .dtors segment.')

		ea_pair = self._fixup_ctors_segment()
		if not ea_pair:
			return False
		ctors_start_ea, ctors_end_ea = ea_pair

		if ctors_start_ea != ida_idaapi.BADADDR and ctors_end_ea != ida_idaapi.BADADDR:
			seg = ida_segment.segment_t()
			seg.start_ea = ctors_start_ea
			seg.end_ea = ctors_end_ea
			seg.bitness = seg_bitness
			seg.type = seg_type
			seg.perm = seg_perm

			print('Creating .ctors segment.')
			ida_segment.add_segm_ex(seg, '.ctors', seg_class, 0)
		else:
			ida_kernwin.warning('Unable to find .ctors segment.')

		return True

	def _fixup_ctors_segment(self):
		# static func_ptr __CTOR_LIST__[1] = { (func_ptr) (-1) };
		# ...
		# static func_ptr __CTOR_END__[1] = { (func_ptr) 0 };
		#
		# static void __do_global_ctors_aux()
		# {
		#   func_ptr* p;
		#   for (p = __CTOR_END__ - 1; *p != (func_ptr)-1; --p)
		#     (*p)();
		# }

		dtors_seg = ida_segment.get_segm_by_name('.dtors')
		if not dtors_seg:
			ida_kernwin.warning('Unable to find .dtors segment, cannot fixup .ctors segment.')
			return None
		dtors_start_ea, dtors_end_ea = dtors_seg.start_ea, dtors_seg.end_ea
		if dtors_start_ea == ida_idaapi.BADADDR or dtors_end_ea == ida_idaapi.BADADDR:
			ida_kernwin.warning('Unexpected .dtors segment addresses, cannot fixup .ctors segment.')
			return None

		ea = ida_name.get_name_ea(ida_idaapi.BADADDR, '.init_proc')
		if ea == ida_idaapi.BADADDR:
			ida_kernwin.warning('Unable to find .init_proc, cannot fixup .ctors segment.')
			return None

		preinit_array_end_ea = ida_idaapi.BADADDR
		got_plt_end_ea = ida_idaapi.BADADDR
		cmp_found = mov_found = lea_found = False
		mov_reg = None
		for ea in idautils.FuncItems(ea):
			if not cmp_found and check_insn_format(ea, 'cmp', [(ida_ua.o_reg, None), (ida_ua.o_mem, None)]):
				value = idc.get_operand_value(ea, 1)
				if value != ida_idaapi.BADADDR and preinit_array_end_ea == ida_idaapi.BADADDR:
					preinit_array_end_ea = value
				cmp_found = True
			elif not mov_found and check_insn_format(ea, 'mov', [(ida_ua.o_reg, None), (ida_ua.o_phrase, None)]):
				# XXX: Cannot use ida_ua.print_operand because it returns garbage data.
				mov_reg = idc.print_operand(ea, 1).lower().strip().lstrip('[').rstrip(']')
				mov_found = True
			elif not lea_found and mov_found and mov_reg is not None and check_insn_format(ea, 'lea', [(ida_ua.o_reg, mov_reg), (ida_ua.o_mem, None)]):
				value = idc.get_operand_value(ea, 1)
				if value != ida_idaapi.BADADDR and got_plt_end_ea == ida_idaapi.BADADDR:
					got_plt_end_ea = value
				lea_found = True

		ctors_end_ea = dtors_start_ea - 0x8
		if ida_bytes.get_qword(ctors_end_ea) != 0:
			raise RuntimeError('Unexpected end of constructors table.')
		ida_bytes.create_qword(ctors_end_ea, 0x8, True)
		ida_name.set_name(ctors_end_ea, '__CTOR_END__', ida_name.SN_NOCHECK)
		ctors_start_ea = ctors_end_ea - 0x8
		while ida_bytes.get_qword(ctors_start_ea) != ida_idaapi.BADADDR:
			ctors_start_ea -= 0x8
		ida_bytes.create_qword(ctors_start_ea, 0x8, True)
		ida_name.set_name(ctors_start_ea, '__CTOR_LIST__', ida_name.SN_NOCHECK)
		ctors_end_ea += 0x8

		if preinit_array_end_ea != ida_idaapi.BADADDR:
			ida_name.set_name(preinit_array_end_ea, '_G__preinit_array_end', ida_name.SN_NOCHECK)

		return (ctors_start_ea, ctors_end_ea)

	def _fixup_dtors_segment(self):
		# static func_ptr __DTOR_LIST__[1] = { (func_ptr) (-1) };
		# ...
		# static func_ptr __DTOR_END__[1] = { (func_ptr) 0 };
		#
		# static void __do_global_dtors_aux()
		# {
		#   func_ptr* p;
		#   for (p = __DTOR_LIST__ + 1; *p; ++p)
		#     (*p)();
		# }

		ea = ida_name.get_name_ea(ida_idaapi.BADADDR, '.term_proc')
		if ea == ida_idaapi.BADADDR:
			ida_kernwin.warning('Unable to find .term_proc, cannot fixup .dtors segment.')
			return None

		last_lea_value = None
		for ea in idautils.FuncItems(ea):
			if check_insn_format(ea, 'cmp', [(ida_ua.o_mem, None), (ida_ua.o_imm, 0)]):
				value = idc.get_operand_value(ea, 0)
				if value != ida_idaapi.BADADDR:
					self.got_start_ea = value
			elif check_insn_format(ea, 'lea', [(ida_ua.o_reg, None), (ida_ua.o_mem, None)]):
				value = idc.get_operand_value(ea, 1)
				if value != ida_idaapi.BADADDR:
					last_lea_value = value
			elif check_insn_format(ea, 'add', [(ida_ua.o_reg, None), (ida_ua.o_imm, None)]):
				value = idc.get_operand_value(ea, 1)
				if value != 0x8:
					continue
				if last_lea_value is None:
					raise RuntimeError('Unexpected instructions at .term_proc().')
				break
		if last_lea_value is None:
			raise RuntimeError('Unexpected instructions at .term_proc().')

		# last_lea_value should equals to: __DTOR_LIST__ + 0x10.
		dtors_start_ea = last_lea_value - 0x10
		if ida_bytes.get_qword(dtors_start_ea) != ida_idaapi.BADADDR:
			raise RuntimeError('Unexpected start of destructors table.')
		ida_bytes.create_qword(dtors_start_ea, 0x8, True)
		ida_name.set_name(dtors_start_ea, '__DTOR_LIST__', ida_name.SN_NOCHECK)

		dtors_end_ea = dtors_start_ea + 0x8
		while ida_bytes.get_qword(dtors_end_ea) != 0:
			dtors_end_ea += 0x8
		ida_bytes.create_qword(dtors_end_ea, 0x8, True)
		ida_name.set_name(dtors_end_ea, '__DTOR_END__', ida_name.SN_NOCHECK)
		dtors_end_ea += 0x8

		return (dtors_start_ea, dtors_end_ea)

	def _fixup_bss_segment(self):
		seg = ida_segment.get_segm_by_name('.bss')
		if seg:
			# Segment already exists, skipping it.
			return False

		# We need to find last segment with R/W permissions.
		data_seg = self._find_last_rw_seg()
		if not data_seg:
			return False
		data_segment_name = ida_segment.get_segm_name(data_seg)

		# XXX: ida_bytes.next_that is not working as expected and returns address after .bss.
		bss_start_ea = ida_idaapi.BADADDR
		ea = data_seg.start_ea
		while ea != ida_idaapi.BADADDR and ea < data_seg.end_ea:
			if not idc.is_loaded(ea):
				bss_start_ea = ea
				break
			ea = ida_bytes.next_addr(ea)
		if bss_start_ea == ida_idaapi.BADADDR:
			return False
		bss_end_ea = data_seg.end_ea

		print('Creating .bss segment.')
		seg = ida_segment.segment_t()
		seg.start_ea = bss_start_ea
		seg.end_ea = bss_end_ea
		seg.type = ida_segment.SEG_BSS
		seg.bitness = data_seg.bitness
		seg.perm = data_seg.perm
		ida_segment.add_segm_ex(seg, '.bss', ida_segment.get_segm_class(data_seg), ida_segment.ADDSEG_NOSREG)

		return True

	def fixup_func_bounds(self, func, max_func_end_ea):
		end_ea = func.end_ea

		data = ida_bytes.get_bytes(end_ea, len(ps5_elf_plugin_t.UD2_INSN_BYTES))
		if not data or data != ps5_elf_plugin_t.UD2_INSN_BYTES:
			return

		end_ea += len(data)

		print('Setting function 0x%x end to 0x%x (old: 0x%x).' % (func.start_ea, end_ea, func.end_ea))
		func.end_ea = end_ea

		ida_funcs.reanalyze_function(func, func.start_ea, end_ea, False)

	def _fixup_symbols(self):
		print('Fixing up symbols.')

		if not self.symbol_table or not self.symbol_table.is_table_loaded():
			ida_kernwin.warning('Symbol table is not loaded, cannot fixup symbols.')
			return False

		if not self.string_table or not self.string_table.is_loaded():
			ida_kernwin.warning('String table is not loaded, cannot fixup symbols.')
			return False

		for i in range(self.symbol_table.get_num_entries()):
			entry = self.symbol_table.get_entry(i)
			if entry is None:
				ida_kernwin.warning('No entry for symbol table entry #%d.' % i)
				return False

			symbol = SymbolTable.Symbol(entry)

			self.symbols.append(symbol)

			if not symbol.is_object() and not symbol.is_func():
				continue

			mangled_name = self.string_table.get_string(symbol.entry['st_name'])
			if not mangled_name:
				continue

			symbol_name_enc, lid_enc, mid_enc = mangled_name.split('#')
			nid, lid, mid = ObjectInfo.decode_nid(symbol_name_enc), ObjectInfo.decode_obj_id(lid_enc), ObjectInfo.decode_obj_id(mid_enc)

			if mid not in self.modules:
				ida_kernwin.warning('No module with ID: 0x%x' % mid)
				return False
			module_name = self.modules[mid].name

			assert lid in self.libraries
			if lid not in self.libraries:
				ida_kernwin.warning('No library with ID: 0x%x' % lid)
				return False
			library_name = self.libraries[lid].name
			is_export = self.libraries[lid].is_export

			symbol_name = self.nids[nid] if nid in self.nids else nid

			symbol.set_descriptor(module_name, library_name, symbol_name, is_export)

		return True

	def _fixup_plt_segment(self):
		print('Fixing up .plt segment.')

		if not self.jmprel_reloc_table.is_table_loaded():
			ida_kernwin.warning('Jmprel relocation table is not loaded, cannot fixup .plt segment.')
			return False

		if not self.string_table.is_loaded():
			ida_kernwin.warning('String table is not loaded, cannot fixup .plt segment.')
			return False

		jmprel_entry_count = self.jmprel_reloc_table.get_num_entries()

		got_plt_seg = ida_segment.get_segm_by_name('.got.plt')
		if not got_plt_seg:
			ida_kernwin.warning('Unable to find .got.plt segment, cannot fixup .plt segment.')
			return False

		target_ea = got_plt_seg.start_ea + struct.calcsize('Q')
		xrefs = list(idautils.XrefsTo(target_ea, ida_xref.XREF_DATA))
		if not xrefs:
			ida_kernwin.warning('Unable to find xrefs to .got.plt segment, cannot fixup .plt segment.')
			return False
		xref_type, plt_start_ea = xrefs[0].type, xrefs[0].frm
		assert xref_type == ida_xref.dr_R

		base_insns = bytes.fromhex(
			'FF 35 00 00 00 00' + # push cs:<ea>
			'FF 25 00 00 00 00' + # jmp cs:<ea>
			''
		)

		stub_insns = bytes.fromhex(
			'FF 25 00 00 00 00' + # jmp qword ptr [rip]
			'68 00 00 00 00'    + # push 0
			'E9 00 00 00 00'    + # jmp 0x5
			''
		)

		# Segment should start with: push cs:<ea>
		data = ida_bytes.get_bytes(plt_start_ea, 2)
		if data[:2] != base_insns[:2]:
			ida_kernwin.warning('Unexpected .plt segment data, cannot fixup .plt segment.')
			return False

		super_seg = ida_segment.getseg(plt_start_ea)
		if not super_seg:
			ida_kernwin.warning('Unable to find segment which includes .plt, cannot fixup .plt segment.')
			return False

		# Need to find: jmp cs:<ea>
		plt_base_ea = ida_bytes.find_bytes('FF 25', range_start = plt_start_ea + len(base_insns), range_end = super_seg.end_ea, flags = ida_search.SEARCH_DOWN | ida_search.SEARCH_CASE)
		if plt_base_ea == ida_idaapi.BADADDR:
			ida_kernwin.warning('Unable to find .plt base ea, cannot fixup .plt segment.')
			return False

		plt_end_ea = align_up(plt_base_ea + jmprel_entry_count * len(stub_insns), 0x10)

		seg = ida_segment.get_segm_by_name('.plt')
		if not seg:
			new_seg = ida_segment.segment_t()
			new_seg.start_ea = plt_start_ea
			new_seg.end_ea = plt_end_ea
			new_seg.bitness = super_seg.bitness
			new_seg.type = super_seg.type
			new_seg.perm = super_seg.perm

			print('Creating .plt segment.')
			ida_segment.add_segm_ex(new_seg, '.plt', ida_segment.get_segm_class(super_seg), 0)
		else:
			# Segment already exists, skipping it.
			pass

		idaldr_node = ida_netnode.netnode('$ IDALDR node for ids loading $')
		if not idaldr_node:
			ida_kernwin.warning('Unable to find netnode for imports.')

		for i in range(jmprel_entry_count):
			entry = self.jmprel_reloc_table.get_entry(i)
			if entry is None:
				ida_kernwin.warning('No entry for jmprel relocation table entry #%d.' % i)
				return False

			record = JmpRelRelocTable.Record(entry)

			reloc_type = record.get_type()
			if not reloc_type in [JmpRelRelocTable.R_AMD64_JUMP_SLOT]:
				ida_kernwin.warning('Unsupported relocation type 0x%x for jmprel relocation table entry #%d.' % (reloc_type, i))
				return False
			if reloc_type != self.relocation_type:
				ida_kernwin.warning('Mismatched relocation type 0x%x (should be 0x%x) for jmprel relocation table entry #%d.' % (reloc_type, self.relocation_type, i))
				return False

			symbol_idx = record.get_symbol_idx()
			if symbol_idx >= len(self.symbols):
				ida_kernwin.warning('Symbol index #%d out of range for jmprel relocation table entry #%d.' % (symbol_idx, i))
				return False
			symbol = self.symbols[symbol_idx]

			if not symbol.has_descriptor():
				ida_kernwin.warning('Symbol #%d has no descriptor for jmprel relocation table entry #%d.' % (symbol_idx, i))
				return False

			name = symbol.get_name()
			name_ex = symbol.get_name_ex()
			comment = symbol.get_name_comment()

			stub_name = '/B%s' % name_ex
			stub_ptr_name = '/PG%s' % name_ex

			stub_ptr_ea = record.entry['r_offset']
			stub_ea = ida_bytes.get_qword(record.entry['r_offset'])
			func_ea = plt_base_ea + i * len(stub_insns)

			#print('Renaming stub pointer %s to %s at 0x%x.' % (ida_name.get_name(stub_ptr_ea), stub_ptr_name, stub_ptr_ea))
			ida_name.set_name(stub_ptr_ea, stub_ptr_name, ida_name.SN_NOCHECK)
			ida_bytes.set_cmt(stub_ptr_ea, '', False)

			#print('Renaming stub %s to %s at 0x%x.' % (ida_name.get_name(stub_ea), stub_name, stub_ea))
			ida_name.set_name(stub_ea, stub_name, ida_name.SN_NOCHECK)
			ida_bytes.set_cmt(stub_ea, '', False)

			func = ida_funcs.get_func(func_ea)
			if not func:
				ida_funcs.add_func(func_ea, ida_idaapi.BADADDR)

			#print('Renaming function %s to %s at 0x%x.' % (ida_name.get_name(func_ea), name, func_ea))
			ida_name.set_name(func_ea, name, ida_name.SN_NOCHECK)
			ida_bytes.set_cmt(func_ea, comment, False)

			func = ida_funcs.get_func(stub_ea)
			if not func:
				ida_funcs.add_func(stub_ea, ida_idaapi.BADADDR)

			ea = ida_name.get_name_ea(func_ea, name)
			if ea != ida_idaapi.BADADDR:
				func = ida_funcs.get_func(ea)
				if func:
					func.flags |= ida_funcs.FUNC_LIB
					ida_funcs.update_func(func)

			ea = ida_name.get_name_ea(stub_ea, stub_name)
			if ea != ida_idaapi.BADADDR:
				func = ida_funcs.get_func(ea)
				if func:
					func.flags |= ida_funcs.FUNC_LIB
					ida_funcs.update_func(func)

			if idaldr_node:
				idaldr_node.supset_ea(stub_ea, stub_name, ida_netnode.stag)

		return True

	def _fixup_relocations(self):
		print('Fixing up relocations.')

		if not self.rela_reloc_table.is_table_loaded():
			ida_kernwin.warning('Rela relocation table is not loaded, cannot fixup relocations.')
			return False

		if not self.string_table.is_loaded():
			ida_kernwin.warning('String table is not loaded, cannot fixup relocations.')
			return False

		idaldr_node = ida_netnode.netnode('$ IDALDR node for ids loading $')
		if not idaldr_node:
			ida_kernwin.warning('Unable to find netnode for imports.')

		rela_entry_count = self.rela_reloc_table.get_num_entries()

		for i in range(rela_entry_count):
			entry = self.rela_reloc_table.get_entry(i)
			if entry is None:
				ida_kernwin.warning('No entry for rela relocation table entry #%d.' % i)
				return False

			record = RelaRelocTable.Record(entry)
			reloc_type = record.get_type()

			ea, addend = as_uint64(record.entry['r_offset']), as_uint64(record.entry['r_addend'])

			if reloc_type in [RelaRelocTable.R_AMD64_GLOB_DAT, RelaRelocTable.R_AMD64_64]:
				symbol_idx = record.get_symbol_idx()
				if symbol_idx < len(self.symbols):
					symbol = self.symbols[symbol_idx]

					if symbol.has_descriptor():
						name = symbol.get_name()
						if name:
							ea = ida_bytes.get_qword(ea)
							if ea != ida_idaapi.BADADDR:
								#print('Renaming symbol at 0x%x to %s.' % (ea, name))
								ida_name.set_name(ea, name, ida_name.SN_NOCHECK)

								if idaldr_node:
									idaldr_node.supset_ea(ea, name, ida_netnode.stag)
					else:
						# TODO: Is it correct?
						#print('Warning! Symbol #%d has no descriptor for rela relocation table entry #%d.' % (symbol_idx, i))
						continue
				else:
					# TODO: Is it correct?
					#print('Warning! Rela relocation table entry #%d have invalid symbol idx #%d.' % (i, symbol_idx))
					continue
			else:
				print('Warning! Unsupported relocation type 0x%x for rela relocation table entry #%d.' % (reloc_type, i))
				continue

		return True

	def _fixup_exports(self):
		print('Fixing up exports.')

		ea_ordinal_map = {}

		for i in range(ida_entry.get_entry_qty()):
			ordinal = ida_entry.get_entry_ordinal(i)
			ea = ida_entry.get_entry(ordinal)
			ea_ordinal_map[ea] = ordinal

		for i, symbol in enumerate(self.symbols):
			if not symbol.is_export and not symbol.is_object() and not symbol.is_func():
				continue

			if not symbol.has_descriptor():
				continue

			ea, size = symbol.entry['st_value'], symbol.entry['st_size']
			if ea == 0 or ea == ida_idaapi.BADADDR:
				continue

			func = ida_funcs.get_func(ea)
			if not func:
				ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)
				ida_ua.create_insn(ea)
				ida_funcs.add_func(ea, ea + size)

			name = symbol.get_name()

			print('Setting name %s to exported function at 0x%x.' % (name, ea))
			if ea in ea_ordinal_map:
				ordinal = ea_ordinal_map[ea]
				ida_entry.rename_entry(ordinal, name, ida_entry.AEF_UTF8)
			else:
				ida_name.set_name(ea, name)
			ida_bytes.set_cmt(ea, '', False)

		return True

	def _fixup_dynsym_segment(self):
		print('Deleting .dynsym segment.')

		seg = ida_segment.get_segm_by_name('.dynsym')
		if not seg:
			ida_kernwin.warning('Unable to find .dynsym segment, cannot fixup .dynsym segment.')
			return False

		ida_segment.del_segm(seg.start_ea, ida_segment.SEGMOD_KILL | ida_segment.SEGMOD_SILENT)

		return True

	def _mark_noret_funcs(self):
		names = [
			'exit', 'exit1', 'abort',
			'__stack_chk_fail',
			'_ZNSt9bad_allocD0Ev', '_ZNSt9bad_allocD1Ev', '_ZNSt9bad_allocD2Ev', '_ZSt11_Xbad_allocv',
			'_ZNSt16invalid_argumentD0Ev', '_ZNSt16invalid_argumentD1Ev', '_ZNSt16invalid_argumentD2Ev', '_ZSt18_Xinvalid_argumentPKc',
			'_ZNSt12length_errorD0Ev', '_ZNSt12length_errorD1Ev', '_ZNSt12length_errorD2Ev', '_ZSt14_Xlength_errorPKc',
			'_ZNSt12out_of_rangeD0Ev', '_ZNSt12out_of_rangeD1Ev', '_ZNSt12out_of_rangeD2Ev', '_ZSt14_Xout_of_rangePKc',
			'_ZNSt14overflow_errorD0Ev', '_ZNSt14overflow_errorD1Ev', '_ZNSt14overflow_errorD2Ev', '_ZSt16_Xoverflow_errorPKc',
			'_ZNSt13runtime_errorD0Ev', '_ZNSt13runtime_errorD1Ev', '_ZNSt13runtime_errorD2Ev', '_ZSt15_Xruntime_errorPKc',
			'_ZNSt17bad_function_callD0Ev', '_ZNSt17bad_function_callD1Ev', '_ZNSt17bad_function_callD2Ev', '_ZSt19_Xbad_function_callv',
			'_ZNSt11regex_errorD0Ev', '_ZNSt11regex_errorD1Ev', '_ZNSt11regex_errorD2Ev',
			'_ZSt10_Rng_abortPKc',
			'_ZSt19_Throw_future_errorRKSt10error_code',
			'_ZSt25_Rethrow_future_exceptionPv', '_ZSt25_Rethrow_future_exceptionSt13exception_ptr',
		]

		for name in names:
			ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name)
			if ea == ida_idaapi.BADADDR:
				continue

			func = ida_funcs.get_func(ea)
			if not func:
				continue

			func.flags |= ida_funcs.FUNC_NORET
			ida_funcs.update_func(func)

			ida_auto.reanalyze_callers(ea, True)

	def _find_last_rw_seg(self):
		rw_seg = None

		seg, first_seg = ida_segment.get_last_seg(), ida_segment.get_first_seg()
		while seg and seg != first_seg:
			name = ida_segment.get_segm_name(seg)
			sclass = ida_segment.get_segm_class(seg)
			if seg.perm == ida_segment.SEGPERM_READ | ida_segment.SEGPERM_WRITE:
				rw_seg = seg
				break
			seg = ida_segment.get_prev_seg(seg.start_ea)

		return rw_seg

	def post_initial_analysis(self):
		self.elf = ElfUtil()
		if not self.elf.is_inited():
			raise RuntimeError('Netnode for elf is not initialized.')

		print('Performing post initial auto analysis.')

		for i in range(ida_segment.get_segm_qty()):
			seg = ida_segment.getnseg(i)
			if seg:
				self._fixup_segment(seg)

		self._parse_extra_segments()
		self._fixup_segment_perms()
		self._link_segments_with_phdrs()
		self._fixup_padding_segment()
		self._fixup_param_segment()
		self._fixup_data_segment()
		self._fixup_init_fini_segments()
		self._fixup_eh_segments()
		self._fixup_ctors_dtors_segments()
		self._fixup_got_segments()
		self._fixup_bss_segment()
		self._fixup_extra_segments()
		self._fixup_symbols()

		# Allow to rename jump functions now because we did set up correct symbol names already.
		ida_ida.inf_set_rename_jumpfunc(True) # Rename jump functions as J_.

		self._fixup_plt_segment()
		self._fixup_relocations()
		self._fixup_exports()
		self._fixup_dynsym_segment()

		self._mark_noret_funcs()

		if self.soname is not None:
			print('Name: %s' % self.soname)
		if self.orig_file_path is not None:
			print('Original file path: %s' % self.orig_file_path)

		if self.init_proc_ea != ida_idaapi.BADADDR:
			print('Address of .init_proc function: 0x%x' % self.init_proc_ea)
		if self.term_proc_ea != ida_idaapi.BADADDR:
			print('Address of .term_proc function: 0x%x' % self.term_proc_ea)

		if self.needed_modules:
			print('Needed modules: %s' % ', '.join(self.needed_modules))
		for id, info in self.modules.items():
			print('Module #%03d: %s' % (id, repr(info)))
		for id, info in self.libraries.items():
			print('Library #%03d: %s' % (id, repr(info)))

		if self.relocation_type is not None:
			print('Relocation type: 0x%x' % self.relocation_type)

	def run(self, arg):
		ida_kernwin.warning('Running as script is not possible.')
		return False

def PLUGIN_ENTRY():
	return ps5_elf_plugin_t()


