# plugins

repository: BraveCattle/ctf
file: tools/ida_pro/IDA_Pro_7.7/python/examples/hexrays/vds17.py
wanted_name: Structure offsets (IDAPython)
comment: Sample plugin17 for Hex-Rays decompiler
created_at: 2024-01-23 08:17:51+00:00
pushed_at: 2024-10-23 09:31:30+00:00
url:
https://github.com/BraveCattle/ctf/blob/5858d0da5c048eb8468ff8be1d63fbc50d61b601/tools/ida_pro/IDA_Pro_7.7/python/examples/hexrays/vds17.py

content:
"""
summary: using the "Select offsets" widget

description:
  Registers an action opens the "Select offsets" widget
  (select_udt_by_offset() call).

  This effectively repeats the functionality already available
  through Alt+Y.

  Place cursor on the union field and press Shift+T
"""

import ida_idaapi
import ida_hexrays
import ida_lines
import ida_typeinf
import ida_kernwin

# --------------------------------------------------------------------------
class func_stroff_ah_t(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def activate(self, ctx):
        # get the current item
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        vu.get_current_item(ida_hexrays.USE_KEYBOARD)

        # REGION1, will be referenced later
        # check that the current item is a union field
        if not vu.item.is_citem():
            ida_kernwin.warning("Please position the cursor on a union member")
            return 0
        e = vu.item.e
        while True:
            op = e.op
            if op != ida_hexrays.cot_memptr and op != ida_hexrays.cot_memref:
                ida_kernwin.warning("Please position the cursor on a union member")
                return 0
            e = e.x
            if op == ida_hexrays.cot_memptr:
                if e.type.is_union():
                    break
            else:
                if ida_typeinf.remove_pointer(e.type).is_union():
                    break
            if not e.type.is_udt():
                ida_kernwin.warning("Please position the cursor on a union member")
                return 0
        # END REGION1

        # REGION2
        # calculate the member offset
        off = 0
        e = vu.item.e
        while True:
            e2 = e.x
            tif = ida_typeinf.remove_pointer(e2.type)
            if not tif.is_union():
                off += e.m
            e = e2
            if e2.op != ida_hexrays.cot_memptr and e2.op != ida_hexrays.cot_memref:
                break
            if not e2.type.is_udt():
                break
        # END REGION2

        # REGION3
        # go up and collect more member references (in order to calculate the final offset)
        p = vu.item.e
        while True:
            p2 = vu.cfunc.body.find_parent_of(p)
            if p2.op == ida_hexrays.cot_memptr:
                break
            if p2.op == ida_hexrays.cot_memref:
                e2 = p2.cexpr
                tif = ida_typeinf.remove_pointer(e2.x.type)
                if not tif.is_union():
                    off += e2.m
                p = p2
                continue
            if p2.op == ida_hexrays.cot_ref:
                # handle &a.b + N (this expression may appear if the user previously selected
                #                  a wrong field)
                delta = 0
                add = vu.cfunc.body.find_parent_of(p2)
                if add.op == ida_hexrays.cot_cast:
                    add = vu.cfunc.body.find_parent_of(add)
                if add.op == ida_hexrays.cot_add and add.y.op == ida_hexrays.cot_num:
                    delta = add.y.numval()
                    objsize = add.type.get_ptrarr_objsize()
                    nbytes = delta * objsize
                    off += nbytes
            # we could use helpers like WORD/BYTE/... to calculate a more precise offset
            # if ( p2->op == cot_call && (e2->exflags & EXFL_LVALUE) != 0 )
            break
        # END REGION3

        # REGION4
        ea = vu.item.e.ea
        # the item itself may be unaddressable.
        # TODO: find its addressable parent
        if ea == ida_idaapi.BADADDR:
            ida_kernwin.warning("Sorry, the current item is not addressable")
            return 0
        # END REGION4

        # REGION5
        # prepare the text representation for the item,
        # use the neighborhoods of cursor
        line = ida_lines.tag_remove(ida_kernwin.get_custom_viewer_curline(vu.ct, False))
        line_len = len(line)
        x = max(0, vu.cpos.x - 10)
        l = min(10, line_len - vu.cpos.x) + 10
        line = line[x:x+l]
        # END REGION5

        # REGION6
        ops = ida_hexrays.ui_stroff_ops_t()
        op = ops.push_back()
        op.offset = off
        op.text = line
        # END REGION6

        # REGION7
        class set_union_sel_t(ida_hexrays.ui_stroff_applicator_t):
            def __init__(self, ea):
                ida_hexrays.ui_stroff_applicator_t.__init__(self)
                self.ea = ea

            def apply(self, opnum, path, top_tif, spath):
                typename = ida_typeinf.print_tinfo('', 0, 0, ida_typeinf.PRTYPE_1LINE, top_tif, '', '')
                ida_kernwin.msg("User selected %s of type %s\n" % (spath, typename))
                if path.empty():
                    return False
                vu.cfunc.set_user_union_selection(self.ea, path)
                vu.cfunc.save_user_unions()
                return True
        # END REGION7

        # REGION8
        su = set_union_sel_t(ea)
        res = ida_hexrays.select_udt_by_offset(None, ops, su)
        if res != 0:
            # regenerate ctree
            vu.refresh_view(True)
        # END REGION8

        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


# --------------------------------------------------------------------------
# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Structure offsets (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin17 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            print("Hex-rays version %s has been detected, Structure offsets ready to use" % ida_hexrays.get_hexrays_version())
            ida_kernwin.register_action(
                ida_kernwin.action_desc_t(
                    "vds17:strchoose",
                    "Structure offsets",
                    func_stroff_ah_t(),
                    "Shift+T"))
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        pass
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()

"""
# A few notes about the VDS17 sample

You can find two VDS17 samples in the IDA Pro install directory:

    python/examples/hexrays/vds17.py
    plugins/hexrays_sdk/plugins/vds17

The former is an IDAPython plugin and the latter is a C++ IDA Pro plugin.
Actually they have the same functionality.
Just to be more concrete the vds17.py plugin will be used below.


## How the user interface works

Let us suppose that we have the following local types:

    1 C     struct {int c0;int c1;}
    2 U     union {int u0;__int16 u1;}
    3 D     struct {int d0;C d1;U d2;}
    4 E     struct {int e0;D e1;}
    5 res_t struct {int r0;int r1;__int16 r2;}

and the decompiler generates the following pseudocode:

    void __cdecl f(res_t *r)
    {
      r->r0 = e->e1.d1.c0;
      r->r1 = e->e1.d2.u0;
      r->r2 = e->e1.d2.u1;
    }

As we see, it looks good, the decompiler did a really good job.

But let us imagine that we need reference to the ```e1.d2.u0``` union member on the last line.

At first we need to load the VDS17 plugin. For that,
place the cursor at the ```u1``` union member on the last line
and use ```Shift-T```.

The "Structure offsets" dialog appears on the screen.
The left pane of the dialog contains the available local types
and the right pane has the following view:

    [checked] e->e1.d2.u1; | 10h | |

Use the left pane to select ```E```, expand it, and select (Double-click) on the ```u0``` member.

The right pane will change to:

    [checked] e->e1.d2.u1; | 10h | [checke] E.e1.d2.u0 |

Since this is what we want, press the ```OK``` button and the pseudocode gets changed:

    void __cdecl f(res_t *r)
    {
      r->r0 = e->e1.d1.c0;
      r->r1 = e->e1.d2.u0;
      r->r2 = e->e1.d2.u0;
    }


### What do we need this plugin for?

All union members have the same offset in the parent structure.
The decompiler selects the first suitable union member.
Sometimes we need to change this selection and use another union member.


## API details

Let us look into ```python/examples/hexrays/vds17.py```.

* REGION1: check that the cursor points to the union member
* REGION2: calculate the member offset
* REGION3: sometimes there is a need to adjust the offset
* REGION4: the pointed item must be addressable

Then the magic begins:

* REGION5: we need something to show in the ```Operand``` column of the right pane.
    The text around the cursor looks like a good compromise.
* REGION6: we are ready to fill the right pane: we have calculated the offset and
    have prepared the descriptive text for the operand.
* REGION7: it is a callback that informs us about the user selection (will be described later)
* REGION8: activate the "Structure offsets" dialog, let the user select,
    the callback specified above will update the union member, refresh pseudocode


### The devil is in the details

The main part of callback is the ```apply``` method (REGION7).
It receives two arguments:

1. ```opnum``` is the number of the selected operand (line) in the right pane of the dialog.
     We need something to map the line number to our operand.
     In the case of VDS17 ```ea``` performs this role.
2. ```path```  the path that describes the union selection.
3. ```top_tif``` typeinfo of the top-level UDT which user selected
4. ```spath``` the field names path to the selected member

The union selection path is denotes a concrete member inside a UDT.

For structure types there is no need in the union selection path because the member
offset uniquely denotes the desired member.

For unions, on the other hand, the member offset is not enough because all union
members start at the same offset zero. For them, we remember the ordinal number of the
union member in the path. E.g., for the local types given above, the following holds:

* e1.d1.c0 is denoted by an empty path because it does not have any unions
* e1.d2.u0 is denoted by path consisting of 0: we use the first member of U
* e1.d2.u1 is denoted by path consisting of 1: we use the second member of U

You can retrieve the union selection path using API call ```get_user_union_selection```
or apply it using ```set_user_union_selection```.
"""


repository: carsond135/malwareanalysis
file: tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/hexrays/vds17.py
description: NKU恶意代码分析
wanted_name: Structure offsets (IDAPython)
comment: Sample plugin17 for Hex-Rays decompiler
created_at: 2023-06-25 19:16:29+00:00
pushed_at: 2022-12-18 09:34:14+00:00
url:
https://github.com/carsond135/malwareanalysis/blob/45c8bf862a54239ff1e7c91313eccf9cf7d4a35b/tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/hexrays/vds17.py


repository: Ph0en1x-XMU/Ph0en1x-Team
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2022-06-02 11:06:16+00:00
pushed_at: 2022-06-02 11:11:40+00:00
file:
Reverse/ida/IDA Pro 7.5.201028 (x86,x64,ARM,ARM64,PPC,PPC64,MIPS)/python/examples/core/install_user_defined_prefix.py

url:
https://github.com/Ph0en1x-XMU/Ph0en1x-Team/blob/a8d80aaea8227e418f5826bf05ef69a1bc8d84f1/Reverse/ida/IDA%20Pro%207.5.201028%20(x86,x64,ARM,ARM64,PPC,PPC64,MIPS)/python/examples/core/install_user_defined_prefix.py

content:
from __future__ import print_function

import ida_lines
import ida_idaapi

PREFIX = ida_lines.SCOLOR_INV + ' ' + ida_lines.SCOLOR_INV

class my_user_prefix_t(ida_lines.user_defined_prefix_t):
    def get_user_defined_prefix(self, ea, insn, lnnum, indent, line):
        if (ea % 2 == 0) and indent == -1:
            return PREFIX
        else:
            return ""


class prefix_plugin_t(ida_idaapi.plugin_t):
    flags = 0
    comment = "This is a user defined prefix sample plugin"
    help = "This is help"
    wanted_name = "user defined prefix"
    wanted_hotkey = ""

    def __init__(self):
        self.prefix = None

    def init(self):
        self.prefix = my_user_prefix_t(8)
        print("prefix installed")
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        self.prefix = None
        print("prefix uninstalled!")


def PLUGIN_ENTRY():
    return prefix_plugin_t()


repository: BraveCattle/ctf
file: tools/ida_pro/IDA_Pro_7.7/plugins/keypatch.py
wanted_name: Keypatch Patcher
comment: Keypatch plugin for IDA Pro (using Keystone framework)
created_at: 2024-01-23 08:17:51+00:00
pushed_at: 2024-10-23 09:31:30+00:00
url:
https://github.com/BraveCattle/ctf/blob/5858d0da5c048eb8468ff8be1d63fbc50d61b601/tools/ida_pro/IDA_Pro_7.7/plugins/keypatch.py

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

# Keypatch IDA Plugin, powered by Keystone Engine (http://www.keystone-engine.org).
# By Nguyen Anh Quynh & Thanh Nguyen, 2018.

# Keypatch is released under the GPL v2. See COPYING for more information.
# Find docs & latest version at http://keystone-engine.org/keypatch

# This IDA plugin includes 3 tools inside: Patcher, Fill Range & Search.
# Access to these tools via menu "Edit | Keypatch", or via right-click popup menu "Keypatch".

# Hotkey Ctrl-Alt-K opens either Patcher or "Fill Range" window, depending on context.
#  - If there is no code selection, hotkey opens Patcher dialog
#  - If a range of code is selected, hotkey opens "Fill Range" dialog

# To revert (undo) the last patching, choose menu "Edit | Keypatch | Undo last patching".
# To check for update version, choose menu "Edit | Keypatch | Check for update".

'''
    Do not use shortcut(Ctrk + Alt + K) for debugging, otherwise debugging features will not be available.
'''
#########################################################################################################
is_debug = False
if is_debug:
    '''
        Install pydevd:

        1. sudo pip install pydevd

        or

        2. Install pycharm-debug.egg, Ensure to use pycharm pro
        https://www.jetbrains.com/help/pycharm/remote-debugging-with-product.html

        # import site
        # site.addsitedir("/usr/local/lib/python2.7/site-packages")
    '''
    try:
        import pydevd

        pydevd.settrace(host='localhost',
                        port=51234,
                        stdoutToServer=True,
                        stderrToServer=True
                        )
    except Exception as e:
        print(e)
#########################################################################################################

import os
import re
import json
from keystone import *
import idc
import idaapi
import six

################################ IDA 6/7 Compatibility import ###########################################
if idaapi.IDA_SDK_VERSION >= 700:
    import ida_search
    from idc import (
        get_operand_type,
        print_operand,
        get_item_end,
        get_item_head,
        get_item_size,
        get_sreg,
        has_value,
        get_full_flags,
        get_wide_byte,
        get_func_attr,
        set_func_end,
        warning,
        get_screen_ea
    )
    from idaapi import (
        get_bytes,
        plan_and_wait,
        is_mapped,
        Choose,
        jumpto
    )
else:
    from idc import (
        GetOpType as get_operand_type,
        GetOpnd as print_operand,
        ItemEnd as get_item_end,
        ItemHead as get_item_head,
        ItemSize as get_item_size,
        GetReg as get_sreg,
        hasValue as has_value,
        GetFlags as get_full_flags,
        Byte as get_wide_byte,
        GetFunctionAttr as get_func_attr,
        Jump as jumpto,
        Warning as warning,
        ScreenEA as get_screen_ea
    )
    from idaapi import (
        get_many_bytes as get_bytes,
        analyze_area as plan_and_wait,
        isEnabled as is_mapped,
        func_setend as set_func_end,
        Choose2 as Choose
    )
#########################################################################################################

# bleeding-edge version
# on a new release, this should be sync with VERSION_STABLE file
VERSION = "2.2"


MAX_INSTRUCTION_STRLEN = 256
MAX_ENCODING_LEN = 40
MAX_ADDRESS_LEN = 40
ENCODING_ERR_OUTPUT = "..."
KP_GITHUB_VERSION = "https://raw.githubusercontent.com/keystone-engine/keypatch/master/VERSION_STABLE"
KP_HOMEPAGE = "http://keystone-engine.org/keypatch"

X86_NOP = "\x90"

# Configuration file
KP_CFGFILE = os.path.join(idaapi.get_user_idadir(), "keypatch.cfg")

# save all the info on patching
patch_info = []

################################ Python2/3 Compatibility function #######################################
# Convert unicode/string/bytes to string
def to_string(s):
  # python3 bytes
  if six.PY3 and isinstance(s, bytes):
      return s.decode('latin-1')
  # python2 unicode
  elif six.PY2 and isinstance(s, six.text_type):
      return s.encode('utf-8')
  return str(s)

def to_hexstr(buf, sep=' '):
    # for python3 bytes
    if six.PY3 and isinstance(buf, bytes):
        return sep.join("{0:02x}".format(c) for c in buf).upper()
    return sep.join("{0:02x}".format(ord(c)) for c in buf).upper()
#########################################################################################################

################################ IDA 6/7 Compatibility function #########################################
def get_dtype(ea, op_idx):
    if idaapi.IDA_SDK_VERSION >= 700:
        insn = idaapi.insn_t()
        idaapi.decode_insn(insn, ea)
        dtype = insn.ops[op_idx].dtype
        dtyp_size = idaapi.get_dtype_size(dtype)
    else:
        dtype = idaapi.cmd.Operands[op_idx].dtyp
        dtyp_size = idaapi.get_dtyp_size(dtype)
    return dtype, dtyp_size

def set_comment(ea, comment):
    if idaapi.IDA_SDK_VERSION >= 700:
        idc.set_cmt(ea, comment, 0)
    else:
        idc.MakeComm(ea, comment)

def get_comment(ea):
    if idaapi.IDA_SDK_VERSION >= 700:
        return idc.get_cmt(ea, 0)
    return idc.Comment(ea)

def read_range_selection():
    if idaapi.IDA_SDK_VERSION >= 700:
        return idaapi.read_range_selection(None)
    return idaapi.read_selection()
#########################################################################################################

# return a normalized code, or None if input is invalid
def convert_hexstr(code):
    # normalize code
    code = code.lower()
    code = code.replace(' ', '')    # remove space
    code = code.replace('h', '')    # remove trailing 'h' in 90h
    code = code.replace('0x', '')   # remove 0x
    code = code.replace('\\x', '')  # remove \x
    code = code.replace(',', '')    # remove ,
    code = code.replace(';', '')    # remove ;
    code = code.replace('"', '')    # remove "
    code = code.replace("'", '')    # remove '
    code = code.replace("+", '')    # remove +

    # single-digit hexcode?
    if len(code) == 1 and ((code >= '0' and code <= '9') or (code >= 'a' and code <= 'f')):
        # stick 0 in front (so 'a' --> '0a')
        code = '0' + code

    # odd-length is invalid
    if len(code) % 2 != 0:
        return None

    try:
        hex_data = code.decode('hex')
        # we want a list of int
        return [ord(i) for i in hex_data]
    except:
        # invalid hex
        return None


# download a file from @url, then return (result, file-content)
# return (0, content) on success, or ({1|2}, None) on download failure
def url_download(url):
    # Import for python2
    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)


def get_name_value(_from, name):
    """
    Fixed: the return value truncated(32 bit) of get_name_value function that analyzed 64 bit binary file about ida64 for win.

    eg:
    type == idaapi.NT_BYTE
    (type, value) = idaapi.get_name_value(idc.BADADDR, "wcschr") # ida64 for win

    value = 0x14003d3f0L is correct  ida64 > 7.x for macOS
    value = 0x4003d3f0L is truncated ida64 >= 6.x for win, ida64 == 6.x for macOS

    :param _from: ea
    :param name: name string
    :return: tuple
    """
    name = to_string(name)
    (typ, value) = idaapi.get_name_value(_from, name)
    if typ == idaapi.NT_BYTE:  # typ is byte name (regular name)
        value = idaapi.get_name_ea(_from, name)
    return (typ, value)


## Main Keypatch class
class Keypatch_Asm:
    # supported architectures
    arch_lists = {
        "X86 16-bit": (KS_ARCH_X86, KS_MODE_16),                # X86 16-bit
        "X86 32-bit": (KS_ARCH_X86, KS_MODE_32),                # X86 32-bit
        "X86 64-bit": (KS_ARCH_X86, KS_MODE_64),                # X86 64-bit
        "ARM": (KS_ARCH_ARM, KS_MODE_ARM),                      # ARM
        "ARM Thumb": (KS_ARCH_ARM, KS_MODE_THUMB),              # ARM Thumb
        "ARM64 (ARMV8)": (KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN),# ARM64
        "Hexagon": (KS_ARCH_HEXAGON, KS_MODE_BIG_ENDIAN),       # Hexagon
        "Mips32": (KS_ARCH_MIPS, KS_MODE_MIPS32),               # Mips32
        "Mips64": (KS_ARCH_MIPS, KS_MODE_MIPS64),               # Mips64
        "PowerPC 32": (KS_ARCH_PPC, KS_MODE_PPC32),             # PPC32
        "PowerPC 64": (KS_ARCH_PPC, KS_MODE_PPC64),             # PPC64
        "Sparc 32": (KS_ARCH_SPARC, KS_MODE_SPARC32),           # Sparc32
        "Sparc 64": (KS_ARCH_SPARC, KS_MODE_SPARC64),           # Sparc64
        "SystemZ": (KS_ARCH_SYSTEMZ, KS_MODE_BIG_ENDIAN),       # SystemZ
    }

    endian_lists = {
        "Little Endian": KS_MODE_LITTLE_ENDIAN,                 # little endian
        "Big Endian": KS_MODE_BIG_ENDIAN,                       # big endian
    }

    syntax_lists = {
        "Intel": KS_OPT_SYNTAX_INTEL,
        "Nasm": KS_OPT_SYNTAX_NASM,
        "AT&T": KS_OPT_SYNTAX_ATT
    }

    def __init__(self, arch=None, mode=None):
        # update current arch and mode
        self.update_hardware_mode()

        # override arch & mode if provided
        if arch is not None:
            self.arch = arch
        if mode is not None:
            self.mode = mode

        # IDA uses Intel syntax by default
        self.syntax = KS_OPT_SYNTAX_INTEL

    # return Keystone arch & mode (with endianess included)
    @staticmethod
    def get_hardware_mode():
        (arch, mode) = (None, None)

        # heuristically detect hardware setup
        info = idaapi.get_inf_structure()

        try:
            cpuname = info.procname.lower()
        except:
            cpuname = info.procName.lower()

        try:
            # since IDA7 beta 3 (170724) renamed inf.mf -> is_be()/set_be()
            is_be = idaapi.cvar.inf.is_be()
        except:
            # older IDA versions
            is_be = idaapi.cvar.inf.mf
        # print("Keypatch BIG_ENDIAN = %s" %is_be)

        if cpuname == "metapc":
            arch = KS_ARCH_X86
            if info.is_64bit():
                mode = KS_MODE_64
            elif info.is_32bit():
                mode = KS_MODE_32
            else:
                mode = KS_MODE_16
        elif cpuname.startswith("arm"):
            # ARM or ARM64
            if info.is_64bit():
                arch = KS_ARCH_ARM64
                if is_be:
                    mode = KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_LITTLE_ENDIAN
            else:
                arch = KS_ARCH_ARM
                # either big-endian or little-endian
                if is_be:
                    mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("sparc"):
            arch = KS_ARCH_SPARC
            if info.is_64bit():
                mode = KS_MODE_SPARC64
            else:
                mode = KS_MODE_SPARC32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("ppc"):
            arch = KS_ARCH_PPC
            if info.is_64bit():
                mode = KS_MODE_PPC64
            else:
                mode = KS_MODE_PPC32
            if cpuname == "ppc":
                # do not support Little Endian mode for PPC
                mode += KS_MODE_BIG_ENDIAN
        elif cpuname.startswith("mips"):
            arch = KS_ARCH_MIPS
            if info.is_64bit():
                mode = KS_MODE_MIPS64
            else:
                mode = KS_MODE_MIPS32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
            arch = KS_ARCH_SYSTEMZ
            mode = KS_MODE_BIG_ENDIAN

        return (arch, mode)

    def update_hardware_mode(self):
        (self.arch, self.mode) = self.get_hardware_mode()

    # normalize assembly code
    # remove comment at the end of assembly code
    @staticmethod
    def asm_normalize(text):
        text = ' '.join(text.split())
        if text.rfind(';') != -1:
            return text[:text.rfind(';')].strip()

        return text.strip()

    @staticmethod
    # check if input address is valid
    # return
    #       -1  invalid address at target binary
    #        0  type mismatch of input address
    #        1  valid address at target binary
    def check_address(address):
        try:
            if is_mapped(address):
                return 1
            else:
                return -1
        except:
            # invalid type
            return 0

    ### resolve IDA names from input asm code
    # todo: a better syntax parser for all archs
    def ida_resolve(self, assembly, address=idc.BADADDR):
        def _resolve(_op, ignore_kw=True):
            names = re.findall(r"[\$a-z0-9_:\.]+", _op, re.I)

            # try to resolve all names
            for name in names:
                # ignore known keywords
                if ignore_kw and name in ('byte', 'near', 'short', 'word', 'dword', 'ptr', 'offset'):
                    continue

                sym = name

                # split segment reg
                parts = name.partition(':')
                if parts[2] != '':
                    sym = parts[2]

                (typ, value) = get_name_value(address, sym)

                # skip if name doesn't exist or segment / segment registers
                if typ in (idaapi.NT_SEG, idaapi.NT_NONE):
                    continue

                _op = _op.replace(sym, '0x{0:X}'.format(value))

            return _op

        if self.check_address(address) == 0:
            print("Keypatch: WARNING: invalid input address {0}".format(address))
            return assembly

        # for now, we only support IDA name resolve for X86, ARM, ARM64, MIPS, PPC, SPARC
        if not (self.arch in (KS_ARCH_X86, KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_MIPS, KS_ARCH_PPC, KS_ARCH_SPARC)):
            return assembly

        _asm = assembly.partition(' ')
        mnem = _asm[0]
        opers = _asm[2].split(',')

        for idx, op in enumerate(opers):
            _op = list(op.partition('['))
            ignore_kw = True
            if _op[1] == '':
                _op[2] = _op[0]
                _op[0] = ''
            else:
                _op[0] = _resolve(_op[0], ignore_kw=True)
                ignore_kw = False

            _op[2] = _resolve(_op[2], ignore_kw=ignore_kw)

            opers[idx] = ''.join(_op)

        asm = "{0} {1}".format(mnem, ','.join(opers))
        return asm

    # return bytes of instruction or data
    # return None on failure
    def ida_get_item(self, address, hex_output=False):
        if self.check_address(address) != 1:
            # not a valid address
            return (None, 0)

        # return None if address is in the middle of instruction / data
        if address != get_item_head(address):
            return (None, 0)

        size = get_item_size(address)
        item = get_bytes(address, size)

        if item is None:
            return (None, 0)

        if hex_output:
            item = to_hexstr(item)

        return (item, size)

    @staticmethod
    def get_op_dtype_name(ea, op_idx):
        dtyp_lists = {
            idaapi.dt_byte: 'byte',     #  8 bit
            idaapi.dt_word: 'word',     #  16 bit
            idaapi.dt_dword: 'dword',   #  32 bit
            idaapi.dt_float: 'dword',   #  4 byte
            idaapi.dt_double: 'dword',  #  8 byte
            #idaapi.dt_tbyte = 5        #  variable size (ph.tbyte_size)
            #idaapi.dt_packreal = 6         #  packed real format for mc68040
            idaapi.dt_qword: 'qword',   #  64 bit
            idaapi.dt_byte16: 'xmmword',#  128 bit
            #idaapi.dt_code = 9         #  ptr to code (not used?)
            #idaapi.dt_void = 10        #  none
            #idaapi.dt_fword = 11       #  48 bit
            #idaapi.dt_bitfild = 12     #  bit field (mc680x0)
            #idaapi.dt_string = 13      #  pointer to asciiz string
            #idaapi.dt_unicode = 14     #  pointer to unicode string
            #idaapi.dt_3byte = 15       #  3-byte data
            #idaapi.dt_ldbl = 16        #  long double (which may be different from tbyte)
            idaapi.dt_byte32: 'ymmword',# 256 bit
        }

        dtype, dtyp_size = get_dtype(ea, op_idx)
        if dtype == idaapi.dt_tbyte and dtyp_size == 10:
          return 'xword'

        dtyp_name = dtyp_lists.get(dtype, None)

        return dtyp_name

    # return asm instructions from start to end
    def ida_get_disasm_range(self, start, end):
        codes = []
        while start < end:
            asm = self.asm_normalize(idc.GetDisasm(start))
            if asm == None:
                asm = ''
            codes.append(asm)
            start = start + get_item_size(start)

        return codes

    # get disasm from IDA
    # return '' on invalid address
    def ida_get_disasm(self, address, fixup=False):

        def GetMnem(asm):
            sp = asm.find(' ')
            if sp == -1:
                return asm
            return asm[:sp]

        if self.check_address(address) != 1:
            # not a valid address
            return ''

        # return if address is in the middle of instruction / data
        if address != get_item_head(address):
            return ''

        asm = self.asm_normalize(idc.GetDisasm(address))
        # for now, only support IDA syntax fixup for Intel CPU
        if not fixup or self.arch != KS_ARCH_X86:
            return asm

        # KS_ARCH_X86 mode
        # rebuild disasm code from IDA
        i = 0
        mnem = GetMnem(asm)
        if mnem == '' or mnem in ('rep', 'repne', 'repe'):
            return asm

        opers = []
        while get_operand_type(address, i) > 0 and i < 6:
            t = get_operand_type(address, i)
            o = print_operand(address, i)

            if t in (idc.o_mem, idc.o_displ):
                parts = list(o.partition(':'))
                if parts[2] == '':
                    parts[2] = parts[0]
                    parts[0] = ''

                if '[' not in parts[2]:
                    parts[2] = '[{0}]'.format(parts[2])

                o = ''.join(parts)

                if 'ptr ' not in o:
                    dtyp_name = self.get_op_dtype_name(address, i)
                    if dtyp_name != None:
                        o = "{0} ptr {1}".format(dtyp_name, o)

            opers.append(o)
            i += 1

        asm = mnem
        for o in opers:
            if o != '':
                asm = "{0} {1},".format(asm, o)

        asm = asm.strip(',')
        return asm

    # assemble code with Keystone
    # return (encoding, count), or (None, 0) on failure
    def assemble(self, assembly, address, arch=None, mode=None, syntax=None):

        # return assembly with arithmetic equation evaluated
        def eval_operand(assembly, start, stop, prefix=''):
            imm = assembly[start+1:stop]
            try:
                eval_imm = eval(imm)
                if eval_imm > 0x80000000:
                    eval_imm = 0xffffffff - eval_imm
                    eval_imm += 1
                    eval_imm = -eval_imm
                return assembly.replace(prefix + imm, prefix + hex(eval_imm))
            except:
                return assembly

        # IDA uses different syntax from Keystone
        # sometimes, we can convert code to be consumable by Keystone
        def fix_ida_syntax(assembly):

            # return True if this insn needs to be fixed
            def check_arm_arm64_insn(arch, mnem):
                if arch == KS_ARCH_ARM:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return False
                elif arch == KS_ARCH_ARM64:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return mnem in ("stp")
                return False

            # return True if this insn needs to be fixed
            def check_ppc_insn(mnem):
                return mnem in ("stw")

            # replace the right most string occurred
            def rreplace(s, old, new):
                li = s.rsplit(old, 1)
                return new.join(li)

            # convert some ARM pre-UAL assembly to UAL, so Keystone can handle it
            # example: streqb --> strbeq
            def fix_arm_ual(mnem, assembly):
                # TODO: this is not an exhaustive list yet
                if len(mnem) != 6:
                    return assembly

                if (mnem[-1] in ('s', 'b', 'h', 'd')):
                    #print(">> 222", mnem[3:5])
                    if mnem[3:5] in ("cc", "eq", "ne", "hs", "lo", "mi", "pl", "vs", "vc", "hi", "ls", "ge", "lt", "gt", "le", "al"):
                        return assembly.replace(mnem, mnem[:3] + mnem[-1] + mnem[3:5], 1)

                return assembly

            if self.arch != KS_ARCH_X86:
                assembly = assembly.lower()
            else:
                # Keystone does not support immediate 0bh, but only 0Bh
                assembly = assembly.upper()

            # however, 0X must be converted to 0x
            # Keystone should fix this limitation in the future
            assembly = assembly.replace("0X", " 0x")

            _asm = assembly.partition(' ')
            mnem = _asm[0]
            if mnem == '':
                return assembly

            # for PPC, Keystone does not accept registers with 'r' prefix,
            # but only the number behind. lets try to fix that here by
            # removing the prefix 'r'.
            if self.arch == KS_ARCH_PPC:
                for n in range(32):
                    r = " r%u," %n
                    if r in assembly:
                        assembly = assembly.replace(r, " %u," %n)
                for n in range(32):
                    r = "(r%u)" %n
                    if r in assembly:
                        assembly = assembly.replace(r, "(%u)" %n)
                for n in range(32):
                    r = ", r%u" %n
                    if assembly.endswith(r):
                        assembly = rreplace(assembly, r, ", %u" %n)

            if self.arch == KS_ARCH_X86:
                if mnem == "RETN":
                    # replace retn with ret
                    return assembly.replace('RETN', 'RET', 1)
                if 'OFFSET ' in assembly:
                    return assembly.replace('OFFSET ', ' ')
                if mnem in ('CALL', 'JMP') or mnem.startswith('LOOP'):
                    # remove 'NEAR PTR'
                    if ' NEAR PTR ' in assembly:
                        return assembly.replace(' NEAR PTR ', ' ')
                elif mnem[0] == 'J':
                    # JMP instruction
                    if ' SHORT ' in assembly:
                        # remove ' short '
                        return assembly.replace(' SHORT ', ' ')
            elif self.arch in (KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_PPC):
                # *** ARM
                # LDR     R1, [SP+rtld_fini],#4
                # STR     R2, [SP,#-4+rtld_fini]!
                # STR     R0, [SP,#fini]!
                # STR     R12, [SP,#4+var_8]!

                # *** ARM64
                # STP     X29, X30, [SP,#-0x10+var_150]!
                # STR     W0, [X29,#0x150+var_8]
                # LDR     X0, [X0,#(qword_4D6678 - 0x4D6660)]
                # TODO:
                # ADRP    X19, #interactive@PAGE

                # *** PPC
                # stw     r5, 0x120+var_108(r1)

                if self.arch == KS_ARCH_ARM and mode == KS_MODE_THUMB:
                    assembly = assembly.replace('movt.w', 'movt')

                if self.arch == KS_ARCH_ARM:
                    #print(">> before UAL fix: ", assembly)
                    assembly = fix_arm_ual(mnem, assembly)
                    #print(">> after UAL fix: ", assembly)

                if check_arm_arm64_insn(self.arch, mnem) or (("[" in assembly) and ("]" in assembly)):
                    bang = assembly.find('#')
                    bracket = assembly.find(']')
                    if bang != -1 and bracket != -1 and bang < bracket:
                        return eval_operand(assembly, bang, bracket, '#')
                    elif '+0x0]' in assembly:
                        return assembly.replace('+0x0]', ']')
                elif check_ppc_insn(mnem):
                    start = assembly.find(', ')
                    stop = assembly.find('(')
                    if start != -1 and stop != -1 and start < stop:
                        return eval_operand(assembly, start, stop)
            return assembly

        def is_thumb(address):
            return get_sreg(address, 'T') == 1

        if self.check_address(address) == 0:
            return (None, 0)

        # use default syntax, arch and mode if not provided
        if syntax is None:
            syntax = self.syntax
        if arch is None:
            arch = self.arch
        if mode is None:
            mode = self.mode

        if arch == KS_ARCH_ARM and is_thumb(address):
            mode = KS_MODE_THUMB

        try:
            ks = Ks(arch, mode)
            if arch == KS_ARCH_X86:
                ks.syntax = syntax
            encoding, count = ks.asm(fix_ida_syntax(assembly), address)
        except KsError as e:
            # keep the below code for debugging
            #print("Keypatch Error: {0}".format(e))
            #print("Original asm: {0}".format(assembly))
            #print("Fixed up asm: {0}".format(fix_ida_syntax(assembly)))
            encoding, count = None, 0

        return (encoding, count)


    # patch at address, return the number of written bytes & original data
    # this process can fail in some cases
    @staticmethod
    def patch_raw(address, patch_data, size):
        ea = address
        orig_data = ''

        while ea < (address + size):

            if not has_value(get_full_flags(ea)):
                print("Keypatch: FAILED to read data at 0x{0:X}".format(ea))
                break

            orig_byte = get_wide_byte(ea)
            orig_data += chr(orig_byte)
            patch_byte = ord(patch_data[ea - address])

            if patch_byte != orig_byte:
                # patch one byte
                if idaapi.patch_byte(ea, patch_byte) != 1:
                    print("Keypatch: FAILED to patch byte at 0x{0:X} [0x{1:X}]".format(ea, patch_byte))
                    break
            ea += 1
        return (ea - address, orig_data)

    # patch at address, return the number of written bytes & original data
    # on patch failure, we revert to the original code, then return (None, None)
    def patch(self, address, patch_data, size):
        # save original function end to fix IDA re-analyze issue after patching
        orig_func_end = get_func_attr(address, idc.FUNCATTR_END)

        (patched_len, orig_data) = self.patch_raw(address, patch_data, size)

        if size != patched_len:
            # patch failure
            if patched_len > 0:
                # revert the changes
                (rlen, _) = self.patch_raw(address, orig_data, patched_len)
                if rlen == patched_len:
                    print("Keypatch: successfully reverted changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))
                else:
                    print("Keypatch: FAILED to revert changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))

            return (None, None)

        # ask IDA to re-analyze the patched area
        if orig_func_end == idc.BADADDR:
            # only analyze patched bytes, otherwise it would take a lot of time to re-analyze the whole binary
            plan_and_wait(address, address + patched_len + 1)
        else:
            plan_and_wait(address, orig_func_end)

            # try to fix IDA function re-analyze issue after patching
            set_func_end(address, orig_func_end)

        return (patched_len, orig_data)

    # return number of bytes patched
    # return
    #    0  Invalid assembly
    #   -1  PatchByte failure
    #   -2  Can't read original data
    #   -3  Invalid address
    def patch_code(self, address, assembly, syntax, padding, save_origcode, orig_asm=None, patch_data=None, patch_comment=None, undo=False):
        global patch_info

        if self.check_address(address) != 1:
            # not a valid address
            return -3

        orig_comment = get_comment(address)
        if orig_comment is None:
            orig_comment = ''

        nop_comment = ""
        padding_len = 0
        if not undo:
            # we are patching via Patcher
            (orig_encoding, orig_len) = self.ida_get_item(address)
            if (orig_encoding, orig_len) == (None, 0):
                return -2

            (encoding, count) = self.assemble(assembly, address, syntax=syntax)
            if encoding is None:
                return 0

            patch_len = len(encoding)
            patch_data = ''.join(chr(c) for c in encoding)

            if patch_data == orig_encoding:
                #print("Keypatch: no need to patch, same encoding data [{0}] at 0x{1:X}".format(to_hexstr(orig_encoding), address))
                return orig_len

            # for now, only support NOP padding on Intel CPU
            if padding and self.arch == KS_ARCH_X86:
                if patch_len < orig_len:
                    padding_len = orig_len - patch_len
                    patch_len = orig_len
                    patch_data = patch_data.ljust(patch_len, X86_NOP)
                elif patch_len > orig_len:
                    patch_end = address + patch_len - 1
                    ins_end = get_item_end(patch_end)
                    padding_len = ins_end - patch_end - 1

                    if padding_len > 0:
                        patch_len = ins_end - address
                        patch_data = patch_data.ljust(patch_len, X86_NOP)

                if padding_len > 0:
                    nop_comment = "\nKeypatch padded NOP to next boundary: {0} bytes".format(padding_len)

            orig_asm = self.ida_get_disasm_range(address, address + patch_len)
        else:
            # we are reverting the change via "Undo" menu
            patch_len = len(patch_data)

        (plen, p_orig_data) = self.patch(address, patch_data, patch_len)
        if plen is None:
            # failed to patch
            return -1

        if not undo: # we are patching
            new_patch_comment = None
            if save_origcode is True:
                # append original instruction to comments
                if orig_comment == '':
                    new_patch_comment = "Keypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)
                else:
                    new_patch_comment = "\nKeypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)

                new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
                set_comment(address, new_comment)

            if padding_len == 0:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))
            else:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}], with {4} byte(s) NOP padded".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data), padding_len))
            # save this patching for future "undo"
            patch_info.append((address, assembly, p_orig_data, new_patch_comment))
        else:   # we are reverting
            if patch_comment:
                # clean previous IDA comment by replacing it with ''
                new_comment = orig_comment.replace(patch_comment, '')
                set_comment(address, new_comment)

            print("Keypatch: successfully reverted {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))

        return plen

    # fill a range of code [addr_begin, addr_end].
    # return the length of patched area
    # on failure, return 0 = wrong input, -1 = failed to patch
    def fill_code(self, addr_begin, addr_end, assembly, syntax, padding, save_origcode, orig_asm=None):
        # treat input as assembly code first
        (encoding, _) =  self.assemble(assembly, addr_begin, syntax=syntax)

        if encoding is None:
            # input might be a hexcode string. try to convert it to raw bytes
            encoding = convert_hexstr(assembly)

        if encoding is None:
            # invalid input: this is neither assembly nor hexcode string
            return 0

        # save original assembly code before overwritting them
        orig_asm = self.ida_get_disasm_range(addr_begin, addr_end)

        # save original comment at addr_begin
        # TODO: save comments in this range, but how to interleave them?
        orig_comment = get_comment(addr_begin)
        if orig_comment is None:
            orig_comment = ''

        patch_data = ""
        assembly_new = []
        size = addr_end - addr_begin
        # calculate filling data
        encode_chr = ''.join(chr(c) for c in encoding)
        while True:
            if len(patch_data) + len(encode_chr) <= size:
                patch_data = patch_data + encode_chr
                assembly_new += [assembly.strip()]
            else:
                break

        # for now, only support NOP padding on Intel CPU
        if padding and self.arch == KS_ARCH_X86:
            for i in range(size -len(patch_data)):
                assembly_new += ["nop"]
            patch_data = patch_data.ljust(size, X86_NOP)

        (plen, p_orig_data) = self.patch(addr_begin, patch_data, len(patch_data))
        if plen is None:
            # failed to patch
            return -1

        new_patch_comment = ''
        # append original instruction to comments
        if save_origcode is True:
            if orig_comment == '':
                new_patch_comment = "Keypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))
            else:
                new_patch_comment = "\nKeypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))

            new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
            set_comment(addr_begin, new_comment)

        print("Keypatch: successfully filled range [0x{0:X}:0x{1:X}] ({2} bytes) with \"{3}\", replaced \"{4}\"".format(
                    addr_begin, addr_end - 1, addr_end - addr_begin, assembly, '; '.join(orig_asm)))

        # save this modification for future "undo"
        patch_info.append((addr_begin, '\n  '.join(assembly_new), p_orig_data, new_patch_comment))

        return plen


    ### Form helper functions
    @staticmethod
    def dict_to_ordered_list(dictionary):
        l = sorted(list(dictionary.items()), key=lambda t: t[0], reverse=False)
        keys = [i[0] for i in l]
        values = [i[1] for i in l]

        return (keys, values)

    def get_value_by_idx(self, dictionary, idx, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            val = values[idx]
        except IndexError:
            val = default

        return val

    def find_idx_by_value(self, dictionary, value, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            idx = values.index(value)
        except:
            idx = default

        return idx

    def get_arch_by_idx(self, idx):
        return self.get_value_by_idx(self.arch_lists, idx)

    def find_arch_idx(self, arch, mode):
        return self.find_idx_by_value(self.arch_lists, (arch, mode))

    def get_syntax_by_idx(self, idx):
        return self.get_value_by_idx(self.syntax_lists, idx, self.syntax)

    def find_syntax_idx(self, syntax):
        return self.find_idx_by_value(self.syntax_lists, syntax)
    ### /Form helper functions


# Common ancestor form to be derived by Patcher, FillRange & Search
class Keypatch_Form(idaapi.Form):
    # prepare for form initializing
    def setup(self, kp_asm, address, assembly=None):
        self.kp_asm = kp_asm
        self.address = address

        # update ordered list of arch and syntax
        self.syntax_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.syntax_lists)[0]
        self.arch_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.arch_lists)[0]

        # update current arch & mode
        self.kp_asm.update_hardware_mode()

        # find right value for c_arch & c_endian controls
        mode = self.kp_asm.mode
        self.endian_id = 0   # little endian
        if self.kp_asm.mode & KS_MODE_BIG_ENDIAN:
            self.endian_id = 1   # big endian
            mode = self.kp_asm.mode - KS_MODE_BIG_ENDIAN

        self.arch_id = self.kp_asm.find_arch_idx(self.kp_asm.arch, mode)

        self.syntax_id = 0  # to make non-X86 arch happy
        if self.kp_asm.arch == KS_ARCH_X86:
            self.syntax_id = self.kp_asm.find_syntax_idx(self.kp_asm.syntax)

        # get original instruction and bytes
        self.orig_asm = kp_asm.ida_get_disasm(address)
        (self.orig_encoding, self.orig_len) = kp_asm.ida_get_item(address, hex_output=True)
        if self.orig_encoding is None:
            self.orig_encoding = ''

        if assembly is None:
            self.asm = self.kp_asm.ida_get_disasm(self.address, fixup=True)
        else:
            self.asm = assembly

    # update Encoding control
    # return True on success, False on failure
    def _update_encoding(self, arch, mode):
        try:
            syntax = None
            if arch == KS_ARCH_X86:
                syntax_id = self.GetControlValue(self.c_syntax)
                syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

            address = self.GetControlValue(self.c_addr)
            try:
                is_mapped(address)
            except:
                # invalid address value
                address = 0

            assembly = self.GetControlValue(self.c_assembly)
            raw_assembly = self.kp_asm.ida_resolve(assembly, address)
            self.SetControlValue(self.c_raw_assembly, raw_assembly)

            (encoding, count) =  self.kp_asm.assemble(raw_assembly, address, arch=arch,
                                                    mode=mode, syntax=syntax)

            if encoding is None:
                self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                self.SetControlValue(self.c_encoding_len, 0)
                return False
            else:
                text = ""
                for byte in encoding:
                    text += "%02X " % byte
                text.strip()
                if text == "":
                    # error?
                    self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                    return False
                else:
                    self.SetControlValue(self.c_encoding, text.strip())
                    self.SetControlValue(self.c_encoding_len, len(encoding))
                    return True
        except Exception as e:
            print(str(e))
            import traceback
            traceback.print_exc()
            self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
            return False

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        return 1

    # update Patcher & Fillrange controls
    def update_patchform(self, fid):
        self.EnableField(self.c_endian, False)
        self.EnableField(self.c_addr, False)

        (arch, mode) = (self.kp_asm.arch, self.kp_asm.mode)
        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # do not show Endian control
            self.ShowField(self.c_endian, False)
            # allow to choose Syntax
            self.ShowField(self.c_syntax, True)
            self.ShowField(self.c_opt_padding, True)
        else:   # do not show Syntax control for non-X86 mode
            self.ShowField(self.c_syntax, False)
            # for now, we do not support padding for non-X86 archs
            self.ShowField(self.c_opt_padding, False)
            #self.EnableField(self.c_opt_padding, False)

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1

    # update some controls - including Encoding control
    def update_controls(self, arch, mode):
        # Fixup & Encoding-len are read-only controls
        self.EnableField(self.c_raw_assembly, False)
        self.EnableField(self.c_encoding_len, False)

        # Encoding is enable to allow user to select & copy
        self.EnableField(self.c_encoding, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update encoding with live assembling
        self._update_encoding(arch, mode | endian)

        return 1

    # get Patcher/FillRange options
    def get_opts(self, name=None):
        names = self.c_opt_chk.children_names
        val = self.c_opt_chk.value
        opts = {}
        for i in range(len(names)):
            opts[names[i]] = val & (2**i)

        if name != None:
            opts[name] = val

        return opts


# Fill Range form
class Keypatch_FillRange(Keypatch_Form):
    def __init__(self, kp_asm, addr_begin, addr_end, assembly=None, opts=None):
        self.setup(kp_asm, addr_begin, assembly)
        self.addr_end = addr_end

        # create FillRange form
        super(Keypatch_FillRange, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Fill Range

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Start      :{c_addr}>
            <End        :{c_addr_end}>
            <Size       :{c_size}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=addr_begin, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_addr_end': self.NumericInput(value=addr_end - 1, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_size': self.NumericInput(value=addr_end - addr_begin, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':idaapi.Form.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some controls read-only in FillRange mode
        self.EnableField(self.c_size, False)
        self.EnableField(self.c_addr_end, False)

        return self.update_patchform(fid)


# Patcher form
class Keypatch_Patcher(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None, opts=None):
        self.setup(kp_asm, address, assembly)

        # create Patcher form
        super(Keypatch_Patcher, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Patcher

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Address    :{c_addr}>
            <Original   :{c_orig_assembly}>
             <##-   Encode:{c_orig_encoding}>
             <##-   Size  :{c_orig_len}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_assembly': self.StringInput(value=self.orig_asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_encoding': self.StringInput(value=self.orig_encoding[:MAX_ENCODING_LEN], width=MAX_ENCODING_LEN),
            'c_orig_len': self.NumericInput(value=self.orig_len, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':self.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some fields read-only in Patch mode
        self.EnableField(self.c_orig_assembly, False)
        self.EnableField(self.c_orig_encoding, False)
        self.EnableField(self.c_orig_len, False)

        return self.update_patchform(fid)


# Search position chooser
class SearchResultChooser(Choose):
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False):
        super(SearchResultChooser, self).__init__(
            title,
            [["Address", Choose.CHCOL_HEX|40]],
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = items
        self.selcount = 0
        self.modal = modal

    def OnClose(self):
        return

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

    def OnGetLine(self, n):
        res = self.items[n]
        res = [idc.atoa(res[0])]
        return res

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

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


# Search form
class Keypatch_Search(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None):
        self.setup(kp_asm, address, assembly)

        # create Search form
        super(Keypatch_Search, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Search
KEYPATCH:: Search

            {FormChangeCb}
            <A~r~ch       :{c_arch}>
            <E~n~dian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <A~d~dress    :{c_addr}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            """, {
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_arch': self.DropdownListControl(
                          items = self.arch_keys,
                          readonly = True,
                          selval = self.arch_id,
                          width = 32),
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    ################################ IDA 6/7 Compatibility method #######################################
    def find_binary(self, ea):
        if idaapi.IDA_SDK_VERSION >= 700:
            return ida_search.find_binary(ea, idaapi.cvar.inf.max_ea, self.GetControlValue(self.c_encoding), 16, idc.SEARCH_DOWN)
        return idc.FindBinary(ea, idc.SEARCH_DOWN, self.GetControlValue(self.c_encoding))
    #####################################################################################################

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # handle the search button
        if fid == -2:
            address = 0
            addresses = []
            while address != idc.BADADDR:
                address = self.find_binary(address)
                if address == idc.BADADDR:
                    break
                addresses.append([address])
                address = address + 1
            c = SearchResultChooser("Searching for [{0}]".format(self.GetControlValue(self.c_raw_assembly)), addresses)
            c.show()
            return 1

        # only Search mode allows to select arch+mode
        arch_id = self.GetControlValue(self.c_arch)
        (arch, mode) = self.kp_asm.get_arch_by_idx(arch_id)

        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # enable Syntax and disable Endian for x86
            self.ShowField(self.c_syntax, True)
            self.EnableField(self.c_syntax, True)
            self.syntax_id = self.GetControlValue(self.c_syntax)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, 0)
        elif arch in (KS_ARCH_ARM64, KS_ARCH_HEXAGON, KS_ARCH_SYSTEMZ):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        elif (arch == KS_ARCH_PPC) and (mode & KS_MODE_PPC32 != 0):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        else:
            # no Syntax & Endian option
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1


# About form
class About_Form(idaapi.Form):
    def __init__(self, version):
        # create About form
        super(About_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: About

            {FormChangeCb}
            Keypatch IDA plugin v%s, using Keystone Engine v%s.
            (c) Nguyen Anh Quynh + Thanh Nguyen, 2018.

            Keypatch is released under the GPL v2.
            Find more info at http://www.keystone-engine.org/keypatch
            """ %(version, keystone.__version__), {
            '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(KP_HOMEPAGE, new = 2)

        return 1


# Check-for-update form
class Update_Form(idaapi.Form):
    def __init__(self, version, message):
        # create Update form
        super(Update_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: Check for update

            {FormChangeCb}
            Your Keypatch 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(KP_HOMEPAGE, new = 2)

        return 1


try:
    # adapted from pull request #7 by @quangnh89
    class Kp_Menu_Context(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.widget_type == idaapi.BWN_DISASM:
                    return idaapi.AST_ENABLE_FOR_FORM
                else:
                    return idaapi.AST_DISABLE_FOR_FORM
            except Exception as e:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS

    # context menu for Patcher
    class Kp_MC_Patcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.patcher()
            return 1

    # context menu for Fill Range
    class Kp_MC_Fill_Range(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.fill_range()
            return 1

    # context menu for Undo
    class Kp_MC_Undo(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.undo()
            return 1

    # context menu for Search
    class Kp_MC_Search(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1

    # context menu for Check Update
    class Kp_MC_Updater(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.updater()
            return 1

    # context menu for About
    class Kp_MC_About(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.about()
            return 1

except:
    pass

# hooks for popup menu
class Hooks(idaapi.UI_Hooks):
    if idaapi.IDA_SDK_VERSION >= 700:
        # IDA >= 700 right click widget popup
        def finish_populating_widget_popup(self, form, popup):
            if idaapi.get_widget_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass
    else:
        # IDA < 700 right click popup
        def finish_populating_tform_popup(self, form, popup):
            # We'll add our action to all "IDA View-*"s.
            # If we wanted to add it only to "IDA View-A", we could
            # also discriminate on the widget's title:
            #
            #  if idaapi.get_tform_title(form) == "IDA View-A":
            #      ...
            #
            if idaapi.get_tform_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass


# check if we already initialized Keypatch
kp_initialized = False

#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class Keypatch_Plugin_t(idaapi.plugin_t):
    comment = "Keypatch plugin for IDA Pro (using Keystone framework)"
    help = "Find more information on Keypatch at http://keystone-engine.org/keypatch"
    wanted_name = "Keypatch Patcher"
    wanted_hotkey = "Ctrl-Alt-K"
    flags = idaapi.PLUGIN_KEEP

    def load_configuration(self):
        # default
        self.opts = {}

        # load configuration from file
        try:
            f = open(KP_CFGFILE, "rt")
            self.opts = json.load(f)
            f.close()
        except IOError:
            print("Keypatch: FAILED to load config file. Use default setup now.")
        except Exception as e:
            print("Keypatch: FAILED to load config file, with exception: {0}".format(str(e)))

        # use default values if not defined in config file
        if 'c_opt_padding' not in self.opts:
            self.opts['c_opt_padding'] = 1

        if 'c_opt_comment' not in self.opts:
            self.opts['c_opt_comment'] = 2

        self.opts['c_opt_chk'] = self.opts['c_opt_padding'] | self.opts['c_opt_comment']

    def init(self):
        global kp_initialized

        # register popup menu handlers
        try:
            Kp_MC_Patcher.register(self, "Patcher    (Ctrl-Alt-K)")
            Kp_MC_Fill_Range.register(self, "Fill Range")
            Kp_MC_Undo.register(self, "Undo last patching")
            Kp_MC_Search.register(self, "Search")
            Kp_MC_Updater.register(self, "Check for update")
            Kp_MC_About.register(self, "About")
        except:
            pass

        # setup popup menu
        self.hooks = Hooks()
        self.hooks.hook()

        self.opts = None
        if kp_initialized == False:
            kp_initialized = True

            if idaapi.IDA_SDK_VERSION >= 700:
                # Add menu IDA >= 7.0
                idaapi.attach_action_to_menu("Edit/Keypatch/Patcher", Kp_MC_Patcher.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/About", Kp_MC_About.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Check for update", Kp_MC_Updater.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Search", Kp_MC_Search.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Undo last patching", Kp_MC_Undo.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Fill Range", Kp_MC_Fill_Range.get_name(), idaapi.SETMENU_APP)
            else:
                # add Keypatch menu
                menu = idaapi.add_menu_item("Edit/Keypatch/", "Patcher     (Ctrl-Alt-K)", "", 1, self.patcher, None)
                if menu is not None:
                    idaapi.add_menu_item("Edit/Keypatch/", "About", "", 1, self.about, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Check for update", "", 1, self.updater, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Search", "", 1, self.search, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Undo last patching", "", 1, self.undo, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Fill Range", "", 1, self.fill_range, None)
                elif idaapi.IDA_SDK_VERSION < 680:
                    # older IDAPython (such as in IDAPro 6.6) does add new submenu.
                    # in this case, put Keypatch menu in menu Edit \ Patch program
                    # not sure about v6.7, so to be safe we just check against v6.8
                    idaapi.add_menu_item("Edit/Patch program/", "-", "", 0, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: About", "", 0, self.about, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Check for update", "", 0, self.updater, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Search", "", 0, self.search, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Undo last patching", "", 0, self.undo, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Fill Range", "", 0, self.fill_range, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Patcher     (Ctrl-Alt-K)", "", 0, self.patcher, None)

            print("=" * 80)
            print("Keypatch v{0} (c) Nguyen Anh Quynh & Thanh Nguyen, 2018".format(VERSION))
            print("Keypatch is using Keystone v{0}".format(keystone.__version__))
            print("Keypatch Patcher's shortcut key is Ctrl-Alt-K")
            print("Use the same hotkey Ctrl-Alt-K to open 'Fill Range' window on a selected range of code")
            print("To revert (undo) the last patching, choose menu Edit | Keypatch | Undo last patching")
            print("Keypatch Search is available from menu Edit | Keypatch | Search")
            print("Find more information about Keypatch at http://keystone-engine.org/keypatch")

            self.load_configuration()

            print("=" * 80)
            self.kp_asm = Keypatch_Asm()

        return idaapi.PLUGIN_KEEP

    def term(self):
        if self.hooks is not None:
            self.hooks.unhook()
            self.hooks = None

        if self.opts is None:
            return
        #save configuration to file
        try:
            json.dump(self.opts, open(KP_CFGFILE, "wt"))
        except Exception as e:
            print("Keypatch: FAILED to save config file, with exception: {0}".format(str(e)))
        else:
            print("Keypatch: configuration is saved to {0}".format(KP_CFGFILE))

    # null handler
    def menu_null(self):
        pass

    # handler for About menu
    def about(self):
        f = About_Form(VERSION)
        f.Execute()
        f.Free()

    # handler for Check-for-Update menu
    def updater(self):
        (r, content) = url_download(KP_GITHUB_VERSION)
        content = to_string(content)
        if r == 0:
            # find stable version
            sig = 'VERSION_STABLE = "'
            tmp = content[content.find(sig)+len(sig):]
            version_stable = tmp[:tmp.find('"')]

            # compare with the current version
            if version_stable == VERSION:
                f = Update_Form(VERSION, "Good, you are already on the latest stable version!")
                f.Execute()
                # free this form
                f.Free()
            else:
                f = Update_Form(VERSION, "Download latest stable version {0} from http://keystone-engine.org/keypatch".format(version_stable))
                f.Execute()
                # free this form
                f.Free()
        else:
            # fail to download
            warning("ERROR: Keypatch failed to connect to internet (Github). Try again later.")
            print("Keypatch: FAILED to connect to Github to check for latest update. Try again later.")

    # handler for Undo menu
    def undo(self):
        global patch_info
        if len(patch_info) == 0:
            # TODO: disable Undo menu?
            warning("ERROR: Keypatch already got to the last undo patching!")
        else:
            (address, assembly, p_orig_data, patch_comment) = patch_info[-1]

            # undo the patch
            self.kp_asm.patch_code(address, None, None, None, None, orig_asm=[assembly], patch_data=p_orig_data, patch_comment=patch_comment, undo=True)
            del(patch_info[-1])

    # handler for Search menu
    def search(self):
        address = get_screen_ea()
        f = Keypatch_Search(self.kp_asm, address)
        f.Execute()
        f.Free()

    # handler for Patcher menu
    def patcher(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return

        selection, addr_begin, addr_end = read_range_selection()
        if selection:
            # call Fill Range function on this selected code
            return self.fill_range()

        address = get_screen_ea()

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        while True:
            f = Keypatch_Patcher(self.kp_asm, address, assembly=init_assembly, opts=self.opts)
            ok = f.Execute()
            if ok == 1:
                try:
                    syntax = None
                    if f.kp_asm.arch == KS_ARCH_X86:
                        syntax_id = f.c_syntax.value
                        syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                    assembly = f.c_assembly.value
                    self.opts = f.get_opts('c_opt_chk')
                    padding = (self.opts.get("c_opt_padding", 0) != 0)
                    comment = (self.opts.get("c_opt_comment", 0) != 0)

                    raw_assembly = self.kp_asm.ida_resolve(assembly, address)

                    print("Keypatch: attempting to modify \"{0}\" at 0x{1:X} to \"{2}\"".format(
                            self.kp_asm.ida_get_disasm(address), address, assembly))

                    length = self.kp_asm.patch_code(address, raw_assembly, syntax, padding, comment, None)
                    if length > 0:
                        # update start address pointing to the next instruction
                        init_assembly = None
                        address += length
                    else:
                        init_assembly = f.c_assembly.value
                        if length == 0:
                            warning("ERROR: Keypatch found invalid assembly [{0}]".format(assembly))
                        elif length == -1:
                            warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(address))
                        elif length == -2:
                            warning("ERROR: Keypatch can't read original data at 0x{0:X}, try again".format(address))

                except KsError as e:
                    print("Keypatch Error: {0}".format(e))
            else:   # Cancel
                break
            f.Free()

    # handler for Fill Range menu
    def fill_range(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return

        selection, addr_begin, addr_end = read_range_selection()
        if not selection:
            warning("ERROR: Keypatch requires a range to be selected for fill in, try again")
            return

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        f = Keypatch_FillRange(self.kp_asm, addr_begin, addr_end, assembly=init_assembly, opts=self.opts)
        ok = f.Execute()
        if ok == 1:
            try:
                syntax = None
                if f.kp_asm.arch == KS_ARCH_X86:
                    syntax_id = f.c_syntax.value
                    syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                assembly = f.c_assembly.value
                self.opts = f.get_opts('c_opt_chk')
                padding = (self.opts.get("c_opt_padding", 0) != 0)
                comment = (self.opts.get("c_opt_comment", 0) != 0)

                raw_assembly = self.kp_asm.ida_resolve(assembly, addr_begin)

                print("Keypatch: attempting to fill range [0x{0:X}:0x{1:X}] with \"{2}\"".format(
                    addr_begin, addr_end - 1, assembly))

                length = self.kp_asm.fill_code(addr_begin, addr_end, raw_assembly, syntax, padding, comment, None)
                if length == 0:
                    warning("ERROR: Keypatch failed to process this input.")
                    print("Keypatch: FAILED to process this input '{0}'".format(assembly))
                elif length == -1:
                    warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(addr_begin))

            except KsError as e:
                print("Keypatch Error: {0}".format(e))

        # free this form
        f.Free()

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


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


repository: carsond135/malwareanalysis
file: tools/IDA_Pro_7.7/IDA_Pro_7.7/plugins/keypatch.py
description: NKU恶意代码分析
wanted_name: Keypatch Patcher
comment: Keypatch plugin for IDA Pro (using Keystone framework)
created_at: 2023-06-25 19:16:29+00:00
pushed_at: 2022-12-18 09:34:14+00:00
url:
https://github.com/carsond135/malwareanalysis/blob/45c8bf862a54239ff1e7c91313eccf9cf7d4a35b/tools/IDA_Pro_7.7/IDA_Pro_7.7/plugins/keypatch.py

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

# Keypatch IDA Plugin, powered by Keystone Engine (http://www.keystone-engine.org).
# By Nguyen Anh Quynh & Thanh Nguyen, 2018.

# Keypatch is released under the GPL v2. See COPYING for more information.
# Find docs & latest version at http://keystone-engine.org/keypatch

# This IDA plugin includes 3 tools inside: Patcher, Fill Range & Search.
# Access to these tools via menu "Edit | Keypatch", or via right-click popup menu "Keypatch".

# Hotkey Ctrl-Alt-K opens either Patcher or "Fill Range" window, depending on context.
#  - If there is no code selection, hotkey opens Patcher dialog
#  - If a range of code is selected, hotkey opens "Fill Range" dialog

# To revert (undo) the last patching, choose menu "Edit | Keypatch | Undo last patching".
# To check for update version, choose menu "Edit | Keypatch | Check for update".

'''
    Do not use shortcut(Ctrk + Alt + K) for debugging, otherwise debugging features will not be available.
'''
#########################################################################################################
is_debug = False
if is_debug:
    '''
        Install pydevd:

        1. sudo pip install pydevd

        or

        2. Install pycharm-debug.egg, Ensure to use pycharm pro
        https://www.jetbrains.com/help/pycharm/remote-debugging-with-product.html

        # import site
        # site.addsitedir("/usr/local/lib/python2.7/site-packages")
    '''
    try:
        import pydevd

        pydevd.settrace(host='localhost',
                        port=51234,
                        stdoutToServer=True,
                        stderrToServer=True
                        )
    except Exception as e:
        print(e)
#########################################################################################################

import os
import re
import json
from keystone import *
import idc
import idaapi
import six

################################ IDA 6/7 Compatibility import ###########################################
if idaapi.IDA_SDK_VERSION >= 700:
    import ida_search
    from idc import (
        get_operand_type,
        print_operand,
        get_item_end,
        get_item_head,
        get_item_size,
        get_sreg,
        has_value,
        get_full_flags,
        get_wide_byte,
        get_func_attr,
        set_func_end,
        warning,
        get_screen_ea
    )
    from idaapi import (
        get_bytes,
        plan_and_wait,
        is_mapped,
        Choose,
        jumpto
    )
else:
    from idc import (
        GetOpType as get_operand_type,
        GetOpnd as print_operand,
        ItemEnd as get_item_end,
        ItemHead as get_item_head,
        ItemSize as get_item_size,
        GetReg as get_sreg,
        hasValue as has_value,
        GetFlags as get_full_flags,
        Byte as get_wide_byte,
        GetFunctionAttr as get_func_attr,
        Jump as jumpto,
        Warning as warning,
        ScreenEA as get_screen_ea
    )
    from idaapi import (
        get_many_bytes as get_bytes,
        analyze_area as plan_and_wait,
        isEnabled as is_mapped,
        func_setend as set_func_end,
        Choose2 as Choose
    )
#########################################################################################################

# bleeding-edge version
# on a new release, this should be sync with VERSION_STABLE file
VERSION = "2.2"


MAX_INSTRUCTION_STRLEN = 256
MAX_ENCODING_LEN = 40
MAX_ADDRESS_LEN = 40
ENCODING_ERR_OUTPUT = "..."
KP_GITHUB_VERSION = "https://raw.githubusercontent.com/keystone-engine/keypatch/master/VERSION_STABLE"
KP_HOMEPAGE = "http://keystone-engine.org/keypatch"

X86_NOP = "\x90"

# Configuration file
KP_CFGFILE = os.path.join(idaapi.get_user_idadir(), "keypatch.cfg")

# save all the info on patching
patch_info = []

################################ Python2/3 Compatibility function #######################################
# Convert unicode/string/bytes to string
def to_string(s):
  # python3 bytes
  if six.PY3 and isinstance(s, bytes):
      return s.decode('latin-1')
  # python2 unicode
  elif six.PY2 and isinstance(s, six.text_type):
      return s.encode('utf-8')
  return str(s)

def to_hexstr(buf, sep=' '):
    # for python3 bytes
    if six.PY3 and isinstance(buf, bytes):
        return sep.join("{0:02x}".format(c) for c in buf).upper()
    return sep.join("{0:02x}".format(ord(c)) for c in buf).upper()
#########################################################################################################

################################ IDA 6/7 Compatibility function #########################################
def get_dtype(ea, op_idx):
    if idaapi.IDA_SDK_VERSION >= 700:
        insn = idaapi.insn_t()
        idaapi.decode_insn(insn, ea)
        dtype = insn.ops[op_idx].dtype
        dtyp_size = idaapi.get_dtype_size(dtype)
    else:
        dtype = idaapi.cmd.Operands[op_idx].dtyp
        dtyp_size = idaapi.get_dtyp_size(dtype)
    return dtype, dtyp_size

def set_comment(ea, comment):
    if idaapi.IDA_SDK_VERSION >= 700:
        idc.set_cmt(ea, comment, 0)
    else:
        idc.MakeComm(ea, comment)

def get_comment(ea):
    if idaapi.IDA_SDK_VERSION >= 700:
        return idc.get_cmt(ea, 0)
    return idc.Comment(ea)

def read_range_selection():
    if idaapi.IDA_SDK_VERSION >= 700:
        return idaapi.read_range_selection(None)
    return idaapi.read_selection()
#########################################################################################################

# return a normalized code, or None if input is invalid
def convert_hexstr(code):
    # normalize code
    code = code.lower()
    code = code.replace(' ', '')    # remove space
    code = code.replace('h', '')    # remove trailing 'h' in 90h
    code = code.replace('0x', '')   # remove 0x
    code = code.replace('\\x', '')  # remove \x
    code = code.replace(',', '')    # remove ,
    code = code.replace(';', '')    # remove ;
    code = code.replace('"', '')    # remove "
    code = code.replace("'", '')    # remove '
    code = code.replace("+", '')    # remove +

    # single-digit hexcode?
    if len(code) == 1 and ((code >= '0' and code <= '9') or (code >= 'a' and code <= 'f')):
        # stick 0 in front (so 'a' --> '0a')
        code = '0' + code

    # odd-length is invalid
    if len(code) % 2 != 0:
        return None

    try:
        hex_data = code.decode('hex')
        # we want a list of int
        return [ord(i) for i in hex_data]
    except:
        # invalid hex
        return None


# download a file from @url, then return (result, file-content)
# return (0, content) on success, or ({1|2}, None) on download failure
def url_download(url):
    # Import for python2
    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)


def get_name_value(_from, name):
    """
    Fixed: the return value truncated(32 bit) of get_name_value function that analyzed 64 bit binary file about ida64 for win.

    eg:
    type == idaapi.NT_BYTE
    (type, value) = idaapi.get_name_value(idc.BADADDR, "wcschr") # ida64 for win

    value = 0x14003d3f0L is correct  ida64 > 7.x for macOS
    value = 0x4003d3f0L is truncated ida64 >= 6.x for win, ida64 == 6.x for macOS

    :param _from: ea
    :param name: name string
    :return: tuple
    """
    name = to_string(name)
    (typ, value) = idaapi.get_name_value(_from, name)
    if typ == idaapi.NT_BYTE:  # typ is byte name (regular name)
        value = idaapi.get_name_ea(_from, name)
    return (typ, value)


## Main Keypatch class
class Keypatch_Asm:
    # supported architectures
    arch_lists = {
        "X86 16-bit": (KS_ARCH_X86, KS_MODE_16),                # X86 16-bit
        "X86 32-bit": (KS_ARCH_X86, KS_MODE_32),                # X86 32-bit
        "X86 64-bit": (KS_ARCH_X86, KS_MODE_64),                # X86 64-bit
        "ARM": (KS_ARCH_ARM, KS_MODE_ARM),                      # ARM
        "ARM Thumb": (KS_ARCH_ARM, KS_MODE_THUMB),              # ARM Thumb
        "ARM64 (ARMV8)": (KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN),# ARM64
        "Hexagon": (KS_ARCH_HEXAGON, KS_MODE_BIG_ENDIAN),       # Hexagon
        "Mips32": (KS_ARCH_MIPS, KS_MODE_MIPS32),               # Mips32
        "Mips64": (KS_ARCH_MIPS, KS_MODE_MIPS64),               # Mips64
        "PowerPC 32": (KS_ARCH_PPC, KS_MODE_PPC32),             # PPC32
        "PowerPC 64": (KS_ARCH_PPC, KS_MODE_PPC64),             # PPC64
        "Sparc 32": (KS_ARCH_SPARC, KS_MODE_SPARC32),           # Sparc32
        "Sparc 64": (KS_ARCH_SPARC, KS_MODE_SPARC64),           # Sparc64
        "SystemZ": (KS_ARCH_SYSTEMZ, KS_MODE_BIG_ENDIAN),       # SystemZ
    }

    endian_lists = {
        "Little Endian": KS_MODE_LITTLE_ENDIAN,                 # little endian
        "Big Endian": KS_MODE_BIG_ENDIAN,                       # big endian
    }

    syntax_lists = {
        "Intel": KS_OPT_SYNTAX_INTEL,
        "Nasm": KS_OPT_SYNTAX_NASM,
        "AT&T": KS_OPT_SYNTAX_ATT
    }

    def __init__(self, arch=None, mode=None):
        # update current arch and mode
        self.update_hardware_mode()

        # override arch & mode if provided
        if arch is not None:
            self.arch = arch
        if mode is not None:
            self.mode = mode

        # IDA uses Intel syntax by default
        self.syntax = KS_OPT_SYNTAX_INTEL

    # return Keystone arch & mode (with endianess included)
    @staticmethod
    def get_hardware_mode():
        (arch, mode) = (None, None)

        # heuristically detect hardware setup
        info = idaapi.get_inf_structure()

        try:
            cpuname = info.procname.lower()
        except:
            cpuname = info.procName.lower()

        try:
            # since IDA7 beta 3 (170724) renamed inf.mf -> is_be()/set_be()
            is_be = idaapi.cvar.inf.is_be()
        except:
            # older IDA versions
            is_be = idaapi.cvar.inf.mf
        # print("Keypatch BIG_ENDIAN = %s" %is_be)

        if cpuname == "metapc":
            arch = KS_ARCH_X86
            if info.is_64bit():
                mode = KS_MODE_64
            elif info.is_32bit():
                mode = KS_MODE_32
            else:
                mode = KS_MODE_16
        elif cpuname.startswith("arm"):
            # ARM or ARM64
            if info.is_64bit():
                arch = KS_ARCH_ARM64
                if is_be:
                    mode = KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_LITTLE_ENDIAN
            else:
                arch = KS_ARCH_ARM
                # either big-endian or little-endian
                if is_be:
                    mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("sparc"):
            arch = KS_ARCH_SPARC
            if info.is_64bit():
                mode = KS_MODE_SPARC64
            else:
                mode = KS_MODE_SPARC32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("ppc"):
            arch = KS_ARCH_PPC
            if info.is_64bit():
                mode = KS_MODE_PPC64
            else:
                mode = KS_MODE_PPC32
            if cpuname == "ppc":
                # do not support Little Endian mode for PPC
                mode += KS_MODE_BIG_ENDIAN
        elif cpuname.startswith("mips"):
            arch = KS_ARCH_MIPS
            if info.is_64bit():
                mode = KS_MODE_MIPS64
            else:
                mode = KS_MODE_MIPS32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
            arch = KS_ARCH_SYSTEMZ
            mode = KS_MODE_BIG_ENDIAN

        return (arch, mode)

    def update_hardware_mode(self):
        (self.arch, self.mode) = self.get_hardware_mode()

    # normalize assembly code
    # remove comment at the end of assembly code
    @staticmethod
    def asm_normalize(text):
        text = ' '.join(text.split())
        if text.rfind(';') != -1:
            return text[:text.rfind(';')].strip()

        return text.strip()

    @staticmethod
    # check if input address is valid
    # return
    #       -1  invalid address at target binary
    #        0  type mismatch of input address
    #        1  valid address at target binary
    def check_address(address):
        try:
            if is_mapped(address):
                return 1
            else:
                return -1
        except:
            # invalid type
            return 0

    ### resolve IDA names from input asm code
    # todo: a better syntax parser for all archs
    def ida_resolve(self, assembly, address=idc.BADADDR):
        def _resolve(_op, ignore_kw=True):
            names = re.findall(r"[\$a-z0-9_:\.]+", _op, re.I)

            # try to resolve all names
            for name in names:
                # ignore known keywords
                if ignore_kw and name in ('byte', 'near', 'short', 'word', 'dword', 'ptr', 'offset'):
                    continue

                sym = name

                # split segment reg
                parts = name.partition(':')
                if parts[2] != '':
                    sym = parts[2]

                (typ, value) = get_name_value(address, sym)

                # skip if name doesn't exist or segment / segment registers
                if typ in (idaapi.NT_SEG, idaapi.NT_NONE):
                    continue

                _op = _op.replace(sym, '0x{0:X}'.format(value))

            return _op

        if self.check_address(address) == 0:
            print("Keypatch: WARNING: invalid input address {0}".format(address))
            return assembly

        # for now, we only support IDA name resolve for X86, ARM, ARM64, MIPS, PPC, SPARC
        if not (self.arch in (KS_ARCH_X86, KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_MIPS, KS_ARCH_PPC, KS_ARCH_SPARC)):
            return assembly

        _asm = assembly.partition(' ')
        mnem = _asm[0]
        opers = _asm[2].split(',')

        for idx, op in enumerate(opers):
            _op = list(op.partition('['))
            ignore_kw = True
            if _op[1] == '':
                _op[2] = _op[0]
                _op[0] = ''
            else:
                _op[0] = _resolve(_op[0], ignore_kw=True)
                ignore_kw = False

            _op[2] = _resolve(_op[2], ignore_kw=ignore_kw)

            opers[idx] = ''.join(_op)

        asm = "{0} {1}".format(mnem, ','.join(opers))
        return asm

    # return bytes of instruction or data
    # return None on failure
    def ida_get_item(self, address, hex_output=False):
        if self.check_address(address) != 1:
            # not a valid address
            return (None, 0)

        # return None if address is in the middle of instruction / data
        if address != get_item_head(address):
            return (None, 0)

        size = get_item_size(address)
        item = get_bytes(address, size)

        if item is None:
            return (None, 0)

        if hex_output:
            item = to_hexstr(item)

        return (item, size)

    @staticmethod
    def get_op_dtype_name(ea, op_idx):
        dtyp_lists = {
            idaapi.dt_byte: 'byte',     #  8 bit
            idaapi.dt_word: 'word',     #  16 bit
            idaapi.dt_dword: 'dword',   #  32 bit
            idaapi.dt_float: 'dword',   #  4 byte
            idaapi.dt_double: 'dword',  #  8 byte
            #idaapi.dt_tbyte = 5        #  variable size (ph.tbyte_size)
            #idaapi.dt_packreal = 6         #  packed real format for mc68040
            idaapi.dt_qword: 'qword',   #  64 bit
            idaapi.dt_byte16: 'xmmword',#  128 bit
            #idaapi.dt_code = 9         #  ptr to code (not used?)
            #idaapi.dt_void = 10        #  none
            #idaapi.dt_fword = 11       #  48 bit
            #idaapi.dt_bitfild = 12     #  bit field (mc680x0)
            #idaapi.dt_string = 13      #  pointer to asciiz string
            #idaapi.dt_unicode = 14     #  pointer to unicode string
            #idaapi.dt_3byte = 15       #  3-byte data
            #idaapi.dt_ldbl = 16        #  long double (which may be different from tbyte)
            idaapi.dt_byte32: 'ymmword',# 256 bit
        }

        dtype, dtyp_size = get_dtype(ea, op_idx)
        if dtype == idaapi.dt_tbyte and dtyp_size == 10:
          return 'xword'

        dtyp_name = dtyp_lists.get(dtype, None)

        return dtyp_name

    # return asm instructions from start to end
    def ida_get_disasm_range(self, start, end):
        codes = []
        while start < end:
            asm = self.asm_normalize(idc.GetDisasm(start))
            if asm == None:
                asm = ''
            codes.append(asm)
            start = start + get_item_size(start)

        return codes

    # get disasm from IDA
    # return '' on invalid address
    def ida_get_disasm(self, address, fixup=False):

        def GetMnem(asm):
            sp = asm.find(' ')
            if sp == -1:
                return asm
            return asm[:sp]

        if self.check_address(address) != 1:
            # not a valid address
            return ''

        # return if address is in the middle of instruction / data
        if address != get_item_head(address):
            return ''

        asm = self.asm_normalize(idc.GetDisasm(address))
        # for now, only support IDA syntax fixup for Intel CPU
        if not fixup or self.arch != KS_ARCH_X86:
            return asm

        # KS_ARCH_X86 mode
        # rebuild disasm code from IDA
        i = 0
        mnem = GetMnem(asm)
        if mnem == '' or mnem in ('rep', 'repne', 'repe'):
            return asm

        opers = []
        while get_operand_type(address, i) > 0 and i < 6:
            t = get_operand_type(address, i)
            o = print_operand(address, i)

            if t in (idc.o_mem, idc.o_displ):
                parts = list(o.partition(':'))
                if parts[2] == '':
                    parts[2] = parts[0]
                    parts[0] = ''

                if '[' not in parts[2]:
                    parts[2] = '[{0}]'.format(parts[2])

                o = ''.join(parts)

                if 'ptr ' not in o:
                    dtyp_name = self.get_op_dtype_name(address, i)
                    if dtyp_name != None:
                        o = "{0} ptr {1}".format(dtyp_name, o)

            opers.append(o)
            i += 1

        asm = mnem
        for o in opers:
            if o != '':
                asm = "{0} {1},".format(asm, o)

        asm = asm.strip(',')
        return asm

    # assemble code with Keystone
    # return (encoding, count), or (None, 0) on failure
    def assemble(self, assembly, address, arch=None, mode=None, syntax=None):

        # return assembly with arithmetic equation evaluated
        def eval_operand(assembly, start, stop, prefix=''):
            imm = assembly[start+1:stop]
            try:
                eval_imm = eval(imm)
                if eval_imm > 0x80000000:
                    eval_imm = 0xffffffff - eval_imm
                    eval_imm += 1
                    eval_imm = -eval_imm
                return assembly.replace(prefix + imm, prefix + hex(eval_imm))
            except:
                return assembly

        # IDA uses different syntax from Keystone
        # sometimes, we can convert code to be consumable by Keystone
        def fix_ida_syntax(assembly):

            # return True if this insn needs to be fixed
            def check_arm_arm64_insn(arch, mnem):
                if arch == KS_ARCH_ARM:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return False
                elif arch == KS_ARCH_ARM64:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return mnem in ("stp")
                return False

            # return True if this insn needs to be fixed
            def check_ppc_insn(mnem):
                return mnem in ("stw")

            # replace the right most string occurred
            def rreplace(s, old, new):
                li = s.rsplit(old, 1)
                return new.join(li)

            # convert some ARM pre-UAL assembly to UAL, so Keystone can handle it
            # example: streqb --> strbeq
            def fix_arm_ual(mnem, assembly):
                # TODO: this is not an exhaustive list yet
                if len(mnem) != 6:
                    return assembly

                if (mnem[-1] in ('s', 'b', 'h', 'd')):
                    #print(">> 222", mnem[3:5])
                    if mnem[3:5] in ("cc", "eq", "ne", "hs", "lo", "mi", "pl", "vs", "vc", "hi", "ls", "ge", "lt", "gt", "le", "al"):
                        return assembly.replace(mnem, mnem[:3] + mnem[-1] + mnem[3:5], 1)

                return assembly

            if self.arch != KS_ARCH_X86:
                assembly = assembly.lower()
            else:
                # Keystone does not support immediate 0bh, but only 0Bh
                assembly = assembly.upper()

            # however, 0X must be converted to 0x
            # Keystone should fix this limitation in the future
            assembly = assembly.replace("0X", " 0x")

            _asm = assembly.partition(' ')
            mnem = _asm[0]
            if mnem == '':
                return assembly

            # for PPC, Keystone does not accept registers with 'r' prefix,
            # but only the number behind. lets try to fix that here by
            # removing the prefix 'r'.
            if self.arch == KS_ARCH_PPC:
                for n in range(32):
                    r = " r%u," %n
                    if r in assembly:
                        assembly = assembly.replace(r, " %u," %n)
                for n in range(32):
                    r = "(r%u)" %n
                    if r in assembly:
                        assembly = assembly.replace(r, "(%u)" %n)
                for n in range(32):
                    r = ", r%u" %n
                    if assembly.endswith(r):
                        assembly = rreplace(assembly, r, ", %u" %n)

            if self.arch == KS_ARCH_X86:
                if mnem == "RETN":
                    # replace retn with ret
                    return assembly.replace('RETN', 'RET', 1)
                if 'OFFSET ' in assembly:
                    return assembly.replace('OFFSET ', ' ')
                if mnem in ('CALL', 'JMP') or mnem.startswith('LOOP'):
                    # remove 'NEAR PTR'
                    if ' NEAR PTR ' in assembly:
                        return assembly.replace(' NEAR PTR ', ' ')
                elif mnem[0] == 'J':
                    # JMP instruction
                    if ' SHORT ' in assembly:
                        # remove ' short '
                        return assembly.replace(' SHORT ', ' ')
            elif self.arch in (KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_PPC):
                # *** ARM
                # LDR     R1, [SP+rtld_fini],#4
                # STR     R2, [SP,#-4+rtld_fini]!
                # STR     R0, [SP,#fini]!
                # STR     R12, [SP,#4+var_8]!

                # *** ARM64
                # STP     X29, X30, [SP,#-0x10+var_150]!
                # STR     W0, [X29,#0x150+var_8]
                # LDR     X0, [X0,#(qword_4D6678 - 0x4D6660)]
                # TODO:
                # ADRP    X19, #interactive@PAGE

                # *** PPC
                # stw     r5, 0x120+var_108(r1)

                if self.arch == KS_ARCH_ARM and mode == KS_MODE_THUMB:
                    assembly = assembly.replace('movt.w', 'movt')

                if self.arch == KS_ARCH_ARM:
                    #print(">> before UAL fix: ", assembly)
                    assembly = fix_arm_ual(mnem, assembly)
                    #print(">> after UAL fix: ", assembly)

                if check_arm_arm64_insn(self.arch, mnem) or (("[" in assembly) and ("]" in assembly)):
                    bang = assembly.find('#')
                    bracket = assembly.find(']')
                    if bang != -1 and bracket != -1 and bang < bracket:
                        return eval_operand(assembly, bang, bracket, '#')
                    elif '+0x0]' in assembly:
                        return assembly.replace('+0x0]', ']')
                elif check_ppc_insn(mnem):
                    start = assembly.find(', ')
                    stop = assembly.find('(')
                    if start != -1 and stop != -1 and start < stop:
                        return eval_operand(assembly, start, stop)
            return assembly

        def is_thumb(address):
            return get_sreg(address, 'T') == 1

        if self.check_address(address) == 0:
            return (None, 0)

        # use default syntax, arch and mode if not provided
        if syntax is None:
            syntax = self.syntax
        if arch is None:
            arch = self.arch
        if mode is None:
            mode = self.mode

        if arch == KS_ARCH_ARM and is_thumb(address):
            mode = KS_MODE_THUMB

        try:
            ks = Ks(arch, mode)
            if arch == KS_ARCH_X86:
                ks.syntax = syntax
            encoding, count = ks.asm(fix_ida_syntax(assembly), address)
        except KsError as e:
            # keep the below code for debugging
            #print("Keypatch Error: {0}".format(e))
            #print("Original asm: {0}".format(assembly))
            #print("Fixed up asm: {0}".format(fix_ida_syntax(assembly)))
            encoding, count = None, 0

        return (encoding, count)


    # patch at address, return the number of written bytes & original data
    # this process can fail in some cases
    @staticmethod
    def patch_raw(address, patch_data, size):
        ea = address
        orig_data = ''

        while ea < (address + size):

            if not has_value(get_full_flags(ea)):
                print("Keypatch: FAILED to read data at 0x{0:X}".format(ea))
                break

            orig_byte = get_wide_byte(ea)
            orig_data += chr(orig_byte)
            patch_byte = ord(patch_data[ea - address])

            if patch_byte != orig_byte:
                # patch one byte
                if idaapi.patch_byte(ea, patch_byte) != 1:
                    print("Keypatch: FAILED to patch byte at 0x{0:X} [0x{1:X}]".format(ea, patch_byte))
                    break
            ea += 1
        return (ea - address, orig_data)

    # patch at address, return the number of written bytes & original data
    # on patch failure, we revert to the original code, then return (None, None)
    def patch(self, address, patch_data, size):
        # save original function end to fix IDA re-analyze issue after patching
        orig_func_end = get_func_attr(address, idc.FUNCATTR_END)

        (patched_len, orig_data) = self.patch_raw(address, patch_data, size)

        if size != patched_len:
            # patch failure
            if patched_len > 0:
                # revert the changes
                (rlen, _) = self.patch_raw(address, orig_data, patched_len)
                if rlen == patched_len:
                    print("Keypatch: successfully reverted changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))
                else:
                    print("Keypatch: FAILED to revert changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))

            return (None, None)

        # ask IDA to re-analyze the patched area
        if orig_func_end == idc.BADADDR:
            # only analyze patched bytes, otherwise it would take a lot of time to re-analyze the whole binary
            plan_and_wait(address, address + patched_len + 1)
        else:
            plan_and_wait(address, orig_func_end)

            # try to fix IDA function re-analyze issue after patching
            set_func_end(address, orig_func_end)

        return (patched_len, orig_data)

    # return number of bytes patched
    # return
    #    0  Invalid assembly
    #   -1  PatchByte failure
    #   -2  Can't read original data
    #   -3  Invalid address
    def patch_code(self, address, assembly, syntax, padding, save_origcode, orig_asm=None, patch_data=None, patch_comment=None, undo=False):
        global patch_info

        if self.check_address(address) != 1:
            # not a valid address
            return -3

        orig_comment = get_comment(address)
        if orig_comment is None:
            orig_comment = ''

        nop_comment = ""
        padding_len = 0
        if not undo:
            # we are patching via Patcher
            (orig_encoding, orig_len) = self.ida_get_item(address)
            if (orig_encoding, orig_len) == (None, 0):
                return -2

            (encoding, count) = self.assemble(assembly, address, syntax=syntax)
            if encoding is None:
                return 0

            patch_len = len(encoding)
            patch_data = ''.join(chr(c) for c in encoding)

            if patch_data == orig_encoding:
                #print("Keypatch: no need to patch, same encoding data [{0}] at 0x{1:X}".format(to_hexstr(orig_encoding), address))
                return orig_len

            # for now, only support NOP padding on Intel CPU
            if padding and self.arch == KS_ARCH_X86:
                if patch_len < orig_len:
                    padding_len = orig_len - patch_len
                    patch_len = orig_len
                    patch_data = patch_data.ljust(patch_len, X86_NOP)
                elif patch_len > orig_len:
                    patch_end = address + patch_len - 1
                    ins_end = get_item_end(patch_end)
                    padding_len = ins_end - patch_end - 1

                    if padding_len > 0:
                        patch_len = ins_end - address
                        patch_data = patch_data.ljust(patch_len, X86_NOP)

                if padding_len > 0:
                    nop_comment = "\nKeypatch padded NOP to next boundary: {0} bytes".format(padding_len)

            orig_asm = self.ida_get_disasm_range(address, address + patch_len)
        else:
            # we are reverting the change via "Undo" menu
            patch_len = len(patch_data)

        (plen, p_orig_data) = self.patch(address, patch_data, patch_len)
        if plen is None:
            # failed to patch
            return -1

        if not undo: # we are patching
            new_patch_comment = None
            if save_origcode is True:
                # append original instruction to comments
                if orig_comment == '':
                    new_patch_comment = "Keypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)
                else:
                    new_patch_comment = "\nKeypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)

                new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
                set_comment(address, new_comment)

            if padding_len == 0:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))
            else:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}], with {4} byte(s) NOP padded".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data), padding_len))
            # save this patching for future "undo"
            patch_info.append((address, assembly, p_orig_data, new_patch_comment))
        else:   # we are reverting
            if patch_comment:
                # clean previous IDA comment by replacing it with ''
                new_comment = orig_comment.replace(patch_comment, '')
                set_comment(address, new_comment)

            print("Keypatch: successfully reverted {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))

        return plen

    # fill a range of code [addr_begin, addr_end].
    # return the length of patched area
    # on failure, return 0 = wrong input, -1 = failed to patch
    def fill_code(self, addr_begin, addr_end, assembly, syntax, padding, save_origcode, orig_asm=None):
        # treat input as assembly code first
        (encoding, _) =  self.assemble(assembly, addr_begin, syntax=syntax)

        if encoding is None:
            # input might be a hexcode string. try to convert it to raw bytes
            encoding = convert_hexstr(assembly)

        if encoding is None:
            # invalid input: this is neither assembly nor hexcode string
            return 0

        # save original assembly code before overwritting them
        orig_asm = self.ida_get_disasm_range(addr_begin, addr_end)

        # save original comment at addr_begin
        # TODO: save comments in this range, but how to interleave them?
        orig_comment = get_comment(addr_begin)
        if orig_comment is None:
            orig_comment = ''

        patch_data = ""
        assembly_new = []
        size = addr_end - addr_begin
        # calculate filling data
        encode_chr = ''.join(chr(c) for c in encoding)
        while True:
            if len(patch_data) + len(encode_chr) <= size:
                patch_data = patch_data + encode_chr
                assembly_new += [assembly.strip()]
            else:
                break

        # for now, only support NOP padding on Intel CPU
        if padding and self.arch == KS_ARCH_X86:
            for i in range(size -len(patch_data)):
                assembly_new += ["nop"]
            patch_data = patch_data.ljust(size, X86_NOP)

        (plen, p_orig_data) = self.patch(addr_begin, patch_data, len(patch_data))
        if plen is None:
            # failed to patch
            return -1

        new_patch_comment = ''
        # append original instruction to comments
        if save_origcode is True:
            if orig_comment == '':
                new_patch_comment = "Keypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))
            else:
                new_patch_comment = "\nKeypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))

            new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
            set_comment(addr_begin, new_comment)

        print("Keypatch: successfully filled range [0x{0:X}:0x{1:X}] ({2} bytes) with \"{3}\", replaced \"{4}\"".format(
                    addr_begin, addr_end - 1, addr_end - addr_begin, assembly, '; '.join(orig_asm)))

        # save this modification for future "undo"
        patch_info.append((addr_begin, '\n  '.join(assembly_new), p_orig_data, new_patch_comment))

        return plen


    ### Form helper functions
    @staticmethod
    def dict_to_ordered_list(dictionary):
        l = sorted(list(dictionary.items()), key=lambda t: t[0], reverse=False)
        keys = [i[0] for i in l]
        values = [i[1] for i in l]

        return (keys, values)

    def get_value_by_idx(self, dictionary, idx, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            val = values[idx]
        except IndexError:
            val = default

        return val

    def find_idx_by_value(self, dictionary, value, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            idx = values.index(value)
        except:
            idx = default

        return idx

    def get_arch_by_idx(self, idx):
        return self.get_value_by_idx(self.arch_lists, idx)

    def find_arch_idx(self, arch, mode):
        return self.find_idx_by_value(self.arch_lists, (arch, mode))

    def get_syntax_by_idx(self, idx):
        return self.get_value_by_idx(self.syntax_lists, idx, self.syntax)

    def find_syntax_idx(self, syntax):
        return self.find_idx_by_value(self.syntax_lists, syntax)
    ### /Form helper functions


# Common ancestor form to be derived by Patcher, FillRange & Search
class Keypatch_Form(idaapi.Form):
    # prepare for form initializing
    def setup(self, kp_asm, address, assembly=None):
        self.kp_asm = kp_asm
        self.address = address

        # update ordered list of arch and syntax
        self.syntax_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.syntax_lists)[0]
        self.arch_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.arch_lists)[0]

        # update current arch & mode
        self.kp_asm.update_hardware_mode()

        # find right value for c_arch & c_endian controls
        mode = self.kp_asm.mode
        self.endian_id = 0   # little endian
        if self.kp_asm.mode & KS_MODE_BIG_ENDIAN:
            self.endian_id = 1   # big endian
            mode = self.kp_asm.mode - KS_MODE_BIG_ENDIAN

        self.arch_id = self.kp_asm.find_arch_idx(self.kp_asm.arch, mode)

        self.syntax_id = 0  # to make non-X86 arch happy
        if self.kp_asm.arch == KS_ARCH_X86:
            self.syntax_id = self.kp_asm.find_syntax_idx(self.kp_asm.syntax)

        # get original instruction and bytes
        self.orig_asm = kp_asm.ida_get_disasm(address)
        (self.orig_encoding, self.orig_len) = kp_asm.ida_get_item(address, hex_output=True)
        if self.orig_encoding is None:
            self.orig_encoding = ''

        if assembly is None:
            self.asm = self.kp_asm.ida_get_disasm(self.address, fixup=True)
        else:
            self.asm = assembly

    # update Encoding control
    # return True on success, False on failure
    def _update_encoding(self, arch, mode):
        try:
            syntax = None
            if arch == KS_ARCH_X86:
                syntax_id = self.GetControlValue(self.c_syntax)
                syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

            address = self.GetControlValue(self.c_addr)
            try:
                is_mapped(address)
            except:
                # invalid address value
                address = 0

            assembly = self.GetControlValue(self.c_assembly)
            raw_assembly = self.kp_asm.ida_resolve(assembly, address)
            self.SetControlValue(self.c_raw_assembly, raw_assembly)

            (encoding, count) =  self.kp_asm.assemble(raw_assembly, address, arch=arch,
                                                    mode=mode, syntax=syntax)

            if encoding is None:
                self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                self.SetControlValue(self.c_encoding_len, 0)
                return False
            else:
                text = ""
                for byte in encoding:
                    text += "%02X " % byte
                text.strip()
                if text == "":
                    # error?
                    self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                    return False
                else:
                    self.SetControlValue(self.c_encoding, text.strip())
                    self.SetControlValue(self.c_encoding_len, len(encoding))
                    return True
        except Exception as e:
            print(str(e))
            import traceback
            traceback.print_exc()
            self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
            return False

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        return 1

    # update Patcher & Fillrange controls
    def update_patchform(self, fid):
        self.EnableField(self.c_endian, False)
        self.EnableField(self.c_addr, False)

        (arch, mode) = (self.kp_asm.arch, self.kp_asm.mode)
        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # do not show Endian control
            self.ShowField(self.c_endian, False)
            # allow to choose Syntax
            self.ShowField(self.c_syntax, True)
            self.ShowField(self.c_opt_padding, True)
        else:   # do not show Syntax control for non-X86 mode
            self.ShowField(self.c_syntax, False)
            # for now, we do not support padding for non-X86 archs
            self.ShowField(self.c_opt_padding, False)
            #self.EnableField(self.c_opt_padding, False)

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1

    # update some controls - including Encoding control
    def update_controls(self, arch, mode):
        # Fixup & Encoding-len are read-only controls
        self.EnableField(self.c_raw_assembly, False)
        self.EnableField(self.c_encoding_len, False)

        # Encoding is enable to allow user to select & copy
        self.EnableField(self.c_encoding, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update encoding with live assembling
        self._update_encoding(arch, mode | endian)

        return 1

    # get Patcher/FillRange options
    def get_opts(self, name=None):
        names = self.c_opt_chk.children_names
        val = self.c_opt_chk.value
        opts = {}
        for i in range(len(names)):
            opts[names[i]] = val & (2**i)

        if name != None:
            opts[name] = val

        return opts


# Fill Range form
class Keypatch_FillRange(Keypatch_Form):
    def __init__(self, kp_asm, addr_begin, addr_end, assembly=None, opts=None):
        self.setup(kp_asm, addr_begin, assembly)
        self.addr_end = addr_end

        # create FillRange form
        super(Keypatch_FillRange, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Fill Range

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Start      :{c_addr}>
            <End        :{c_addr_end}>
            <Size       :{c_size}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=addr_begin, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_addr_end': self.NumericInput(value=addr_end - 1, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_size': self.NumericInput(value=addr_end - addr_begin, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':idaapi.Form.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some controls read-only in FillRange mode
        self.EnableField(self.c_size, False)
        self.EnableField(self.c_addr_end, False)

        return self.update_patchform(fid)


# Patcher form
class Keypatch_Patcher(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None, opts=None):
        self.setup(kp_asm, address, assembly)

        # create Patcher form
        super(Keypatch_Patcher, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Patcher

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Address    :{c_addr}>
            <Original   :{c_orig_assembly}>
             <##-   Encode:{c_orig_encoding}>
             <##-   Size  :{c_orig_len}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_assembly': self.StringInput(value=self.orig_asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_encoding': self.StringInput(value=self.orig_encoding[:MAX_ENCODING_LEN], width=MAX_ENCODING_LEN),
            'c_orig_len': self.NumericInput(value=self.orig_len, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':self.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some fields read-only in Patch mode
        self.EnableField(self.c_orig_assembly, False)
        self.EnableField(self.c_orig_encoding, False)
        self.EnableField(self.c_orig_len, False)

        return self.update_patchform(fid)


# Search position chooser
class SearchResultChooser(Choose):
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False):
        super(SearchResultChooser, self).__init__(
            title,
            [["Address", Choose.CHCOL_HEX|40]],
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = items
        self.selcount = 0
        self.modal = modal

    def OnClose(self):
        return

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

    def OnGetLine(self, n):
        res = self.items[n]
        res = [idc.atoa(res[0])]
        return res

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

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


# Search form
class Keypatch_Search(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None):
        self.setup(kp_asm, address, assembly)

        # create Search form
        super(Keypatch_Search, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Search
KEYPATCH:: Search

            {FormChangeCb}
            <A~r~ch       :{c_arch}>
            <E~n~dian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <A~d~dress    :{c_addr}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            """, {
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_arch': self.DropdownListControl(
                          items = self.arch_keys,
                          readonly = True,
                          selval = self.arch_id,
                          width = 32),
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    ################################ IDA 6/7 Compatibility method #######################################
    def find_binary(self, ea):
        if idaapi.IDA_SDK_VERSION >= 700:
            return ida_search.find_binary(ea, idaapi.cvar.inf.max_ea, self.GetControlValue(self.c_encoding), 16, idc.SEARCH_DOWN)
        return idc.FindBinary(ea, idc.SEARCH_DOWN, self.GetControlValue(self.c_encoding))
    #####################################################################################################

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # handle the search button
        if fid == -2:
            address = 0
            addresses = []
            while address != idc.BADADDR:
                address = self.find_binary(address)
                if address == idc.BADADDR:
                    break
                addresses.append([address])
                address = address + 1
            c = SearchResultChooser("Searching for [{0}]".format(self.GetControlValue(self.c_raw_assembly)), addresses)
            c.show()
            return 1

        # only Search mode allows to select arch+mode
        arch_id = self.GetControlValue(self.c_arch)
        (arch, mode) = self.kp_asm.get_arch_by_idx(arch_id)

        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # enable Syntax and disable Endian for x86
            self.ShowField(self.c_syntax, True)
            self.EnableField(self.c_syntax, True)
            self.syntax_id = self.GetControlValue(self.c_syntax)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, 0)
        elif arch in (KS_ARCH_ARM64, KS_ARCH_HEXAGON, KS_ARCH_SYSTEMZ):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        elif (arch == KS_ARCH_PPC) and (mode & KS_MODE_PPC32 != 0):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        else:
            # no Syntax & Endian option
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1


# About form
class About_Form(idaapi.Form):
    def __init__(self, version):
        # create About form
        super(About_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: About

            {FormChangeCb}
            Keypatch IDA plugin v%s, using Keystone Engine v%s.
            (c) Nguyen Anh Quynh + Thanh Nguyen, 2018.

            Keypatch is released under the GPL v2.
            Find more info at http://www.keystone-engine.org/keypatch
            """ %(version, keystone.__version__), {
            '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(KP_HOMEPAGE, new = 2)

        return 1


# Check-for-update form
class Update_Form(idaapi.Form):
    def __init__(self, version, message):
        # create Update form
        super(Update_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: Check for update

            {FormChangeCb}
            Your Keypatch 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(KP_HOMEPAGE, new = 2)

        return 1


try:
    # adapted from pull request #7 by @quangnh89
    class Kp_Menu_Context(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.widget_type == idaapi.BWN_DISASM:
                    return idaapi.AST_ENABLE_FOR_FORM
                else:
                    return idaapi.AST_DISABLE_FOR_FORM
            except Exception as e:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS

    # context menu for Patcher
    class Kp_MC_Patcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.patcher()
            return 1

    # context menu for Fill Range
    class Kp_MC_Fill_Range(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.fill_range()
            return 1

    # context menu for Undo
    class Kp_MC_Undo(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.undo()
            return 1

    # context menu for Search
    class Kp_MC_Search(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1

    # context menu for Check Update
    class Kp_MC_Updater(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.updater()
            return 1

    # context menu for About
    class Kp_MC_About(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.about()
            return 1

except:
    pass

# hooks for popup menu
class Hooks(idaapi.UI_Hooks):
    if idaapi.IDA_SDK_VERSION >= 700:
        # IDA >= 700 right click widget popup
        def finish_populating_widget_popup(self, form, popup):
            if idaapi.get_widget_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass
    else:
        # IDA < 700 right click popup
        def finish_populating_tform_popup(self, form, popup):
            # We'll add our action to all "IDA View-*"s.
            # If we wanted to add it only to "IDA View-A", we could
            # also discriminate on the widget's title:
            #
            #  if idaapi.get_tform_title(form) == "IDA View-A":
            #      ...
            #
            if idaapi.get_tform_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass


# check if we already initialized Keypatch
kp_initialized = False

#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class Keypatch_Plugin_t(idaapi.plugin_t):
    comment = "Keypatch plugin for IDA Pro (using Keystone framework)"
    help = "Find more information on Keypatch at http://keystone-engine.org/keypatch"
    wanted_name = "Keypatch Patcher"
    wanted_hotkey = "Ctrl-Alt-K"
    flags = idaapi.PLUGIN_KEEP

    def load_configuration(self):
        # default
        self.opts = {}

        # load configuration from file
        try:
            f = open(KP_CFGFILE, "rt")
            self.opts = json.load(f)
            f.close()
        except IOError:
            print("Keypatch: FAILED to load config file. Use default setup now.")
        except Exception as e:
            print("Keypatch: FAILED to load config file, with exception: {0}".format(str(e)))

        # use default values if not defined in config file
        if 'c_opt_padding' not in self.opts:
            self.opts['c_opt_padding'] = 1

        if 'c_opt_comment' not in self.opts:
            self.opts['c_opt_comment'] = 2

        self.opts['c_opt_chk'] = self.opts['c_opt_padding'] | self.opts['c_opt_comment']

    def init(self):
        global kp_initialized

        # register popup menu handlers
        try:
            Kp_MC_Patcher.register(self, "Patcher    (Ctrl-Alt-K)")
            Kp_MC_Fill_Range.register(self, "Fill Range")
            Kp_MC_Undo.register(self, "Undo last patching")
            Kp_MC_Search.register(self, "Search")
            Kp_MC_Updater.register(self, "Check for update")
            Kp_MC_About.register(self, "About")
        except:
            pass

        # setup popup menu
        self.hooks = Hooks()
        self.hooks.hook()

        self.opts = None
        if kp_initialized == False:
            kp_initialized = True

            if idaapi.IDA_SDK_VERSION >= 700:
                # Add menu IDA >= 7.0
                idaapi.attach_action_to_menu("Edit/Keypatch/Patcher", Kp_MC_Patcher.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/About", Kp_MC_About.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Check for update", Kp_MC_Updater.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Search", Kp_MC_Search.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Undo last patching", Kp_MC_Undo.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Fill Range", Kp_MC_Fill_Range.get_name(), idaapi.SETMENU_APP)
            else:
                # add Keypatch menu
                menu = idaapi.add_menu_item("Edit/Keypatch/", "Patcher     (Ctrl-Alt-K)", "", 1, self.patcher, None)
                if menu is not None:
                    idaapi.add_menu_item("Edit/Keypatch/", "About", "", 1, self.about, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Check for update", "", 1, self.updater, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Search", "", 1, self.search, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Undo last patching", "", 1, self.undo, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Fill Range", "", 1, self.fill_range, None)
                elif idaapi.IDA_SDK_VERSION < 680:
                    # older IDAPython (such as in IDAPro 6.6) does add new submenu.
                    # in this case, put Keypatch menu in menu Edit \ Patch program
                    # not sure about v6.7, so to be safe we just check against v6.8
                    idaapi.add_menu_item("Edit/Patch program/", "-", "", 0, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: About", "", 0, self.about, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Check for update", "", 0, self.updater, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Search", "", 0, self.search, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Undo last patching", "", 0, self.undo, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Fill Range", "", 0, self.fill_range, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Patcher     (Ctrl-Alt-K)", "", 0, self.patcher, None)

            print("=" * 80)
            print("Keypatch v{0} (c) Nguyen Anh Quynh & Thanh Nguyen, 2018".format(VERSION))
            print("Keypatch is using Keystone v{0}".format(keystone.__version__))
            print("Keypatch Patcher's shortcut key is Ctrl-Alt-K")
            print("Use the same hotkey Ctrl-Alt-K to open 'Fill Range' window on a selected range of code")
            print("To revert (undo) the last patching, choose menu Edit | Keypatch | Undo last patching")
            print("Keypatch Search is available from menu Edit | Keypatch | Search")
            print("Find more information about Keypatch at http://keystone-engine.org/keypatch")

            self.load_configuration()

            print("=" * 80)
            self.kp_asm = Keypatch_Asm()

        return idaapi.PLUGIN_KEEP

    def term(self):
        if self.hooks is not None:
            self.hooks.unhook()
            self.hooks = None

        if self.opts is None:
            return
        #save configuration to file
        try:
            json.dump(self.opts, open(KP_CFGFILE, "wt"))
        except Exception as e:
            print("Keypatch: FAILED to save config file, with exception: {0}".format(str(e)))
        else:
            print("Keypatch: configuration is saved to {0}".format(KP_CFGFILE))

    # null handler
    def menu_null(self):
        pass

    # handler for About menu
    def about(self):
        f = About_Form(VERSION)
        f.Execute()
        f.Free()

    # handler for Check-for-Update menu
    def updater(self):
        (r, content) = url_download(KP_GITHUB_VERSION)
        content = to_string(content)
        if r == 0:
            # find stable version
            sig = 'VERSION_STABLE = "'
            tmp = content[content.find(sig)+len(sig):]
            version_stable = tmp[:tmp.find('"')]

            # compare with the current version
            if version_stable == VERSION:
                f = Update_Form(VERSION, "Good, you are already on the latest stable version!")
                f.Execute()
                # free this form
                f.Free()
            else:
                f = Update_Form(VERSION, "Download latest stable version {0} from http://keystone-engine.org/keypatch".format(version_stable))
                f.Execute()
                # free this form
                f.Free()
        else:
            # fail to download
            warning("ERROR: Keypatch failed to connect to internet (Github). Try again later.")
            print("Keypatch: FAILED to connect to Github to check for latest update. Try again later.")

    # handler for Undo menu
    def undo(self):
        global patch_info
        if len(patch_info) == 0:
            # TODO: disable Undo menu?
            warning("ERROR: Keypatch already got to the last undo patching!")
        else:
            (address, assembly, p_orig_data, patch_comment) = patch_info[-1]

            # undo the patch
            self.kp_asm.patch_code(address, None, None, None, None, orig_asm=[assembly], patch_data=p_orig_data, patch_comment=patch_comment, undo=True)
            del(patch_info[-1])

    # handler for Search menu
    def search(self):
        address = get_screen_ea()
        f = Keypatch_Search(self.kp_asm, address)
        f.Execute()
        f.Free()

    # handler for Patcher menu
    def patcher(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return

        selection, addr_begin, addr_end = read_range_selection()
        if selection:
            # call Fill Range function on this selected code
            return self.fill_range()

        address = get_screen_ea()

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        while True:
            f = Keypatch_Patcher(self.kp_asm, address, assembly=init_assembly, opts=self.opts)
            ok = f.Execute()
            if ok == 1:
                try:
                    syntax = None
                    if f.kp_asm.arch == KS_ARCH_X86:
                        syntax_id = f.c_syntax.value
                        syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                    assembly = f.c_assembly.value
                    self.opts = f.get_opts('c_opt_chk')
                    padding = (self.opts.get("c_opt_padding", 0) != 0)
                    comment = (self.opts.get("c_opt_comment", 0) != 0)

                    raw_assembly = self.kp_asm.ida_resolve(assembly, address)

                    print("Keypatch: attempting to modify \"{0}\" at 0x{1:X} to \"{2}\"".format(
                            self.kp_asm.ida_get_disasm(address), address, assembly))

                    length = self.kp_asm.patch_code(address, raw_assembly, syntax, padding, comment, None)
                    if length > 0:
                        # update start address pointing to the next instruction
                        init_assembly = None
                        address += length
                    else:
                        init_assembly = f.c_assembly.value
                        if length == 0:
                            warning("ERROR: Keypatch found invalid assembly [{0}]".format(assembly))
                        elif length == -1:
                            warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(address))
                        elif length == -2:
                            warning("ERROR: Keypatch can't read original data at 0x{0:X}, try again".format(address))

                except KsError as e:
                    print("Keypatch Error: {0}".format(e))
            else:   # Cancel
                break
            f.Free()

    # handler for Fill Range menu
    def fill_range(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return

        selection, addr_begin, addr_end = read_range_selection()
        if not selection:
            warning("ERROR: Keypatch requires a range to be selected for fill in, try again")
            return

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        f = Keypatch_FillRange(self.kp_asm, addr_begin, addr_end, assembly=init_assembly, opts=self.opts)
        ok = f.Execute()
        if ok == 1:
            try:
                syntax = None
                if f.kp_asm.arch == KS_ARCH_X86:
                    syntax_id = f.c_syntax.value
                    syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                assembly = f.c_assembly.value
                self.opts = f.get_opts('c_opt_chk')
                padding = (self.opts.get("c_opt_padding", 0) != 0)
                comment = (self.opts.get("c_opt_comment", 0) != 0)

                raw_assembly = self.kp_asm.ida_resolve(assembly, addr_begin)

                print("Keypatch: attempting to fill range [0x{0:X}:0x{1:X}] with \"{2}\"".format(
                    addr_begin, addr_end - 1, assembly))

                length = self.kp_asm.fill_code(addr_begin, addr_end, raw_assembly, syntax, padding, comment, None)
                if length == 0:
                    warning("ERROR: Keypatch failed to process this input.")
                    print("Keypatch: FAILED to process this input '{0}'".format(assembly))
                elif length == -1:
                    warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(addr_begin))

            except KsError as e:
                print("Keypatch Error: {0}".format(e))

        # free this form
        f.Free()

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


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


repository: Ph0en1x-XMU/Ph0en1x-Team
file: Reverse/ida/IDA Pro 7.5.201028 (x86,x64,ARM,ARM64,PPC,PPC64,MIPS)/python/examples/hexrays/vds11.py
wanted_name: Optimize goto chains (IDAPython)
comment: Sample plugin11 for Hex-Rays decompiler
created_at: 2022-06-02 11:06:16+00:00
pushed_at: 2022-06-02 11:11:40+00:00
url:
https://github.com/Ph0en1x-XMU/Ph0en1x-Team/blob/a8d80aaea8227e418f5826bf05ef69a1bc8d84f1/Reverse/ida/IDA%20Pro%207.5.201028%20(x86,x64,ARM,ARM64,PPC,PPC64,MIPS)/python/examples/hexrays/vds11.py

content:
#
#      Hex-Rays Decompiler project
#      Copyright (c) 2007-2020 by Hex-Rays, support@hex-rays.com
#      ALL RIGHTS RESERVED.
#
#      Sample plugin for Hex-Rays Decompiler.
#      It installs a custom block optimization rule:
#
#        goto L1     =>        goto L2
#        ...
#      L1:
#        goto L2
#
#      In other words we fix a goto target if it points to a chain of gotos.
#      This improves the decompiler output in some cases.
#
#      This is a rewrite in Python of the vds11 example that comes with hexrays sdk.
#

import ida_bytes
import ida_range
import ida_kernwin
import ida_hexrays
import ida_typeinf
import ida_idaapi

class goto_optimizer_t(ida_hexrays.optblock_t):
    def func(self, blk):
        if self.handle_goto_chain(blk):
            return 1
        return 0

    def handle_goto_chain(self, blk):
        mgoto = blk.tail
        if not mgoto or mgoto.opcode != ida_hexrays.m_goto:
            return False

        visited = []
        t0 = mgoto.l.b
        i = t0
        mba = blk.mba

        # follow the goto chain
        while True:
            if i in visited:
                return False
            visited.append(i)
            b = mba.get_mblock(i)
            m2 = ida_hexrays.getf_reginsn(b.head)
            if not m2 or m2.opcode != ida_hexrays.m_goto:
                break
            i = m2.l.b

        if i == t0:
            return False # not a chain

        # all ok, found a goto chain
        mgoto.l.b = i # jump directly to the end of the chain

        # fix the successor/predecessor lists
        blk.succset[0] = i
        mba.get_mblock(i).predset.add(blk.serial)
        mba.get_mblock(t0).predset._del(blk.serial)

        # since we changed the control flow graph, invalidate the use/def chains.
        # stricly speaking it is not really necessary in our plugin because
        # we did not move around any microcode operands.
        mba.mark_chains_dirty()

        # it is a good idea to verify microcode after each change
        # however, it may be time consuming, so comment it out eventually
        mba.verify(True);
        return True

# --------------------------------------------------------------------------
# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Optimize goto chains (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin11 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.optimizer = goto_optimizer_t()
            self.optimizer.install()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.optimizer.remove()
    def run(self, arg):
        if arg == 1:
            return self.optimizer.remove()
        elif arg == 2:
            return self.optimizer.install()

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: panda-re/panda
file: panda/plugins/ida_taint2/hexrays_ida_taint2.py
description: Platform for Architecture-Neutral Dynamic Analysis
wanted_name: PANDA: Pseudocode ida_taint2
comment: Use ida_taint2 output to color tainted pseudocode lines
created_at: 2013-08-23 19:26:40+00:00
pushed_at: 2025-05-12 01:55:37+00:00
url:
https://github.com/panda-re/panda/blob/cc0a43e0347918d73dc1dfef9e3eb360432ec843/panda/plugins/ida_taint2/hexrays_ida_taint2.py

content:
"""
IDAPython plugin to ingest an ida_taint2 report to pseudocode windows.
The initial selections can be reused in a new pseudocode window, or the
selections can be changed.
"""
import csv
import idaapi
import ida_idaapi
import ida_kernwin
import ida_hexrays
import ida_lines


from copy import copy

from PyQt5.QtWidgets import (QDialog, QHBoxLayout, QLabel, QPushButton,
                             QVBoxLayout, QCheckBox, QFileDialog)

# assumes the ida_taint2_common.py file is placed in a plugins subfolder called
# common
from common.ida_taint2_common import (MaxTCNSelectDialog, ProcessSelectDialog,
                                      skip_csv_header)

INST_COLOR = 0x55AAFF

# index of process ID on row in ida_taint2 output file
PID_INDEX = 1

# index of taint compute number on row in ida_taint2 output file
TCN_INDEX = 4

filename = ""
selected_process = None
tainted_pcs = set()
has_tcns = False
process_tcn = None
maxtcn_displayed = None

class HIT2_ReuseDialog(QDialog):
    REUSE_SETTINGS = 1
    GET_NEW_SETTINGS = 2
    CANCEL_REQUEST = 3
    
    def __init__(self):
        super(HIT2_ReuseDialog, self).__init__()
        
        self.setWindowTitle("Reuse ida_taint2 Settings?")
        
        btn_ok = QPushButton("OK")
        btn_ok.clicked.connect(self.accept)
        btn_ok.setDefault(True)
        
        btn_cancel = QPushButton("Cancel")
        btn_cancel.clicked.connect(self.reject)
        
        self.lbl_file = QLabel()
        self.lbl_file.setText("File:  " + filename)
        self.lbl_pid = QLabel()
        self.lbl_pid.setText("Process:  " + selected_process['process_name'] +
                             "(" + str(selected_process['process_id']) + ")")
        
        self.lbl_tcn = QLabel()
        if (process_tcn is None):
            self.lbl_tcn.setText(
                "No taint compute numbers are available in this file.")
        else:
            self.lbl_tcn.setText("Displaying maximum taint compute number of " +
                                 str(maxtcn_displayed) + " out of " +
                                 str(process_tcn) + ".") 
        
        self.chkbx_reuse = QCheckBox("Reuse Above Settings")
        self.chkbx_reuse.setChecked(True)
        
        hbox = QHBoxLayout()
        hbox.addStretch(1)
        hbox.addWidget(btn_ok)
        hbox.addWidget(btn_cancel)
        
        vbox = QVBoxLayout()
        vbox.addWidget(self.lbl_file)
        vbox.addWidget(self.lbl_pid)
        vbox.addWidget(self.lbl_tcn)
        vbox.addWidget(self.chkbx_reuse)
        vbox.addLayout(hbox)
        
        self.setLayout(vbox)
        
    def isReuseSettings(self):
        return self.chkbx_reuse.isChecked()
        
    @classmethod
    def askToReuse(cls):
        rd = cls()
        if (QDialog.Accepted == rd.exec_()):
            if (rd.isReuseSettings()):
                return HIT2_ReuseDialog.REUSE_SETTINGS
            else:
                return HIT2_ReuseDialog.GET_NEW_SETTINGS
        else:
            return HIT2_ReuseDialog.CANCEL_REQUEST


class hexrays_ida_taint2_t(ida_idaapi.plugin_t):
    # PLUGIN_PROC means to load plugin when load new database, and unload it
    # when the database is closed
    flags = ida_idaapi.PLUGIN_PROC
    comment = "Use ida_taint2 output to color tainted pseudocode lines"
    help = "Use Alt-F5 to select process to taint with"
    wanted_name = "PANDA:  Pseudocode ida_taint2"
    wanted_hotkey = "Alt-F5"
    
    def tag_addrcode(self, s):
        # the eas associated with a line of pseudocode are encoded in
        # invisible COLOR_ADDR tags in the line of text
        if ((idaapi.COLOR_ON == s[0]) and (chr(idaapi.COLOR_ADDR) == s[1])):
            return True
        else:
            return False
        
    def clear_colors(self, cfunc):
        # clear the background colors on all the pseudocode lines
        # as may be applying colors for different process
        sv = cfunc.get_pseudocode()
        for i in range(len(sv)):
            # background color of -1 means to use the default
            sv[i].bgcolor = 0xFFFFFFFF
            
    def color_eas(self, cfunc, tainted_pcs):
        # the plugins/bap/utils/hexrays.py file found at
        # https://github.com/BinaryAnalysisPlatform/bap-ida-python/ was
        # invaluable in determining how to extract the effective addresses
        # from each pseudocode line
        sv = cfunc.get_pseudocode()
        anchor = ida_hexrays.ctree_anchor_t()
        for i in range(len(sv)):
            curline = copy(sv[i].line)
            while (len(curline) > 0):
                skipcode_index = ida_lines.tag_skipcode(curline)
                if (0 == skipcode_index):
                    # no code found, go to next character
                    curline = curline[1:]
                else:
                    if (self.tag_addrcode(curline)):
                        addr_tag = int(curline[2:skipcode_index], 16)
                        anchor.value = addr_tag
                        if (anchor.is_citem_anchor() and
                            not anchor.is_blkcmt_anchor()):
                            address = cfunc.treeitems.at(addr_tag).ea
                            if (address != ida_idaapi.BADADDR):
                                if (address in tainted_pcs):
                                    sv[i].bgcolor = INST_COLOR
                    curline = curline[skipcode_index:]

    def get_maxtcn(self, maxtcn_for_pid):
        global has_tcns
        global maxtcn_displayed
        global process_tcn
        global selected_process
        
        if (not has_tcns):
            process_tcn = None
            maxtcn_displayed = None
        else:
            sspid = str(selected_process['process_id'])
            new_process_tcn = maxtcn_for_pid[sspid]
            if (new_process_tcn > 0):
                new_maxtcn_displayed = MaxTCNSelectDialog.getMaxTCN(
                    selected_process['process_name'],
                    selected_process['process_id'], new_process_tcn)
                if (None == new_maxtcn_displayed):
                    # user changed mind - cancel entire selection
                    selected_process = None
                    process_tcn = None
                    maxtcn_displayed = None
                else:
                    process_tcn = new_process_tcn
                    maxtcn_displayed = new_maxtcn_displayed
            else:
                # only one TCN available for this process, 0
                process_tcn = 0
                maxtcn_displayed = 0
                
    def get_new_settings(self):
        global filename
        global has_tcns
        global selected_process
        
        # get output of ida_taint2 plugin to determine which lines get colored
        filename, _ = QFileDialog.getOpenFileName(None, "Open file", ".",
            "CSV Files (*.csv)")
        if filename == "":
            return
        ida_kernwin.msg("hexrays_ida_taint2:  file selected is " + filename + "\n")
            
        processes = set()
        maxtcn_for_pid = dict()
        input_file = open(filename, "r")
        reader = csv.reader(input_file)
        has_tcns = skip_csv_header(reader, True)

        for row in reader:
            processes.add((row[0], int(row[PID_INDEX])))
            if (has_tcns):
                cur_tcn = int(row[TCN_INDEX])
                if (row[PID_INDEX] not in maxtcn_for_pid):
                    maxtcn_for_pid[row[PID_INDEX]] = cur_tcn
                elif (cur_tcn > maxtcn_for_pid[row[PID_INDEX]]):
                    maxtcn_for_pid[row[PID_INDEX]] = cur_tcn
        input_file.close()
            
        selected_process = ProcessSelectDialog.selectProcess(processes,
            maxtcn_for_pid)
        if (None == selected_process):
            return
            
        self.get_maxtcn(maxtcn_for_pid)
            
    def update_taint_info(self):
        global filename
        global has_tcns
        global maxtcn_displayed
        global selected_process
        global tainted_pcs
        
        selected_pid = selected_process['process_id']
        input_file = open(filename, "r")
        reader = csv.reader(input_file)
        tainted_pcs = set()
        # skip header
        skip_csv_header(reader, False)
        for row in reader:
            pid = int(row[PID_INDEX])
            pc = int(row[2], 16)
           
            if pid != selected_pid:
                continue
                
            if (has_tcns):
                cur_tcn = int(row[TCN_INDEX])
                if (cur_tcn > maxtcn_displayed):
                    continue
                
            if (pc not in tainted_pcs):
                tainted_pcs.add(pc)
        input_file.close()
            
    def color_pseudocode(self, widget, clear_old):
        global filename
        global selected_process
        global tainted_pcs
        
        vu = ida_hexrays.get_widget_vdui(widget)
        
        cfunc = vu.cfunc
        if cfunc is None:
            ida_kernwin.msg("hexrays_ida_taint2:  Widget has no " +
                "decompiled pseudocode!\n")
            return True
            
        if (0 == len(tainted_pcs)):
            self.get_new_settings()
            if (filename == ""):
                return        
            if (None == selected_process):
                return
                    
            self.update_taint_info()
        else:
            ida_kernwin.msg("hexrays_ida_taint2:  reusing " + filename +
                            " and process " + selected_process['process_name'] +
                            " (" + str(selected_process['process_id']) + ")\n")
        
        if (clear_old):
            self.clear_colors(cfunc)
            
        if (len(tainted_pcs) > 0):
            self.color_eas(cfunc, tainted_pcs)
            ida_kernwin.refresh_idaview_anyway()
        else:
            if (clear_old):
                ida_kernwin.refresh_idaview_anyway()
            ida_kernwin.msg("hexrays_ida_taint2:  no tainted PCs found " +
                            "for selected process\n")
            
        return 1
    
    def init(self):
        if (idaapi.init_hexrays_plugin()):
            # the plugin will terminate right away if return PLUGIN_OK instead
            return ida_idaapi.PLUGIN_KEEP
        else:
            # PLUGIN_SKIP means don't load the plugin
            return ida_idaapi.PLUGIN_SKIP

    def term(self):
        # nothing to do when plugin is unloaded
        pass
        
    def run(self, arg):
        global tainted_pcs
        # this is called when select the plugin from the Edit>Plugins menu
        curwidget = ida_kernwin.get_current_widget()
        if (ida_kernwin.BWN_PSEUDOCODE == ida_kernwin.get_widget_type(curwidget)):
            reuse = HIT2_ReuseDialog.GET_NEW_SETTINGS
            clear_old = False
            if (len(tainted_pcs) > 0):
                reuse = HIT2_ReuseDialog.askToReuse()
                if (HIT2_ReuseDialog.GET_NEW_SETTINGS == reuse):
                    tainted_pcs.clear()
                    # need to clear old colors in case changing process on the
                    # same decompiled function that colored before
                    clear_old = True
            if (HIT2_ReuseDialog.CANCEL_REQUEST != reuse):
                self.color_pseudocode(curwidget, clear_old)
        else:
            ida_kernwin.msg("Current window is not a pseudocode window\n")

def PLUGIN_ENTRY():
    return hexrays_ida_taint2_t()


repository: grayhatacademy/ida
file: plugins/rizzo/rizzo.py
wanted_name: Rizzo
comment: Function signature
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/rizzo/rizzo.py

content:
##########################################################################
# An IDAPython plugin that generates "fuzzy" function signatures that can
# be shared and applied amongst different IDBs.
#
# There are multiple sets of signatures that are generated:
#
#   o "Formal" signatures, where functions must match exactly
#   o "Fuzzy" signatures, where functions must only resemble each other
#     in terms of data/call references.
#   o String-based signatures, where functions are identified based on
#     unique string references.
#   o Immediate-based signatures, where functions are identified based
#     on immediate value references.
#
# These signatures are applied based on accuracy, that is, formal
# signatures are applied first, then string and immediate based
# signatures, and finally fuzzy signatures.
#
# Further, functions are identified based on call references. Consider,
# for example, two functions, one named 'foo', the other named 'bar'.
# The 'foo' function is fairly unique and a reliable signature is easily
# generated for it, but the 'bar' function is more difficult to reliably
# identify. However, 'foo' calls 'bar', and thus once 'foo' is identified,
# 'bar' can also be identified by association.
#
# Craig Heffner
# @devttys0
##########################################################################

from __future__ import print_function
import idc
import idaapi
import idautils

import os
import time
import pickle       # http://natashenka.ca/pickle/
import collections
import hashlib

from shims import ida_shims


class RizzoSignatures(object):
    '''
    Simple wrapper class for storing signature info.
    '''
    SHOW = []

    def __init__(self):
        self.fuzzy = {}
        self.formal = {}
        self.strings = {}
        self.functions = {}
        self.immediates = {}

        self.fuzzydups = set()
        self.formaldups = set()
        self.stringdups = set()
        self.immediatedups = set()

    def show(self):
        if not self.SHOW:
            return

        print("\n\nGENERATED FORMAL SIGNATURES FOR:")
        for (key, ea) in self.formal.items():
            func = RizzoFunctionDescriptor(self.formal, self.functions, key)
            if func.name in self.SHOW:
                print(func.name)

        print("\n\nGENERATED FUZZY SIGNATURES FOR:")
        for (key, ea) in self.fuzzy.items():
            func = RizzoFunctionDescriptor(self.fuzzy, self.functions, key)
            if func.name in self.SHOW:
                print(func.name)


class RizzoStringDescriptor(object):
    '''
    Wrapper class for easily accessing necessary string information.
    '''

    def __init__(self, string):
        self.ea = string.ea
        self.value = str(string)
        self.xrefs = [x.frm for x in idautils.XrefsTo(self.ea)]


class RizzoBlockDescriptor(object):
    '''
    Code block info is stored in tuples, which minimize pickle storage space.
    This class provides more Pythonic (and sane) access to values of interest
    for a given block.
    '''

    def __init__(self, block):
        self.formal = block[0]
        self.fuzzy = block[1]
        self.immediates = block[2]
        self.functions = block[3]

    def match(self, nblock, fuzzy=False):
        # TODO: Fuzzy matching at the block level gets close, but produces a
        #  higher number of false positives; for example, it confuses
        #  hmac_md5 with hmac_sha1.
        return (self.formal == nblock.formal and
                len(self.immediates) == len(nblock.immediates) and
                len(self.functions) == len(nblock.functions))


class RizzoFunctionDescriptor(object):
    '''
    Function signature info is stored in dicts and tuples, which minimize pickle
    storage space. This class provides more Pythonic (and sane) access to
    values of interest for a given function.
    '''

    def __init__(self, signatures, functions, key):
        self.ea = signatures[key]
        self.name = functions[self.ea][0]
        self.blocks = functions[self.ea][1]


class Rizzo(object):
    '''
    Workhorse class which performs the primary logic and functionality.
    '''

    DEFAULT_SIGNATURE_FILE = "rizzo.sig"

    def __init__(self, sigfile=None):
        if sigfile:
            self.sigfile = sigfile
        else:
            self.sigfile = self.DEFAULT_SIGNATURE_FILE

        # Useful for identifying string xrefs from individual instructions
        self.strings = {}
        for string in idautils.Strings():
            self.strings[string.ea] = RizzoStringDescriptor(string)

        start = time.time()
        self.signatures = self.generate()
        end = time.time()

        print("Generated %d formal signatures and %d fuzzy signatures for %d " \
              "functions in %.2f seconds." % (len(self.signatures.formal),
                                              len(self.signatures.fuzzy),
                                              len(self.signatures.functions),
                                              (end-start)))

    def save(self):
        print(("Saving signatures to %s..." % self.sigfile), end=' ')
        fp = open(self.sigfile, "wb")
        pickle.dump(self.signatures, fp)
        fp.close()
        print("done.")

    def load(self):
        print(("Loading signatures from %s..." % self.sigfile), end=' ')
        fp = open(self.sigfile, "rb")
        sigs = pickle.load(fp)
        fp.close()
        print("done.")
        return sigs

    def sighash(self, value):
        h = hashlib.md5()
        h.update(value.encode("utf-8"))
        hash_value = h.hexdigest()
        del(h)
        return hash_value

    def block(self, block):
        '''
        Returns a tuple:
        ([formal, block, signatures], [fuzzy, block, signatures],
        set([unique, immediate, values]), [called, function, names])
        '''
        formal = []
        fuzzy = []
        functions = []
        immediates = []

        ea = ida_shims.start_ea(block)
        while ea < ida_shims.end_ea(block):
            insn = ida_shims.decode_insn(ea)

            # Get a list of all data/code refs from the current instruction
            drefs = [x for x in idautils.DataRefsFrom(ea)]
            crefs = [x for x in idautils.CodeRefsFrom(ea, False)]

            # Add all instruction mnemonics to the formal block hash
            formal.append(ida_shims.print_insn_mnem(ea))

            # If this is a call instruction, be sure to note the name of the
            # function being called. This is used to apply call-based
            # signatures to functions.
            #
            # For fuzzy signatures, we can't use the actual name or EA of the
            # function, but rather just want to note that a function call was
            # made.
            #
            # Formal signatures already have the call instruction mnemonic,
            # which is more specific than just saying that a call was made.
            if idaapi.is_call_insn(ea):
                for cref in crefs:
                    func_name = ida_shims.get_name(cref)
                    if func_name:
                        functions.append(func_name)
                        fuzzy.append("funcref")
            # If there are data references from the instruction, check to see
            # if any of them are strings. These are looked up in the
            # pre-generated strings dictionary.
            #
            # String values are easily identifiable, and are used as part of
            # both the fuzzy and the formal signatures.
            #
            # It is more difficult to determine if non-string values are
            # constants or not; for both fuzzy and formal signatures, just use
            # "data" to indicate that some data was referenced.
            elif drefs:
                for dref in drefs:
                    if dref in self.strings:
                        formal.append(self.strings[dref].value)
                        fuzzy.append(self.strings[dref].value)
                    else:
                        formal.append("dataref")
                        fuzzy.append("dataref")
            # If there are no data or code references from the instruction, use
            # every operand as part of the formal signature.
            #
            # Fuzzy signatures are only concerned with interesting immediate
            # values, that is, values that are greater than 65,535, are not
            # memory addresses, and are not displayed as negative values.
            elif not drefs and not crefs:
                ops = ida_shims.get_operands(insn)
                for n in range(0, len(ops)):
                    opnd_text = ida_shims.print_operand(ea, n)
                    formal.append(opnd_text)
                    if ops[n].type == idaapi.o_imm and \
                            not opnd_text.startswith('-'):
                        if ops[n].value >= 0xFFFF:
                            if ida_shims.get_full_flags(ops[n].value) == 0:
                                fuzzy.append(str(ops[n].value))
                                immediates.append(ops[n].value)

            ea = ida_shims.next_head(ea)

        return (self.sighash(''.join(formal)),
                self.sighash(''.join(fuzzy)),
                immediates, functions)

    def function(self, func):
        '''
        Returns a list of blocks.
        '''
        blocks = []

        for block in idaapi.FlowChart(func):
            blocks.append(self.block(block))

        return blocks

    def generate(self):
        signatures = RizzoSignatures()

        # Generate unique string-based function signatures
        for (ea, string) in self.strings.items():
            # Only generate signatures on reasonably long strings with one xref
            if len(string.value) >= 8 and len(string.xrefs) == 1:
                func = idaapi.get_func(string.xrefs[0])
                if func:
                    str_hash = self.sighash(string.value)

                    # Check for and remove string duplicate signatures (the same
                    # string can appear more than once in an IDB).
                    # If no duplicates, add this to the string signature dict.
                    if str_hash in signatures.strings:
                        del signatures.strings[str_hash]
                        signatures.stringdups.add(str_hash)
                    elif str_hash not in signatures.stringdups:
                        signatures.strings[str_hash] = ida_shims.start_ea(func)

        # Generate formal, fuzzy, and immediate-based function signatures
        for ea in idautils.Functions():
            func = idaapi.get_func(ea)
            if func:
                # Generate a signature for each block in this function
                blocks = self.function(func)

                # Build function-wide formal and fuzzy signatures by simply
                # concatenating the individual function block signatures.
                formal = self.sighash(''.join([str(e)
                                               for (e, f, i, c) in blocks]))
                fuzzy = self.sighash(''.join([str(f)
                                              for (e, f, i, c) in blocks]))

                # Add this signature to the function dictionary.
                start_ea = ida_shims.start_ea(func)
                signatures.functions[start_ea] = (ida_shims.get_name(start_ea),
                                                  blocks)

                # Check for and remove formal duplicate signatures.
                # If no duplicates, add this to the formal signature dict.
                if formal in signatures.formal:
                    del signatures.formal[formal]
                    signatures.formaldups.add(formal)
                elif formal not in signatures.formaldups:
                    signatures.formal[formal] = ida_shims.start_ea(func)

                # Check for and remove fuzzy duplicate signatures.
                # If no duplicates, add this to the fuzzy signature dict.
                if fuzzy in signatures.fuzzy:
                    del signatures.fuzzy[fuzzy]
                    signatures.fuzzydups.add(fuzzy)
                elif fuzzy not in signatures.fuzzydups:
                    signatures.fuzzy[fuzzy] = ida_shims.start_ea(func)

                # Check for and remove immediate duplicate signatures.
                # If no duplicates, add this to the immediate signature dict.
                for (e, f, immediates, c) in blocks:
                    for immediate in immediates:
                        if immediate in signatures.immediates:
                            del signatures.immediates[immediate]
                            signatures.immediatedups.add(immediate)
                        elif immediate not in signatures.immediatedups:
                            signatures.immediates[immediate] = \
                                ida_shims.start_ea(func)

        # These need not be maintained across function calls,
        # and only add to the size of the saved signature file.
        signatures.fuzzydups = set()
        signatures.formaldups = set()
        signatures.stringdups = set()
        signatures.immediatedups = set()

        # DEBUG
        signatures.show()

        return signatures

    def match(self, extsigs):
        fuzzy = {}
        formal = {}
        strings = {}
        immediates = {}

        # Match formal function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.formal.items():
            if extsig in self.signatures.formal:
                new_fun = RizzoFunctionDescriptor(
                    extsigs.formal, extsigs.functions, extsig)
                curr_fun = RizzoFunctionDescriptor(
                    self.signatures.formal, self.signatures.functions, extsig)
                formal[curr_fun] = new_fun
        end = time.time()
        print("Found %d formal matches in %.2f seconds." % (len(formal),
                                                            (end-start)))

        # Match fuzzy function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.fuzzy.items():
            if extsig in self.signatures.fuzzy:
                curr_fun = RizzoFunctionDescriptor(
                    self.signatures.fuzzy, self.signatures.functions, extsig)
                new_fun = RizzoFunctionDescriptor(
                    extsigs.fuzzy, extsigs.functions, extsig)
                # Only accept this as a valid match if the functions have the
                # same number of basic code blocks
                if len(curr_fun.blocks) == len(new_fun.blocks):
                    fuzzy[curr_fun] = new_fun
        end = time.time()
        print("Found %d fuzzy matches in %.2f seconds." % (len(fuzzy),
                                                           (end-start)))

        # Match string based function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.strings.items():
            if extsig in self.signatures.strings:
                curr_fun = RizzoFunctionDescriptor(
                    self.signatures.strings, self.signatures.functions, extsig)
                new_fun = RizzoFunctionDescriptor(
                    extsigs.strings, extsigs.functions, extsig)
                strings[curr_fun] = new_fun
        end = time.time()
        print("Found %d string matches in %.2f seconds." % (len(strings),
                                                            (end-start)))

        # Match immediate baesd function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.immediates.items():
            if extsig in self.signatures.immediates:
                curr_fun = RizzoFunctionDescriptor(
                    self.signatures.immediates, self.signatures.functions,
                    extsig)
                new_fun = RizzoFunctionDescriptor(
                    extsigs.immediates, extsigs.functions, extsig)
                immediates[curr_fun] = new_fun
        end = time.time()
        print("Found %d immediate matches in %.2f seconds." % (len(immediates),
                                                               (end-start)))

        # Return signature matches in the order we want them applied
        # The second tuple of each match is set to True if it is a fuzzy match,
        # e.g.:   ((match, fuzzy), (match, fuzzy), ...)
        return ((formal, False),
                (strings, False),
                (immediates, False),
                (fuzzy, True))

    def rename(self, ea, name):
        # Don't rely on the name in curfunc, it could have already been renamed
        curname = ida_shims.get_name(ea)
        # Don't rename if the name is a special identifier, or if the ea has
        # already been named
        if curname.startswith('sub_') and \
                name.split('_')[0] not in \
                ['sub', 'loc', 'unk', 'dword', 'word', 'byte']:
            # Don't rename if the name already exists in the IDB
            if ida_shims.get_name_ea_simple(name) == idc.BADADDR:
                if ida_shims.set_name(ea, name):
                    ida_shims.set_func_flags(
                        ea, (ida_shims.get_func_flags(ea) | idc.FUNC_LIB))
                    return 1
        return 0

    def apply(self, extsigs):
        count = 0

        start = time.time()

        # This applies formal matches first, then fuzzy matches
        for (match, fuzzy) in self.match(extsigs):
            # Keeps track of all function names that we've identified candidate
            # functions for
            rename = {}

            for (curfunc, newfunc) in match.items():
                if newfunc not in rename:
                    rename[newfunc.name] = []

                # Attempt to rename this function
                rename[newfunc.name].append(curfunc.ea)

                bm = {}
                duplicates = set()

                # Search for unique matching code blocks inside this function
                for nblock in newfunc.blocks:
                    nblock = RizzoBlockDescriptor(nblock)
                    for cblock in curfunc.blocks:
                        cblock = RizzoBlockDescriptor(cblock)

                        if cblock.match(nblock, fuzzy):
                            if cblock in bm:
                                del bm[cblock]
                                duplicates.add(cblock)
                            elif cblock not in duplicates:
                                bm[cblock] = nblock

                # Rename known function calls from each unique code block
                for (cblock, nblock) in bm.items():
                    for n in range(0, len(cblock.functions)):
                        ea = ida_shims.get_name_ea_simple(cblock.functions[n])
                        if ea != idc.BADADDR:
                            if nblock.functions[n] in rename:
                                rename[nblock.functions[n]].append(ea)
                            else:
                                rename[nblock.functions[n]] = [ea]

                # Rename the identified functions
                for (name, candidates) in rename.items():
                    if candidates:
                        winner = \
                            collections.Counter(candidates).most_common(1)[0][0]
                        count += self.rename(winner, name)

        end = time.time()
        print("Renamed %d functions in %.2f seconds." % (count, (end-start)))


def RizzoBuild(sigfile=None):
    print("Building Rizzo signatures, this may take a few minutes...")
    start = time.time()
    r = Rizzo(sigfile)
    r.save()
    end = time.time()
    print("Built signatures in %.2f seconds" % (end-start))


def RizzoApply(sigfile=None):
    print("Applying Rizzo signatures, this may take a few minutes...")
    start = time.time()
    r = Rizzo(sigfile)
    s = r.load()
    r.apply(s)
    end = time.time()
    print("Signatures applied in %.2f seconds" % (end-start))


def rizzo_produce(arg=None):
    fname = ida_shims.ask_file(1, "*.riz", "Save signature file as")
    if fname:
        if '.' not in fname:
            fname += ".riz"
        RizzoBuild(fname)


def rizzo_load(arg=None):
    fname = ida_shims.ask_file(0, "*.riz", "Load signature file")
    if fname:
        RizzoApply(fname)

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

        def activate(self, ctx):
            rizzo_produce()
            return 1

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


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

        def activate(self, ctx):
            rizzo_load()
            return 1

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


class RizzoPlugin(idaapi.plugin_t):
    flags = 0
    comment = "Function signature"
    help = ""
    wanted_name = "Rizzo"
    wanted_hotkey = ""
    produce_action_name = 'producerizzo:action'
    load_action_name = 'loadrizzo:action'
    menu_name = "Rizzo signature file..."
    produce_tooltip = "Produce rizzo signature file."
    load_tooltip = "Load rizzo signature file."
    menu_tab = 'File/'
    menu_context = []

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            produce_desc = idaapi.action_desc_t(self.produce_action_name,
                                                self.menu_name,
                                                ProduceRizzoAction(),
                                                self.wanted_hotkey,
                                                self.produce_tooltip,
                                                199)

            load_desc = idaapi.action_desc_t(self.load_action_name,
                                             self.menu_name,
                                             LoadRizzoAction(),
                                             self.wanted_hotkey,
                                             self.load_tooltip,
                                             199)

            idaapi.register_action(produce_desc)
            idaapi.register_action(load_desc)

            idaapi.attach_action_to_menu(
                os.path.join(self.menu_tab, 'Produce file/'),
                self.produce_action_name,
                idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu(
                os.path.join(self.menu_tab, 'Load file/'),
                self.load_action_name,
                idaapi.SETMENU_APP)
        else:
            self.menu_context.append(
                idaapi.add_menu_item(
                    os.path.join(self.menu_tab, 'Load file/'),
                    "Rizzo signature file...", "", 0, rizzo_load, (None,)))

            self.menu_context.append(
                idaapi.add_menu_item(
                    os.path.join(self.menu_tab, 'Produce file/'),
                    "Rizzo signature file...", "", 0, rizzo_produce, (None,)))

        return idaapi.PLUGIN_KEEP

    def term(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.detach_action_from_menu(
                self.menu_tab, self.produce_action_name)
            idaapi.detach_action_from_menu(
                self.menu_tab, self.load_action_name)
        else:
            if self.menu_context is not None:
                idaapi.del_menu_item(self.menu_context)
        return None

    def run(self, arg):
        return None

    def rizzo_script(self):
        idaapi.IDAPython_ExecScript(self.script, globals())


def PLUGIN_ENTRY():
    return RizzoPlugin()


repository: gaasedelen/tenet
file: plugins/tenet/integration/ida_loader.py
description: A Trace Explorer for Reverse Engineers
wanted_name: Tenet
comment: Trace Explorer
created_at: 2021-04-20 13:56:48+00:00
pushed_at: 2023-10-23 13:53:22+00:00
url:
https://github.com/gaasedelen/tenet/blob/62c0bec073fcbb31c5be52c9ea236b6ea2bcc418/plugins/tenet/integration/ida_loader.py

content:
import time
import logging

import ida_idaapi
import ida_kernwin

from tenet.util.log import pmsg
from tenet.integration.ida_integration import TenetIDA

logger = logging.getLogger("Tenet.IDA.Loader")

#------------------------------------------------------------------------------
# IDA Plugin Loader
#------------------------------------------------------------------------------
#
#    This file contains a stub 'plugin' class for the plugin as required by
#    IDA Pro. Practically speaking, there should be little to *no* logic placed
#    in this file because it is disassembler-specific.
#
#    When IDA Pro is starting up, it will import all python files placed in its
#    root plugin folder. It will then attempt to call PLUGIN_ENTRY() on each of
#    the imported 'plugins'. We import PLUGIN_ENTRY into tenet_plugin.py
#    so that IDA can see it.
#
#    PLUGIN_ENTRY() is expected to return a plugin object (TenetIDAPlugin)
#    derived from ida_idaapi.plugin_t. IDA will register the plugin, and
#    interface with the plugin object to load / unload the plugin at certain
#    times, per its configuration (flags, hotkeys).
#
#    There should be virtually no reason for you to modify this file.
#

def PLUGIN_ENTRY():
    """
    Required plugin entry point for IDAPython Plugins.
    """
    return TenetIDAPlugin()

class TenetIDAPlugin(ida_idaapi.plugin_t):
    """
    The IDA plugin stub for Tenet.
    """

    #
    # Plugin flags:
    # - PLUGIN_MOD:  The plugin may modify the database
    # - PLUGIN_PROC: Load/unload the plugin when an IDB opens / closes
    # - PLUGIN_HIDE: Hide the plugin from the IDA plugin menu
    #

    flags = ida_idaapi.PLUGIN_PROC | ida_idaapi.PLUGIN_MOD | ida_idaapi.PLUGIN_HIDE
    comment = "Trace Explorer"
    help = ""
    wanted_name = "Tenet"
    wanted_hotkey = ""

    #--------------------------------------------------------------------------
    # IDA Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """

        try:
            self.core = TenetIDA()
            self.core.load()
        except Exception as e:
            pmsg("Failed to initialize Tenet")
            logger.exception("Exception details:")

        #
        # we return PLUGIN_KEEP here regardless of success/failure. this is to
        # ensure that IDA will not try to reload the plugin again.
        #

        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        ida_kernwin.warning("Tenet cannot be run as a script in IDA.")

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """
        logger.debug("IDA term started...")

        start = time.time()
        logger.debug("-"*50)

        try:
            self.core.unload()
            self.core = None
        except Exception as e:
            logger.exception("Failed to cleanly unload Tenet from IDA.")

        end = time.time()
        logger.debug("-"*50)

        logger.debug("IDA term done... (%.3f seconds...)" % (end-start))


repository: BinaryAnalysisPlatform/bap-ida-python
file: plugins/bap/plugins/bap_comments.py
description: integration with IDA
wanted_name: BAP: View BAP Attributes
created_at: 2016-06-10 18:00:33+00:00
pushed_at: 2021-05-27 15:13:33+00:00
url:
https://github.com/BinaryAnalysisPlatform/bap-ida-python/blob/d8d4679de2f50bb75f556419565821d95404034e/plugins/bap/plugins/bap_comments.py

content:
import idaapi
import idc

from bap.utils import bap_comment, ida


class Attributes(idaapi.Choose2):
    def __init__(self, comms):
        idaapi.Choose2.__init__(self, 'Select an attribute', [
            ['name', 8],
            ['addr', 8],
            ['data', 64]
        ])
        self.comms = [
            [name, '{:#x}'.format(addr), ' '.join(data)]
            for (name, addr, data) in comms
        ]

    def OnClose(self):
        pass

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

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


class BapComment(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    help = 'propagate comments to IDA Views'
    comment = ''
    wanted_name = 'BAP: View BAP Attributes'
    wanted_hotkey = 'Shift-B'

    def __init__(self):
        self.comms = {}

    def init(self):
        ida.comment.register_handler(self.update)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        comms = {}
        for addr in ida.addresses():
            comm = idaapi.get_cmt(addr, 0)
            if comm:
                try:
                    parsed = bap_comment.parse(comm)
                    if parsed:
                        for (name, data) in parsed.items():
                            comms[(addr, name)] = data
                except:
                    idc.Message("BAP> failed to parse string {0}\n{1}".
                                format(comm, str(sys.exc_info()[1])))
        comms = [(name, addr, data)
                 for ((addr, name), data) in comms.items()]
        attrs = Attributes(comms)
        choice = attrs.Show(modal=True)
        if choice >= 0:
            idc.Jump(comms[choice][1])

    def term(self):
        pass

    def update(self, ea, key, value):
        """Add key=values to comm string at EA."""
        cmt = idaapi.get_cmt(ea, 0)
        comm = cmt and bap_comment.parse(cmt) or {}
        values = comm.setdefault(key, [])
        if value and value != '()' and value not in values:
            values.append(value)
        idaapi.set_cmt(ea, bap_comment.dumps(comm), 0)


def PLUGIN_ENTRY():
    return BapComment()


repository: Kvancy/ComprehendAI
file: ComprehendAI.py
wanted_name: ComprehendAI
comment: AI-based Reverse Analysis Plugin
created_at: 2025-03-04 13:18:45+00:00
pushed_at: 2025-04-16 13:15:42+00:00
description:
An AI plugin for assisting IDA reverse analysis, which facilitates quickly summarizing the functions of code and accelerates the analysis efficiency.

url:
https://github.com/Kvancy/ComprehendAI/blob/c5ea78935508dc8e0e619b78bef3902b8e9022ce/ComprehendAI.py

content:
import traceback
import idaapi
import idc
import idautils
import ida_xref
import json
import os
import asyncio

from idaapi import action_handler_t, UI_Hooks
from threading import Lock,Thread,Event
from openai import OpenAI
from enum import Enum

class TaskType(Enum):
    ANALYSIS = 1
    CUSTOM_QUERY = 2
    CUSTON_QUERY_WITH_CODE = 3

class QueryStatus(Enum):
    SUCCESS = 1
    FAILED = 2
    STOPPED = 3

#处理配置文件
class ConfigManager:
    _instance = None
    _lock = Lock()
    
    def __new__(cls):
        with cls._lock:
            if not cls._instance:
                cls._instance = super().__new__(cls)
                cls._instance._initialize()
            return cls._instance
    
    def _initialize(self):
        self.script_dir = os.path.dirname(os.path.abspath(__file__))
        self.config_path = os.path.join(self.script_dir, 'config.json')
        self.config = self._load_config()
        self.openai_client = self._create_openai_client()
        
    def _load_config(self):
        try:
            with open(self.config_path, "r") as f:
                return json.load(f)
        except Exception as e:
            raise RuntimeError(f"Failed to load config: {str(e)}")
    
    def _create_openai_client(self):
        return OpenAI(
            api_key=self.config["openai"]["api_key"],
            base_url=self.config["openai"]["base_url"]
        )
    
    @property
    def model_name(self):
        return self.config["openai"]["model"]
    
    @property
    def client(self):
        return self.openai_client
#处理反汇编代码提取
class DisassemblyProcessor:
    def __init__(self, max_depth=2):
        self.max_depth = max_depth
        self._lock = Lock()
        self._reset_state()
        
    def _reset_state(self):
        with self._lock:
            self.processed_funcs = set()
            self.func_disasm_list = []
    
    def get_current_function_disasm(self):
        self._reset_state()
        
        current_ea = idc.get_screen_ea()
        func_start = idc.get_func_attr(current_ea, idc.FUNCATTR_START)
        
        if func_start == idaapi.BADADDR:
            raise ValueError("Failed to locate function start address")
            
        self._process_function(func_start, self.max_depth)
        return "\n".join(self.func_disasm_list)
    
    def _process_function(self, func_ea, depth):
        if func_ea in self.processed_funcs or depth < 0:
            return
            
        with self._lock:
            self.processed_funcs.add(func_ea)
        
        try:
            decompiled = str(idaapi.decompile(func_ea))
            with self._lock:
                self.func_disasm_list.append(decompiled)
        except Exception as e:
            print(f"Decompilation failed for {hex(func_ea)}: {str(e)}")
        
        for callee in self._get_callees(func_ea):
            self._process_function(callee, depth - 1)
    
    def _get_callees(self, func_ea):
        callees = set()
        for ea in range(func_ea, idc.get_func_attr(func_ea, idc.FUNCATTR_END)):
            for xref in idautils.XrefsFrom(ea):
                if xref.type in [ida_xref.fl_CN, ida_xref.fl_CF]:
                    callee_ea = xref.to
                    if idc.get_func_attr(callee_ea, idc.FUNCATTR_START) == callee_ea:
                        callees.add(callee_ea)
        return callees
#处理openai接口相关
class AIService:
    def __init__(self):
        self.config = ConfigManager()
        self.stop_event = Event()

    def ask_ai(self, prompt, ai_isRunning:Lock):
        messages = [{"role": "user", "content": prompt}]
        print("ComprehendAI output:")
        self.stop_event.clear() #初始化事件
            
        result = self._request_openai(messages)
        ai_isRunning.release()  # 分析完成，无论成功失败都需释放锁

        match result:
            case QueryStatus.SUCCESS:
                print("\r✅ 分析完成！")
            case QueryStatus.FAILED:
                print("\r❌ 分析失败，请重试")
            case QueryStatus.STOPPED:
                print("\r✅ 分析暂停")
                

    def _request_openai(self,messages):
        reasoning_content = ""
        answer_content = ""
        is_answering = False
        try:
            completion = self.config.client.chat.completions.create(
            model=self.config.model_name,
            messages=messages,
            stream=True,
            )
            for chunk in completion:
                if self.stop_event.is_set():
                    raise StopIteration("任务被停止")

                # 如果chunk.choices为空，则打印usage
                if not chunk.choices:
                    print("\nUsage:")
                    print(chunk.usage)
                else:
                    delta = chunk.choices[0].delta
                    # 打印思考过程
                    if hasattr(delta, 'reasoning_content') and delta.reasoning_content != None:
                        print(delta.reasoning_content, end='', flush=True)
                        reasoning_content += delta.reasoning_content
                    else:
                        # 开始回复
                        if delta.content != "" and is_answering is False:
                            print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n")
                            is_answering = True
                        answer_content += delta.content
                        
            print(answer_content)        
            return QueryStatus.SUCCESS
        
        except StopIteration as e:
            print(f"Error occurred: {e}")
            return QueryStatus.STOPPED

        except Exception as e:
            print(f"Error occurred: {e}")
            traceback.print_exc()
            return QueryStatus.FAILED
        
#处理用户接口
class AnalysisHandler:

    def __init__(self):
        self.disassembler = DisassemblyProcessor()
        self.ai_service = AIService()
        self.ai_isRunning = Lock()
        self.prompt = """
你是一名人工智能逆向工程专家。
我会提供你一些反汇编代码，其中首个函数是你需要分析并总结成报告的函数，
其余函数是该函数调用的一些子函数。
分析要求：
重点描述主函数功能，并对核心行为进行推测；
简要描述子函数功能

输出要求：
主函数功能：...
行为推测：...
子函数功能：...
纯文本输出。

下面是你要分析的反汇编代码：
"""
    def set_analysis_depth(self, depth):
        self.disassembler.max_depth = depth
    
    def _create_analysis_prompt(self, disassembly):
        return self.prompt + disassembly
    
    def _create_analysis_custom_query(self, disassembly,question):
        return question + disassembly
    
    def create_ai_task(self,taskType,question=""):

        match taskType:
            case TaskType.ANALYSIS:
                disassembly = self.disassembler.get_current_function_disasm()
                prompt = self._create_analysis_prompt(disassembly)
                self.async_task(prompt)
            case TaskType.CUSTOM_QUERY:
                self.async_task(question)    
            case TaskType.CUSTON_QUERY_WITH_CODE:
                disassembly = self.disassembler.get_current_function_disasm()
                prompt = self._create_analysis_custom_query(disassembly,question)
                self.async_task(prompt)
        
        
    def async_task(self,question):
        print(question)
        if self.ai_isRunning.acquire(blocking=False):
            task = Thread(target=self.ai_service.ask_ai,args=(question,self.ai_isRunning,)) 
            task.start()
            
        else:
            print("\r❌ 当前AI正在处理任务,请稍后尝试")
    
    def stop(self):
        self.ai_service.stop_event.set()

#处理插件框架
class ComprehendAIPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = "AI-based Reverse Analysis Plugin"
    help = "Perform AI-based analysis on binary code"
    wanted_name = "ComprehendAI"
    wanted_hotkey = "Ctrl+Shift+A"

    ACTION_DEFINITIONS = [
        ("AI_analysis:Analysis", "Analysis", "执行非阻塞型AI分析"),
        ("AI_analysis:SetDepth", "Set analysis depth", "设置分析深度"),
        ("AI_analysis:SetPrompt", "Set your own prompt", "自定义prompt"),
        ("AI_analysis:CustomQueryWithCode", "Ask AI with code", "结合代码自定义提问"),
        ("AI_analysis:CustomQuery", "Ask AI", "自定义提问"),
        ("AI_analysis:Stop", "Stop", "停止"),
    ]

    def init(self):
        self.ui_hook = self.MenuHook()
        self.ui_hook.hook()
        
        self.handler = AnalysisHandler()
        self._register_actions()
        
        print("ComprehendAI initialized")
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        self.ui_hook.unhook()
        self._unregister_actions()
        print("ComprehendAI unloaded")

    def _register_actions(self):
        for action_id, label, tooltip in self.ACTION_DEFINITIONS:
            action_desc = idaapi.action_desc_t(
                action_id,
                label,
                self.MenuCommandHandler(action_id,self.handler),
                None,
                tooltip,
                0
            )
            idaapi.register_action(action_desc)

    def _unregister_actions(self):
        for action_id, _, _ in self.ACTION_DEFINITIONS:
            idaapi.unregister_action(action_id)

    class MenuHook(UI_Hooks):
        def finish_populating_widget_popup(self, form, popup):
            if idaapi.get_widget_type(form) in (idaapi.BWN_DISASM, idaapi.BWN_PSEUDOCODE):
                for action_id, _, _ in ComprehendAIPlugin.ACTION_DEFINITIONS:
                    idaapi.attach_action_to_popup(form, popup, action_id, "ComprehendAI/", idaapi.SETMENU_APP)

    class MenuCommandHandler(action_handler_t):
        def __init__(self, action_id,handler:AnalysisHandler):
            super().__init__()
            self.action_id = action_id
            self.handler = handler
    
        def activate(self, ctx):     

            match self.action_id:
                case "AI_analysis:Analysis":
                    self.handler.create_ai_task(TaskType.ANALYSIS)
                case "AI_analysis:CustomQuery":
                    question = idaapi.ask_text(0, "", "输入问题")
                    if question:
                        self.handler.create_ai_task(TaskType.CUSTOM_QUERY, question)
                case "AI_analysis:SetDepth":
                    new_depth = idaapi.ask_long(2, "设置分析深度 (默认2):")
                    if new_depth is not None:
                        self.handler.set_analysis_depth(new_depth)       
                case "AI_analysis:SetPrompt":
                    yourPrompt = idaapi.ask_text(0, f"{self.handler.prompt}", "输入问题")
                    self.handler.prompt = yourPrompt
                case "AI_analysis:CustomQueryWithCode":
                    question = idaapi.ask_text(0, "", "输入问题")
                    if question:
                        self.handler.create_ai_task(TaskType.CUSTON_QUERY_WITH_CODE, question)
                case "AI_analysis:Stop":
                    self.handler.stop()
            return 1

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

def PLUGIN_ENTRY():
    return ComprehendAIPlugin()


repository: thalium/ida_kmdf
file: wdf_plugin/wdf_plugin.py
wanted_name: Wdfalyzer
comment: Detect wdf version and apply associated wdf types
created_at: 2023-05-09 20:29:02+00:00
pushed_at: 2025-03-02 12:33:07+00:00
url:
https://github.com/thalium/ida_kmdf/blob/01eb81aad4bf9df17294305b611fb9abacc2179e/wdf_plugin/wdf_plugin.py

content:
import idaapi
import idc
import wdf as wdf


class WdfPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "Detect wdf version and apply associated wdf types"
    wanted_name = "Wdfalyzer"
    wanted_hotkey = "Alt-F11"

    def init(self):
        self.run()

    def run(self):
        if wdf.load_wdf():
            wdf.apply_wdf()
            idaapi.auto_wait()

    def term(self):
        pass


def PLUGIN_ENTRY():
    return WdfPlugin()


repository: SouhailHammou/IDARay-Plugin
file: idaray.py
wanted_name: IDARay Plugin
comment: Scans the database with multiple YARA files
created_at: 2018-10-26 10:14:33+00:00
pushed_at: 2018-11-16 11:25:21+00:00
description:
IDARay is an IDA Pro plugin that matches the database against multiple YARA files which themselves may contain multiple rules.

url:
https://github.com/SouhailHammou/IDARay-Plugin/blob/e471852be4c54b222293fdb64d06ad3cdcedd5cf/idaray.py

content:
'''
IDARay - A simple IDA plugin that matches the database against multiple YARA files.
Author : Souhail Hammou

Plugin tested under IDA Pro 6.8
Feel free to copy or edit the plugin.
'''
from idc import *
from idaapi import *
import yara
import os

class YARAScan() :
	def __init__(self, directory) :
		self.directory = directory
		self.rules_dict = self.compile_all_rules()
		self.matches = {}
	
	def compile_all_rules(self) :
		rules = {}
		try :
			for filename in os.listdir(self.directory) :
				try :
					#We can't just use filepaths with yara.compile because we're also interested
					#in the filename. As a result, we'll match file by file.
					rules.update( { filename : yara.compile(os.path.join(self.directory, filename)) })
				except :
					print "IDARay : Compilation Error ! Please make sure your directory only contains YARA files."
		except :
			print "IDARay : Error while accessing directory."
		return rules
		
	def scan(self) :
		if not self.rules_dict :
			return 0
		
		#Join all segments
		data = ''
		prev_seg = 0
		for seg in Segments() :
			if prev_seg and SegEnd(prev_seg) < seg :
				data += chr(0xFF) * (seg - SegEnd(prev_seg))
			for ea in range(seg, SegEnd(seg)):
					data += chr(Byte(ea))
			prev_seg = seg
		
		#Match
		for entry in self.rules_dict :
			rule = self.rules_dict[entry]
			filename = entry
			m = rule.match(data=data)
			
			if m :
				self.matches.update( { filename : m } )
		
		#Display the results
		FilesChooserView(self.matches)
		return 1

class SingleDetailsView(idaapi.Choose2) :
	def __init__(self, rule) :
		idaapi.Choose2.__init__(self,
		"Matched strings for rule : %s" % rule.rule,
		[ ["Offset (double-click to follow)", 20 | Choose2.CHCOL_HEX], ["Matched Strings/Patterns", 50 | Choose2.CHCOL_PLAIN] ],
		Choose2.CH_MODAL)
		
		self.strings = rule.strings
		self.Show()
	
	def OnClose(self):
		return

	def OnGetLine(self, n):
		matched = self.strings[n][2]
		
		#Some logic to distinguish strings from patterns
		#is this a pattern and not a string ?
		if not all( ord(c) < 0x7F and ord(c) > 0x1F for c in matched ) :
			bytes = matched
			matched = '{'
			for b in bytes : 
				matched += ' %02x'.upper() % ord(b)
			matched += ' }'
		return [ str(hex(self.strings[n][0])).strip('L').upper().replace('X','x'), matched ]

	def OnGetSize(self):
		return len(self.strings)
		
	def OnSelectLine(self, n) :
		idc.Jump(FirstSeg() + self.strings[n][0])

		
class MultipleDetailsView(idaapi.Choose2) :
	def __init__(self, filename ,yara_matches) :
		idaapi.Choose2.__init__(self,
		"Matches for %s" % filename,
		[ ["Matched Rules (double-click for details)", 30 | Choose2.CHCOL_PLAIN]],
		Choose2.CH_MODAL)
		
		self.yara_matches = yara_matches
		self.Show()
	
	def OnClose(self):
		return

	def OnGetLine(self, n):
		match = self.yara_matches[n]
		line = [match.rule]
		return line

	def OnGetSize(self):
		return len(self.yara_matches)
	
	def OnSelectLine(self, n) :
		#Go to the single details view
		SingleDetailsView(self.yara_matches[n])
	
		
class FilesChooserView(idaapi.Choose2) :
	def __init__(self, matches) :
		idaapi.Choose2.__init__(self,
		"Matched YARA Files",
		[ ["Filename (Double-click for details)", 30 | Choose2.CHCOL_PLAIN], ["Number of matched rules", 30 | Choose2.CHCOL_PLAIN] ],
		Choose2.CH_MODAL)
		
		self.matches = matches
		self.Show()
			
	def OnClose(self):
		return

	def OnGetLine(self, n):
		filename = key = list(self.matches)[n]
		
		#Filename || Number of matched rules from the file
		return [filename, str(len(self.matches[key]))]

	def OnGetSize(self):
		return len(self.matches)
	
	def OnSelectLine(self, n) :
		filename = key = list(self.matches)[n]
		
		if len(self.matches[key]) > 1 :
			#If rules within multiple files match, display this view
			MultipleDetailsView(filename, self.matches[key])
		else :
			#If the matched rules are within a single file, skip the file selection view
			SingleDetailsView(self.matches[key][0])
		return 1
		
class ChooseFiles(idaapi.Form) :
	def __init__(self) :
		Form.__init__(self, 
		r"""Choose a directory containing YARA files

		Welcome to IDARay !
		IDARay matches the database against multiple 
		YARA files which themselves may contain multiple rules.
		
		Please select a directory containing YARA files only.
		YARA Directory		<:{dir}>
		""",
		{
		'dir' : Form.DirInput(swidth=50),
		})
		
		self.Compile()

class idaray_handler(idaapi.action_handler_t):
	#http://www.hexblog.com/?p=886
	def __init__(self) :
		idaapi.action_handler_t.__init__(self)
	
	def activate(self, ctx) :
		#Open the choose file dialog
		form = ChooseFiles()
		res = form.Execute()
		
		if res == 1 :
			if form.dir.value :
				directory = form.dir.value
				yara_scan = YARAScan(directory)
				yara_scan.scan()
			else :
				print "IDARay : Please choose a directory !"
		return 1
		
	def update(self, ctx) :
		return idaapi.AST_ENABLE_ALWAYS	

		
class idaray_plugin_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_UNL
	comment = "Scans the database with multiple YARA files"
	wanted_name = "IDARay Plugin"
	help = "help"
	wanted_hotkey = ""
	
	def init(self):
		global idaray
		
		if 'idaray' not in globals() :
			idaray = 1
			self.initialize_menu()
		
		return idaapi.PLUGIN_KEEP

	def run(self, arg):
		idaray_handler().activate(None)

	def term(self):
		pass

	def initialize_menu(self):
		print '''------------------------------------------------------
IDARay plugin - Souhail Hammou (2018)
------------------------------------------------------'''
		
def PLUGIN_ENTRY():
    return idaray_plugin_t()


repository: FriedAppleTeam/FRAPL
file: Framework/FridaLink/FridaLink.py
description: FRAPL Framework
wanted_name: Frida Link
created_at: 2016-10-24 10:30:03+00:00
pushed_at: 2016-12-26 13:17:42+00:00
url:
https://github.com/FriedAppleTeam/FRAPL/blob/89c14d57e0cc77b915fe1e95f60e9e1847699103/Framework/FridaLink/FridaLink.py

content:
#
#  FridaLink.py
#  FridaLink IDA  
#
#  Created by Alexander Hude on 16/11/15.
#  Copyright (c) 2015 FriedApple. All rights reserved.
#

import SocketServer
import threading
import pickle
import select
import socket
import json

# IDA Python SDK
from idaapi import *
from idc import *

# FridaLink modules

import FridaLink.Common.Config
import FridaLink.Utils.Logging
import FridaLink.Core.Modules as FrLModules
import FridaLink.UI.MainMenu as IdaMainMenu
import FridaLink.UI.PopupMenus as IdaPopupMenus
import FridaLink.UI.SettingsDialog

import FridaLink.Core.IdaEngine
import FridaLink.Core.FridaEngine
import FridaLink.Core.BacktraceEngine
import FridaLink.Core.HookEngine
import FridaLink.Core.ReplaceImplEngine
import FridaLink.Core.MemoryEngine
import FridaLink.Core.DebugEngine
import FridaLink.Core.AppLogEngine
import FridaLink.DB.DBProtocol

# Load current FridaLink config and settings
reload(FridaLink.Common.Config)
reload(FridaLink.Settings.SettingsStorage)

# Reload modules in debug mode
if FridaLink.Common.Config.DevModeEnabled() == True:
	reload(FridaLink.Utils.Logging)
	reload(FridaLink.Core.Modules)
	reload(FridaLink.UI.MainMenu)
	reload(FridaLink.UI.PopupMenus)
	reload(FridaLink.UI.SettingsDialog)
	reload(FridaLink.Core.IdaEngine)
	reload(FridaLink.Core.FridaEngine)
	reload(FridaLink.Core.BacktraceEngine)
	reload(FridaLink.Core.HookEngine)
	reload(FridaLink.Core.ReplaceImplEngine)
	reload(FridaLink.Core.MemoryEngine)
	reload(FridaLink.Core.DebugEngine)
	reload(FridaLink.Core.AppLogEngine)
	reload(FridaLink.DB.DBProtocol)

from FridaLink.Utils.Logging import fl_log as fl_log
from FridaLink.Settings.SettingsStorage import SettingsStorage as FrLSettings
from FridaLink.UI.SettingsDialog import SettingsDialog as FrLSettingsDlg
from FridaLink.UI.LoadModuleDialog import LoadModuleDialog as FrLLoadModuleDlg
from FridaLink.Common.MessageTypes import *

# *** CALLABLES

class TargetInfoCallable:
	def __init__(self, engine, platform, arch):
		self.engine = engine
		self.platform = platform
		self.arch = arch
	def __call__(self):
		self.request = self.engine.handleTargetInfo(self.platform, self.arch)

class ModulesResponseCallable:
	def __init__(self, engine, modules):
		self.engine = engine
		self.modules = modules
	def __call__(self):
		self.engine.handleModulesResponse(self.modules)

class BacktraceCallable:
	def __init__(self, engine, platform, arch, request):
		self.engine = engine
		self.platform = platform
		self.arch = arch
		self.request = request
	def __call__(self):
		self.request = self.engine.handleBacktraceRequest(self.platform, self.arch, self.request)

class HookCallable:
	def __init__(self, engine, platform, arch, response):
		self.engine = engine
		self.platform = platform
		self.arch = arch
		self.response = response
	def __call__(self):
		self.engine.handleHookResponse(self.platform, self.arch, self.response)

class ReplaceCallable:
	def __init__(self, engine, response):
		self.engine = engine
		self.response = response
	def __call__(self):
		self.engine.handleReplaceResponse(self.response)

class MemoryCallable:
	def __init__(self, engine, platform, arch, response):
		self.engine = engine
		self.platform = platform
		self.arch = arch
		self.response = response
	def __call__(self):
		self.engine.handleMemoryFetchResponse(self.platform, self.arch, self.response)

class FraplLogCallable:
	def __init__(self, engine, log_type, log_entry):
		self.engine = engine
		self.log_type = log_type
		self.log_entry = log_entry
	def __call__(self):
		self.engine.handleFraplLog(self.log_type, self.log_entry)

class TargetLogCallable:
	def __init__(self, engine, log_header, log_entry):
		self.engine = engine
		self.log_header = log_header
		self.log_entry = log_entry
	def __call__(self):
		self.engine.handleTargetLog(self.log_header, self.log_entry)

class QueryCallable:
	def __init__(self, engine, db_id, query):
		self.engine = engine
		self.db_id = db_id
		self.query = query
	def __call__(self):
		self.engine.handleDbQuery(self.db_id, self.query)

# *** TCP MESSAGE HANDLER

class MessageHandler(SocketServer.BaseRequestHandler):

	def sendAck(self):
		outJSON = json.dumps({
			"req_id": kFridaLink_Ack, 
		})
		self.request.sendto(outJSON, self.client_address)

	def handleFridaRequest(self, message):
		decoder = json.JSONDecoder()
		while message:
			try:
				inJSON, idx = decoder.raw_decode(message)
			except Exception as e:
				fl_log("FridaLink: unable to decode JSON message: %s\n-------%s\n-------\n" % (str(e), message))
				return
			
			if inJSON['req_id'] == kFridaLink_TargetInfo:
				targetInfo = TargetInfoCallable(self.server.engine, inJSON['platform'], inJSON['arch'])
				execute_sync(targetInfo, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_ModulesResponse:
				modules = ModulesResponseCallable(self.server.engine, inJSON['response'])
				execute_sync(modules, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_HookResponse:
				hook = HookCallable(self.server.engine, inJSON['platform'], inJSON['arch'], inJSON['response'])
				execute_sync(hook, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_ReplaceResponse:
				replace = ReplaceCallable(self.server.engine, inJSON['response'])
				execute_sync(replace, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_FetchMemResponse:
				hook = MemoryCallable(self.server.engine, inJSON['platform'], inJSON['arch'], inJSON['response'])
				execute_sync(hook, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_FraplLogEntry:
				fraplLog = FraplLogCallable(self.server.engine, inJSON['log_type'], inJSON['log_entry'])
				execute_sync(fraplLog, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_TargetLogEntry:
				targetLog = TargetLogCallable(self.server.engine, inJSON['log_header'], inJSON['log_entry'])
				execute_sync(targetLog, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_DBQuery:
				query = QueryCallable(self.server.engine, inJSON['db_id'], inJSON['query'])
				execute_sync(query, 0)
				self.sendAck()
			elif inJSON['req_id'] == kFridaLink_ProcessBacktrace:
				backtrace = BacktraceCallable(self.server.engine, inJSON['platform'], inJSON['arch'], inJSON['request'])
				execute_sync(backtrace, 0)
				outJSON = json.dumps({
					"req_id": kFridaLink_ProcessBacktrace, 
					"request": backtrace.request
				})
				self.request.sendto(outJSON, self.client_address)
			message = message[idx:].lstrip()            

	def handle(self):
		# allow only one FridaLink connection
		if self.server.engine.clientSocket is not None:
			return

		fl_log("FridaLink: established with " + str(self.client_address) + "\n")
		self.server.engine.clientSocket = self.request
		self.server.engine.clientAddress = self.client_address
		self.server.engine.clientConnected();        
		while True:
			header = self.request.recv(8)
			if not header:
				break
			try:
				dataLength = int(header)
			except Exception as e:
				fl_log("FridaLink: invalid header %s" % header)

			msg = ""
			while dataLength != 0:
				bytesToRead = dataLength if dataLength < 4096 else 4096
				data = self.request.recv(bytesToRead)
				msg += data
				dataLength -= bytesToRead

			self.handleFridaRequest(msg)

		self.server.engine.clientSocket = None
		self.server.engine.clientAddress = None
		fl_log("FridaLink: closed with " + str(self.client_address) + "\n")
		self.request.close()

# *** PLUGIN CLASS

class FridaLinkPlugin(	plugin_t,
						FridaLink.Core.IdaEngine.IdaEngineProtocol,
						FridaLink.Core.FridaEngine.FridaEngineProtocol,
						FridaLink.Core.BacktraceEngine.BacktraceEngineProtocol,
						FridaLink.Core.HookEngine.HookEngineProtocol,
						FridaLink.Core.ReplaceImplEngine.ReplaceImplEngineProtocol,
						FridaLink.Core.MemoryEngine.MemoryEngineProtocol,
						FridaLink.Core.DebugEngine.DebugEngineProtocol,
						FridaLink.Core.AppLogEngine.AppLogEngineProtocol,
						FridaLink.DB.DBProtocol.DBProtocol
					  ):

	flags = PLUGIN_KEEP
	comment = ""

	help = "Plugin That uses Frida API"
	wanted_name = "Frida Link"
	wanted_hotkey = ""

	def init(self):
		# FIXME: make init method for modules
		super(FridaLinkPlugin, self).__init__()
		fl_log("FridaLink: loaded\n")
		self.server = None
		self.serverThread = None 
		self.serverStopped = threading.Event()
		self.serverRunning = False

		self.clientSocket = None
		self.clientAddress = None
		return PLUGIN_KEEP

	def clientConnected(self):
		self.requestModules()
		self.syncMemoryRegions()
		self.deleteAllHooks()
		self.syncIdbHooks()
		self.syncImportHooks()
		self.deleteAllReplaced()
		self.syncReplacedFuncs();

	def run(self, arg):
		fl_log("FridaLink: run(%d)\n" % arg)

	def run(self, arg = 0):
		fl_log("FridaLink: run()\n")
		IdaMainMenu.RegisterMenuActions(self)
		IdaPopupMenus.RegisterMenuActions(self)
		self.handleBuildImport()
		self.startServer()

	def termThread(self):
		if self.clientSocket is not None:
			self.clientSocket.shutdown(socket.SHUT_RDWR)

		if self.serverRunning == True:
	 		self.server.shutdown()
	 		self.serverStopped.wait(3)
	 		self.serverThread = None

		fl_log("FridaLink: unloaded\n")

	def term(self):	# asynchronous unload (external, UI_Hook::term)

		# disable logging to prevert deadlocks
		FrLSettings().setLogEnabled(False);

		if self.clientSocket is not None:
			self.clientSocket.shutdown(socket.SHUT_RDWR)

		if self.serverRunning == True:
	 		self.server.shutdown()

	 	self.closeAllDBs();

		self.resetHookedColors();
		self.resetReplacedColors();

		IdaPopupMenus.UnregisterMenuActions()
		IdaMainMenu.UnregisterMenuActions()

		request_refresh(0xFFFFFFFF) # IWID_ALL

	def unloadPlugin(self):	# synchronous unload (internal, Main Menu)

		self.resetHookedColors();
		self.resetReplacedColors();

		IdaPopupMenus.UnregisterMenuActions()
		IdaMainMenu.UnregisterMenuActions()

		request_refresh(0xFFFFFFFF) # IWID_ALL

		# do not block UI while stopping server
		t = threading.Thread(target=self.termThread)
		t.start()

	def doNothing(self):
		pass

	def logToggle(self):
		if FrLSettings().getLogEnabled():
			fl_log("FridaLink: Debug Log disabled\n")
		FrLSettings().setLogEnabled(not FrLSettings().getLogEnabled())
		if FrLSettings().getLogEnabled():
			fl_log("FridaLink: Debug Log enabled\n")

	def loadModule(self):
		loadModuleDlg = FrLLoadModuleDlg()
		loadModuleDlg.Compile()
		loadModuleDlg.filePath.value = "*"
		ok = loadModuleDlg.Execute()
	
		if ok == 1:
			path = loadModuleDlg.filePath.value
			if path != "*":
				module = os.path.basename(path)
				FrLModules.LoadModule(None, module, path)
	
		loadModuleDlg.Free()

	def overwriteSymbolToggle(self):
		FrLSettings().setOverwriteSymbolName(not FrLSettings().getOverwriteSymbolName())
		if FrLSettings().getOverwriteSymbolName():
			fl_log("FridaLink: Symbol Overwrite enabled\n")
		else:
			fl_log("FridaLink: Symbol Overwrite disabled\n")

	def loadProject(self):
		if self.clientSocket == None:
			fl_log("FridaLink: Frida not connected\n");
			return

		filePath = AskFile(0, "*.flp", "Open FridaLink project")
		if filePath is None:
			return
		with open(filePath, 'rb') as file:
			settings = pickle.load(file)
			hooks = pickle.load(file)
			repl = pickle.load(file)
			memory = pickle.load(file)
			frida = pickle.load(file)

			FrLSettings().restore(settings)
			self.deleteAllHooks();
			self.deleteAllReplaced();
			self.restoreHookData(hooks)
			self.restoreReplacedFuncData(repl)
			self.restoreMemoryData(memory)
			self.restoreFridaData(frida)

			file.close()
			fl_log("FridaLink: project loaded\n")

	def saveProject(self):
		filePath = AskFile(1, "*.flp", "Save FridaLink project")
		if filePath is None:
			return
		with open(filePath, 'wb') as file:
			pickle.dump(FrLSettings(), file, pickle.HIGHEST_PROTOCOL)
			pickle.dump(self.backupHookData(), file, pickle.HIGHEST_PROTOCOL)
			pickle.dump(self.backupReplacedFuncData(), file, pickle.HIGHEST_PROTOCOL)
			pickle.dump(self.backupMemoryData(), file, pickle.HIGHEST_PROTOCOL)
			pickle.dump(self.backupFridaData(), file, pickle.HIGHEST_PROTOCOL)
			file.close()
			fl_log("FridaLink: project saved\n")

	def showSettings(self):
		settingsDlg = FrLSettingsDlg()
		settingsDlg.Compile()
		settingsDlg.host.value = FrLSettings().getHost()
		settingsDlg.port.value = FrLSettings().getPort()
		settingsDlg.cpuctx_cols.value = FrLSettings().getCpuContextColumns()
		ok = settingsDlg.Execute()
		if ok == 1:
			FrLSettings().setHost(settingsDlg.host.value)
			FrLSettings().setPort(settingsDlg.port.value)
			if FrLSettings().getCpuContextColumns() != settingsDlg.cpuctx_cols.value:
				FrLSettings().setCpuContextColumns(settingsDlg.cpuctx_cols.value)
				self.refreshCpuViews()

	def serverTask(self):
		try:
			self.server = SocketServer.ThreadingTCPServer(
				(FrLSettings().getHost(), FrLSettings().getPort()), 
				MessageHandler
			)
			self.server.engine = self
	
			fl_log("FridaLink: SERVER STARTED!\n")

			self.serverStopped.clear()
			self.serverRunning = True
			self.server.serve_forever()
			self.serverRunning = False

			self.server.server_close()
			self.server = None

			fl_log("FridaLink: SERVER STOPPED!\n")
			self.serverStopped.set()
		except Exception, e:
			fl_log("FridaLink: Unable to start server!\n")
			self.server = None
			raise e

	def startServer(self):
		FrLModules.SanityChecks()
		FrLModules.GetNextModuleBase()

		if FrLSettings().getOverwriteSymbolName():
			fl_log("FridaLink: Symbol Overwrite enabled\n")
		else:
			fl_log("FridaLink: Symbol Overwrite disabled\n")

		if self.server is None:
			self.serverThread = threading.Thread(target=self.serverTask)
			self.serverThread.start()

	def stopServer(self):
		if self.serverRunning == True:
			self.server.shutdown()
	 		self.serverStopped.wait(3)
	 		self.serverThread = None

#def PLUGIN_ENTRY():
#    return FridaLinkPlugin()
if __name__ == '__main__':
	fridaLink = FridaLinkPlugin()
	fridaLink.init()
	fridaLink.run()


repository: ZjW1nd/Better-printk-IDA9
file: better_printk.py
description: Better printk display plugin for IDA (9.0).
wanted_name: Printk Optimizer
comment: Optimize printk calls in Linux kernel modules
created_at: 2025-01-21 13:50:32+00:00
pushed_at: 2025-02-28 08:22:05+00:00
url:
https://github.com/ZjW1nd/Better-printk-IDA9/blob/bf83c56496ed077c9f9d2a2767730ecc3f4096f9/better_printk.py

content:
# Created By Zj_W1nd
# Repo: https://github.com/ZjW1nd/Better-printk-IDA9
import ida_idaapi
import ida_bytes
import ida_funcs
import ida_hexrays
import idc
import ida_lines
import ida_typeinf

KERN_LEVELS = {
    '0': "EMERG",
    '1': "ALERT",
    '2': "CRIT",
    '3': "ERR",
    '4': "WARNING",
    '5': "NOTICE",
    '6': "INFO",
    '7': "DEBUG"
}

class PrintkVisitor(ida_hexrays.ctree_visitor_t):
    """负责遍历和优化printk调用的访问者类"""
    def __init__(self, printk_func, modifications):
        ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST | ida_hexrays.CV_INSNS)
        self.printk_func = printk_func
        self.modifications = modifications

    def visit_insn(self, insn):
        # 检查指令类型
        if insn.op != ida_hexrays.cit_expr:
            return 0
        # 获取表达式
        expr = insn.cexpr
        if expr.op != ida_hexrays.cot_call:
            return 0
        # 获取被调用函数
        called_func = expr.x
        if called_func.op != ida_hexrays.cot_obj:
            return 0
        if called_func.obj_ea != self.printk_func:
            return 0
        # 获取第一个参数
        if len(expr.a) < 1:
            return 0
        arg = expr.a[0]
        if arg.op != ida_hexrays.cot_ref:
            return 0
        # 获取字符串地址
        str_addr = arg.x.obj_ea
        # 读取字符串内容
        if ida_bytes.get_byte(str_addr) != 0x1:
            return 0
        
        level = chr(ida_bytes.get_byte(str_addr + 1))
        if level not in KERN_LEVELS:
            return 0
        # 读取实际字符串
        str_start = str_addr + 2
        str_end = str_start
        while ida_bytes.get_byte(str_end) != 0:
            str_end += 1
        actual_str = ida_bytes.get_bytes(str_start, str_end - str_start).decode('ascii')
        # 处理回车
        actual_str = actual_str.replace('\n', '\\n')
        # 生成优化后的字符串
        level_str = KERN_LEVELS[level]
        optimized_str = f'KERN_{level_str}, "{actual_str}"'
        self.modifications.append(optimized_str)
        print(f"Found printk at {insn.ea}, extract string: {optimized_str}")
        idc.set_cmt(insn.ea, optimized_str, 0) # 汇编页面的注释
        return 0

class PrintkOptimizer(ida_hexrays.Hexrays_Hooks):
    """负责管理printk优化过程的主类"""
    def __init__(self):
        super(PrintkOptimizer, self).__init__()
        self.printk_func = None
        self.modifications = [] # 存储需要修改的printk

    def find_printk(self):
        """查找printk函数的地址"""
        func = ida_funcs.get_next_func(0)
        while func:
            func_name = ida_funcs.get_func_name(func.start_ea)
            if func_name == "printk":
                self.printk_func = func.start_ea
                return
            func = ida_funcs.get_next_func(func.start_ea)
        raise Exception("printk function not found.")

    def optimize_printk(self, cfunc):
        """优化printk调用"""
        if self.printk_func is None:
            try:
                self.find_printk()
            except Exception as e:
                print(f"Better-printk: {e}")
                return False
        
        # 创建并应用访问者
        visitor = PrintkVisitor(self.printk_func, self.modifications)
        visitor.apply_to(cfunc.body, None)
        return True
    
    def modify_printk_line(self, sl, ord):
        """伪代码行的注释，按序匹配"""
        # 实际修改太麻烦了，而且还需要hook双击动作，尝试修改后发现很僵硬，双击跳转参数也没了
        # 因此选择在伪代码行尾添加注释
        # todo: 匹配是按照printk在反汇编Ctree中遍历的先后顺序来的
        # 如果输出反编译伪代码时printk调用行与Ctree遍历顺序不一致则会有问题
        # 汇编是100%正确的
        # Actually modifying the code is too troublesome, and it also requires hooking double-click actions. 
        # After attempting to modify it, I found it to be very rigid, and the double-click jump parameters were lost.
        # Therefore, I chose to add comments at the end of the pseudocode lines.
        # todo: The matching is based on the order in which printk is traversed in the disassembled Ctree.
        # If the order of printk call lines in the decompiled pseudocode output does not match the Ctree traversal order, there will be problems.
        # The assembly is 100% correct.
        l : str = sl.line
        optimized_str=self.modifications[ord]
        # 在行尾添加注释
        comment = f"  /* {optimized_str} */"
        # 设置注释颜色为灰色
        colored_comment = ida_lines.COLSTR(comment, ida_lines.SCOLOR_AUTOCMT)
        # 将注释添加到行尾
        sl.line = l + colored_comment

    def func_printed(self, cfunc):
        """Hex-Rays回调函数，在函数反编译完成后调用"""
        # 优化 printk 调用
        self.optimize_printk(cfunc)
        ord=0
        for sl in cfunc.get_pseudocode():
            if("printk" in sl.line):
                self.modify_printk_line(sl,ord)
                ord+=1
        self.modifications.clear()
        return 0

class PrintkOptimizerPlugin(ida_idaapi.plugin_t):
    """Module provided by ida examples"""
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Printk Optimizer"
    wanted_hotkey = ""
    comment = "Optimize printk calls in Linux kernel modules"
    help = "This plugin optimizes printk calls by decoding kernel log levels and messages"

    def init(self):
        print(">>> PrintkOptimizerPlugin: Init called.")
        if not ida_hexrays.init_hexrays_plugin():
            print("Hex-Rays plugin not found, skipping PrintkOptimizer initialization.")
            return ida_idaapi.PLUGIN_SKIP
        self.hooks = PrintkOptimizer()
        self.hooks.hook()
        return ida_idaapi.PLUGIN_KEEP

    def term(self):
        if hasattr(self, 'hooks'):
            self.hooks.unhook()
        return

    def run(self, arg):
        print(">>> PrintkOptimizer: run() is invoked.")
        return 0

def PLUGIN_ENTRY():
    return PrintkOptimizerPlugin()

# 允许更好查看数据，右键指针选择Better printk即可
# Warning: 对于低于9.0版本的ida（个人使用8.3）来说，相关api存在问题，可以注释掉使用
# Warning: For ida below 9.0 (I use 8.3), there are problems with the related api
class PrintkDataType(ida_bytes.data_type_t):
    def __init__(self):
        ida_bytes.data_type_t.__init__(
            self,
            "Better_printk_string",
            1,  # 最小大小
            "Better printk string",
            None,
            "printk_str")

    def calc_item_size(self, ea, maxsize):
        # 检查是否是结构体成员
        tif = ida_typeinf.tinfo_t()
        if tif.get_udm_by_tid(None, ea) != -1:
            return 1

        # 检查第一个字节是否为 0x1
        if ida_bytes.get_byte(ea) != 0x1:
            return 0

        # 检查第二个字节是否为有效的日志等级
        level = chr(ida_bytes.get_byte(ea + 1))
        if level not in KERN_LEVELS:
            return 0

        # 计算字符串长度
        str_start = ea + 2
        str_end = str_start
        while ida_bytes.get_byte(str_end) != 0:
            str_end += 1
            if str_end - ea > maxsize:
                return 0

        return str_end - ea + 1  # 包括结束符

class PrintkDataFormat(ida_bytes.data_format_t):
    FORMAT_NAME = "py_printk_string_format"
    def __init__(self):
        ida_bytes.data_format_t.__init__(
            self,
            PrintkDataFormat.FORMAT_NAME)

    def printf(self, value, current_ea, operand_num, dtid):
        # 检查第一个字节是否为 0x1
        if value[0] != 0x1:
            return None

        # 获取日志等级
        level = chr(value[1])
        if level not in KERN_LEVELS:
            return None

        # 获取实际字符串
        str_start = 2
        str_end = len(value)
        actual_str = value[str_start:str_end].decode('ascii', errors='replace')
        # 处理回车
        actual_str = actual_str.replace('\n', '\\n')
        # 生成格式化字符串
        level_str = KERN_LEVELS[level]
        return f"KERN_{level_str} '{actual_str}'\x00"

# 注册自定义数据类型和格式
new_formats = [
    (PrintkDataType(), PrintkDataFormat()),
]

if ida_bytes.find_custom_data_type(PrintkDataFormat.FORMAT_NAME) == -1:
    if not ida_bytes.register_data_types_and_formats(new_formats):
        print("Failed to register printk data type!")
    else:
        print("printk data type registered successfully!")


repository: LAC-Japan/IDA_Plugin_AntiDebugSeeker
file: sources/python/ver1.0/AntiDebugSeeker.py
description: Automatically identify and extract potential anti-debugging techniques used by malware.
wanted_name: AntiDebugSeeker
comment: Anti Debug Detection Plugin
created_at: 2023-09-11 02:29:26+00:00
pushed_at: 2024-11-22 03:45:20+00:00
url:
https://github.com/LAC-Japan/IDA_Plugin_AntiDebugSeeker/blob/070e26d115734a988c53a39bd8ee7bcec7e9a957/sources/python/ver1.0/AntiDebugSeeker.py

content:
import idaapi
import ida_bytes
import idc
import idautils
import ida_kernwin
import ida_lines
import os
import re
import json

class AntiDebugResults(ida_kernwin.Choose):
    def __init__(self, title, items):
        ida_kernwin.Choose.__init__(self, title, [["Category Name", 20 | ida_kernwin.Choose.CHCOL_PLAIN],
                                                 ["Possible Anti-Debug API", 20 | ida_kernwin.Choose.CHCOL_PLAIN],
                                                 ["Address", 15 | ida_kernwin.Choose.CHCOL_HEX],
                                                 ["Possible Anti-Debug Technique", 25 | ida_kernwin.Choose.CHCOL_PLAIN],
                                                 ["Address", 15 | ida_kernwin.Choose.CHCOL_HEX]])
        self.items = items

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

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

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

    def OnGetLineAttr(self, n):
        black = 0x000000
        gray = 0xC0C0C0
        if self.items[n][0] == "" and self.items[n][1] == "":
            return [ida_kernwin.Choose.CHCOL_PLAIN | black, ida_kernwin.Choose.CHCOL_PLAIN | black]
        elif n % 2 == 0:
            return [ida_kernwin.Choose.CHCOL_PLAIN | gray, ida_kernwin.Choose.CHCOL_PLAIN | gray]
        else:
            return [ida_kernwin.Choose.CHCOL_PLAIN | gray, ida_kernwin.Choose.CHCOL_PLAIN | gray]

    def OnSelectLine(self, n):
        col2_addr = self.items[n][2]
        col4_addr = self.items[n][4]

        if col2_addr:
            idaapi.jumpto(int(col2_addr, 16))
        if col4_addr:
            idaapi.jumpto(int(col4_addr, 16))

    def OnDblClick(self, n):
        widget = ida_kernwin.find_widget("Anti Debug Detection Results")
        if not widget:
            return

        sel = self.GetEmbSelection()
        if not sel:
            return

        col = sel[0].y
        addr = self.items[n][col]

        if addr:
            idaapi.jumpto(int(addr, 16))
        return 1

    def show(self):
        self.Show()

class ConfigEditorForm(ida_kernwin.Form):
    def __init__(self, file_path):
        self.file_path = file_path
        try:
            with open(file_path, 'r') as f:
                self.file_contents = f.read()
        except Exception as e:
            print("Error reading file: %s" % str(e))
            self.file_contents = ""

        F = ida_kernwin.Form
        F.__init__(self, r"""STARTITEM {id:txtInput}
BUTTON YES* Save
BUTTON CANCEL Cancel
Edit configuration
{FormChangeCb}
<##File contents:{txtInput}>
        """, {
            'txtInput': F.MultiLineTextControl(text=self.file_contents),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange),
        })

    def OnFormChange(self, fid):
        return 1

    def OnButtonOk(self, code=0):
        try:
            with open(self.file_path, 'w') as f:
                f.write(self.GetControlValue(self.txtInput).text)
            return 1
        except Exception as e:
            print("Error writing to file: %s" % str(e))
            return 0

    def OnFormChange(self, fid):
        if fid == -2:
            self.OnButtonOk()
        return 1


class AntiDebugPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Anti Debug Detection Plugin"
    help = "Detects the use of Anti Debug Technique"
    wanted_name = "AntiDebugSeeker"
    wanted_hotkey = "Ctrl-Shift-D"

    def init(self):
        idaapi.msg("AntiDebugSeeker initialized. Ctrl-Shift-D Start Analysis.\n")
        self.edit_config_hotkey_ctx = ida_kernwin.add_hotkey("Ctrl-Shift-E", self.edit_config)
        self.saved_results = []
        if self.edit_config_hotkey_ctx is None:
            idaapi.msg("Failed to register hotkey for config editor. It might be already in use.\n")
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        if arg == "edit":
            self.edit_config()
        elif arg == "details":
            self.show_anti_debug_details()
        else:
            idaapi.msg("AntiDebugSeeker is running...\n")
            self.detect_anti_debug_functions()

    def save_anti_debug_results(self, result_win):
        self.saved_results = result_win

    def collect_anti_debug_results(self):
        return self.saved_results

    def term(self):
        idaapi.msg("AntiDebugSeeker terminated.\n")
        idaapi.msg("Edit anti_debug.config : Switch Other tab and Press Ctrl+Shift+E.\n")

    def load_anti_debug_descriptions(self, file_name):
        
        plugin_dir = os.path.dirname(os.path.realpath(__file__))
        file_path = os.path.join(plugin_dir, file_name)
        with open(file_path, 'r') as f:
            return json.load(f)

    def edit_config(self):
        plugin_dir = os.path.dirname(os.path.realpath(__file__))
        file_path = os.path.join(plugin_dir, "anti_debug.config")
        editor = ConfigEditorForm(file_path)
        editor.Compile()
        editor.Execute()

    def load_config(self, file_name):
        config_data = {
            "Anti_Debug_API": {},
            "Anti_Debug_Technique": [],
            "default_search_range": None
        }

        plugin_dir = os.path.dirname(os.path.realpath(__file__))
        file_path = os.path.join(plugin_dir, file_name)

        current_section = None
        current_category = None

        with open(file_path, "r") as f:
            for line in f:
                line = line.strip()

                if line.startswith("###") and line.endswith("###"):
                    current_section = line[3:-3]
                    current_category = None
                elif line.startswith("[") and line.endswith("]"):
                    current_category = line[1:-1]
                    if current_section == "Anti_Debug_API":
                        config_data[current_section][current_category] = []
                    elif current_section == "Anti_Debug_Technique":
                        config_data["Anti_Debug_Technique"].append({
                            "name": current_category,
                            "search_keyword": None,
                            "nearby_keyword": None,
                            "nearby_keyword2": None,
                            "search_range": config_data["default_search_range"]
                        })
                elif current_category and line and current_section in ["Anti_Debug_API", "Anti_Debug_Technique"]:
                    if current_section == "Anti_Debug_API":
                        config_data[current_section][current_category].append(line)
                    elif current_section == "Anti_Debug_Technique":
                        entry = config_data["Anti_Debug_Technique"][-1]
                        if line.startswith("search_range="):
                            entry["search_range"] = int(line.split("=")[1])
                        else:
                            if entry["search_keyword"] is None:
                                entry["search_keyword"] = line
                            elif entry["nearby_keyword"] is None:
                                entry["nearby_keyword"] = line
                            elif entry["nearby_keyword2"] is None:
                                entry["nearby_keyword2"] = line
                elif line.startswith("default_search_range=") and current_section == "Anti_Debug_Technique":
                    config_data["default_search_range"] = int(line.split("=")[1])

        return config_data



    def search_asm(self, search_keyword, nearby_keyword=None, nearby_keyword2=None, search_range=None):
        api_names = set(name for _, name in idautils.Names())
        one_keyword_match = []
        matching_sets = []

        def search_keyword_in_code(keyword, ea, disasm_line, end_addr):
            found = False
            keyword_addr = None
            
            keyword_is_int = False
            try:
                keyword_as_int = int(keyword)
                keyword_is_int = True
            except ValueError:
                pass
            
            regex_matched = re.search(r'\b' + re.escape(str(keyword)) + r'(h)?\b', disasm_line)
            comment = ida_bytes.get_cmt(ea, 0)
            
            def check_found(keyword_is_int, regex_matched):
                return (not keyword_is_int) or (keyword_is_int and regex_matched)

            if keyword in api_names:
                keyword_addr = idc.get_name_ea_simple(keyword)

                if keyword_addr != idaapi.BADADDR and ea in list(idautils.CodeRefsTo(keyword_addr, 0)):
                    found = True
                    keyword_addr = ea

            if comment and keyword in comment:
               found = check_found(keyword_is_int, regex_matched)
                
            if keyword in disasm_line:
               found = check_found(keyword_is_int, regex_matched)

            if not found:
                ref_addr = idc.get_operand_value(ea, 0)
                if ref_addr != idaapi.BADADDR:
                    strtype = idaapi.get_str_type(ref_addr)
                    if strtype == 0xFFFFFFFF:
                        return False, None
                    ref_string = ida_bytes.get_strlit_contents(ref_addr, ida_bytes.get_item_size(ref_addr), int(strtype))
                    if ref_string and keyword in ref_string.decode('utf-8', errors='ignore'):
                        found = check_found(keyword_is_int, regex_matched)
                        
            if found:
               if keyword_addr != ea:
                  keyword_addr = ea
                        
            return found, keyword_addr
            
        entry_point = idc.get_inf_attr(idc.INF_START_IP)  
        entry_section = idaapi.getseg(entry_point)
        
        for seg_ea in idautils.Segments():
            if seg_ea != entry_section.start_ea:  
                continue
               
            ea = seg_ea
            end_addr = idc.get_segm_end(seg_ea)

            while ea < end_addr:
                disasm_line = idc.GetDisasm(ea)
                search_found, search_addr = search_keyword_in_code(search_keyword, ea, disasm_line, end_addr)

                if search_found:
                    base_addr = search_addr if search_keyword in api_names else ea
                    if not nearby_keyword and not nearby_keyword2:
                        one_keyword_match.append(hex(base_addr))
                    else:
                        next_ea = base_addr
                        while next_ea - base_addr < search_range and next_ea < end_addr:
                            disasm_next_line = idc.GetDisasm(next_ea)
                            found_nearby, nearby_addr = search_keyword_in_code(nearby_keyword, next_ea, disasm_next_line, end_addr)
                            if found_nearby:
                                if not nearby_keyword2:
                                    matching_sets.append((hex(base_addr), hex(next_ea), None))
                                else:
                                    next_ea2 = next_ea
                                    while next_ea2 - next_ea < search_range and next_ea2 < end_addr:
                                        disasm_next_line2 = idc.GetDisasm(next_ea2)
                                        found_nearby2, nearby_addr2 = search_keyword_in_code(nearby_keyword2, next_ea2, disasm_next_line2, end_addr)
                                        if found_nearby2:
                                            matching_sets.append((hex(base_addr), hex(next_ea), hex(next_ea2)))
                                        next_ea2 = idc.next_head(next_ea2)
                            next_ea = idc.next_head(next_ea)
                ea = idc.next_head(ea)

        if one_keyword_match:
           print(f"[Result] Found {len(one_keyword_match)} match(es) for First_keyword '{search_keyword}' under the rule of specifying one keyword.")
           return one_keyword_match
            
        if matching_sets:
           print(f"[Result] Found {len(matching_sets)} matching set(s).")
           return matching_sets
        else:
            return None, None, None

    def detect_anti_debug_functions(self):
        results = []
        color = 0x98FB98
        color2 = 0xAAD4FF

        def set_item_color(address, color):
            try:
                idaapi.set_item_color(int(address, 16), color)
            except TypeError as e:
                print(f"Error: {e}")

        def add_comment(address, comment_name):
            existing_comment = idaapi.get_cmt(address, 0)
            comment_description = anti_debug_techniques_descriptions.get(comment_name, "")
            comment_text = comment_name if existing_comment is None else existing_comment + " | " + comment_name
            if comment_description:
                comment_text += " - " + comment_description
            idaapi.set_cmt(address, comment_text, 0)

        config_data = self.load_config("anti_debug.config")
        anti_debug_functions = config_data["Anti_Debug_API"]
        search_patterns = config_data["Anti_Debug_Technique"]
        default_search_range = config_data["default_search_range"]
        anti_debug_techniques_descriptions = self.load_anti_debug_descriptions("anti_debug_techniques_descriptions.json")
        
        def is_address_in_target_sections(address, func, entry_section):
           if entry_section is not None:
               section_start = entry_section.start_ea
               section_end = entry_section.end_ea
               if section_start <= address <= section_end:
                   return True

           if func is None:
               return False

           func_start = func.start_ea
           func_end = func.end_ea

           code_section = idaapi.getseg(func_start)
           if code_section is not None and code_section.perm & idaapi.SEGPERM_EXEC:
               return True

           return False

            
        entry_point = idc.get_inf_attr(idc.INF_START_IP)
        entry_section = idaapi.getseg(entry_point)

        for category, functions in anti_debug_functions.items():
            for func_name in functions:
                func_addr = idc.get_name_ea_simple(func_name)

                if func_addr == idaapi.BADADDR:
                    continue

                xrefs = list(idautils.XrefsTo(func_addr))
                
                if xrefs:
                    idaapi.msg(f"{func_name} function found.\n")
                else:
                    continue

                processed_addresses = set()
                for xref in xrefs:
                    address = xref.frm
                    address_str = f"0x{address:X}"
                    containing_func = idaapi.get_func(xref.frm)
                    
                    if address_str in processed_addresses:
                       continue
                    
                    if is_address_in_target_sections(address, containing_func, entry_section):
                        idaapi.msg(f"  {address_str}\n")
                        set_item_color(address_str, color)
                        add_comment(xref.frm, category)
                        results.append([category, func_name, address_str, "", ""])
                        processed_addresses.add(address_str)

        start_addr = idaapi.cvar.inf.min_ea
        end_addr = idaapi.cvar.inf.max_ea

        for pattern in search_patterns:
            search_keyword = pattern["search_keyword"]
            nearby_keyword = pattern.get("nearby_keyword", None)
            nearby_keyword2 = pattern.get("nearby_keyword2", None)
            search_range = pattern.get("search_range", default_search_range)

            asm_address = nearby_address = nearby2_address = None

            if nearby_keyword2:
                result = self.search_asm(search_keyword, nearby_keyword, nearby_keyword2, search_range=search_range)
            elif nearby_keyword:
                result = self.search_asm(search_keyword, nearby_keyword, search_range=search_range)
            else:
                result = self.search_asm(search_keyword)
                
            if result:
                if isinstance(result, list):
                    print(f"Found for pattern {pattern['name']}.")
                    for element in result: 
                        if element is not None:
                           if isinstance(element, tuple):
                               asm_address, nearby_address, nearby2_address = element + (None,) * (3 - len(element))
                               for address in [asm_address, nearby_address, nearby2_address]:
                                   if address is not None:
                                       address_int = int(address, 16)
                                       set_item_color(address, color2) 
                                       
                               if asm_address is not None:
                                   address_int = int(asm_address, 16)
                                   add_comment(address_int, pattern["name"])
                                   results.append(["", "", "", pattern["name"], str(asm_address), ""])

                           elif isinstance(element, str):
                              address_int = int(element, 16)
                              set_item_color(element, color2)
                              add_comment(address_int, pattern["name"])
                              results.append(["", "", "", pattern["name"], str(element), ""])
                              
                elif isinstance(result, tuple):
                    if result == (None, None, None):
                       print(f"Nothing Found for pattern {pattern['name']}.")

        sorted_results = sorted(results, key=lambda x: (int(x[2], 16) if x[2] else int(x[4], 16) if x[4] else float("inf")))
        self.save_anti_debug_results(sorted_results)

        result_win = AntiDebugResults("Anti Debug Detection Results", sorted_results)
        result_win.show()

def PLUGIN_ENTRY():
    return AntiDebugPlugin()


repository: airbus-cert/etwbreaker
file: etwbreaker.py
description: An IDA plugin to deal with Event Tracing for Windows (ETW)
wanted_name: ETWBreaker
created_at: 2020-05-29 14:10:42+00:00
pushed_at: 2022-07-08 09:20:29+00:00
url:
https://github.com/airbus-cert/etwbreaker/blob/9fb7a6cd883bd23cbf328a2b152cb3a72417f087/etwbreaker.py

content:
"""
ETWBreaker is an IDA plugin that find all references about ETW into a module

Currently ETWBreaker can work with Manifest-based ETW and Tracelogging.

For Manifest-based ETW, ETWBreaker will parse module ressources to find and parse the manifest
and show all possible events handled by the current module.

It tries to parse also all tracelogging events, which are more hidden than Manifest-based ones.

In the end ETWBreaker can generate a conditional breakpoint to dynamically analyze events.
"""

import idaapi
import idautils
import ida_dbg
import ida_name
import idc
import sys
import struct
import textwrap
from io import BytesIO
from PyQt5 import QtCore, QtWidgets, QtGui
from typing import List, Tuple

__author__ = "Airbus CERT"


class ETWBreakerException(Exception):
    """
    Base exception for all exception of ETW breaker
    """
    def __init__(self, message):
        super().__init__(message)


class ETWBreakerWevtTemplateNotFound(ETWBreakerException):
    """
    The WEVT_TEMPLATE ressource was not found
    """
    def __init__(self):
        super().__init__("WEVT_TEMPLATE resource not found.")


class ETWBreakerTLNotFound(ETWBreakerException):
    """
    The tracelogging magic was not found
    """
    def __init__(self):
        super().__init__("Trace logging not found")


class ETWBreakerUnexpectedToken(ETWBreakerException):
    """
    During parsing an unexpected token was found.
    Please open an issue on Github.
    """
    def __init__(self, expected, found):
        super().__init__("Unexpected token. Expected %s, found %s"%(expected, found))


class Stream(BytesIO):
    """
    A wrapper that is nicer to understand
    """
    def read_u32(self) -> int:
        return struct.unpack("<I", self.read(4))[0]

    def read_u16(self) -> int:
        return struct.unpack("<H", self.read(2))[0]

    def read_u8(self) -> int:
        return struct.unpack("<B", self.read(1))[0]

    def read_u64(self) -> int:
        return struct.unpack("<Q", self.read(8))[0]


class Guid:
    """
    A global unique identifier
    """
    def __init__(self, raw):
        self.raw = raw

    def __str__(self):
        Data1 = struct.unpack("<I", self.raw[0:4])[0]
        Data2 = struct.unpack("<H", self.raw[4:6])[0]
        Data3 = struct.unpack("<H", self.raw[6:8])[0]
        Data4 = self.raw[8:16]
        return "%08x-%04x-%04x-%s-%s" % (Data1, Data2, Data3, "".join("%02x" % x for x in Data4[0:2]), "".join("%02x" % x for x in Data4[2:]))



class ETWBreaker(idaapi.ida_idaapi.plugin_t):
    """
    This is the main plugin class
    """
    comment = ""
    help = ""
    flags = idaapi.PLUGIN_MOD
    wanted_name = 'ETWBreaker'
    wanted_hotkey = 'Ctrl-Shift-L'
    hxehook = None

    def init(self):
        """
        Init plugin function
        """
        if idc.get_inf_attr(idc.INF_FILETYPE) != idc.FT_PE:
            # skip if it's not a PE
            return idaapi.PLUGIN_SKIP
        ETWBreaker.log("'%s' loaded. %s activates/deactivates synchronization." % (ETWBreaker.wanted_name, ETWBreaker.wanted_hotkey))
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        Launch when you press Ctrl-Shift-L
        """
        ETWBreaker.log("%s is enabled now." % ETWBreaker.wanted_name)
        providers = []

        # Manifest based provider
        try:
            providers += parse_manifest(find_wevt_template(*find_segment(".rsrc")[0]))
        except IndexError:
            ETWBreaker.log("Please consider reloading the file and check the 'Load resources' checkbox")
        except ETWBreakerException as e:
            ETWBreaker.log(str(e))

        # Tracelogging
        for segment in find_segment(".rdata"):
            try:
                providers += parse_tracelogging(find_tracelogging_meta(*segment))
            except ETWBreakerException as e:
                ETWBreaker.log(str(e))

        ETWResultsForm(providers).show()


    def term(self):
        ETWBreaker.log("%s unloaded." % (ETWBreaker.wanted_name))

    def log(message):
        idaapi.msg("[%s] %s\n" % (ETWBreaker.wanted_name, message))


class Event:
    """
    An ETW event
    """
    def __init__(self, event_id: int, version: int, channel: int, level: int, opcode: int, task: int, keyword: int):
        self.event_id = event_id
        self.version = version
        self.channel = channel
        self.level = level
        self.opcode = opcode
        self.task = task
        self.keyword = keyword

    def find_symbol(self) -> str:
        """
        Try to find a symbol associated to the event

        This is based on the event header signature
        Most of then are into .rdata segment and some of them have a name
        """
        pattern = struct.pack("<HBBBBHQ", self.event_id, self.version, self.channel, self.level, self.opcode, self.task, self.keyword)
        for start, end in find_segment(".rdata"):
            offset = idc.get_bytes(start, end - start).find(pattern)
            if offset == -1:
                continue

            name = ida_name.get_name(start+offset)
            if name is None:
                continue

            return name
        return None


class Channel:
    """
    Channel is a pure ETW concept
    """
    def __init__(self, identifier: int, name: str):
        """
        :ivar identifier: unique identifier of the channel
        :ivar name: name of the channel, generally include the provider name
        """
        self.identifier = identifier
        self.name = name

    def __str__(self):
        return self.name


class Provider:
    """
    An ETW Provider is defined by a unique GUID
    and a list of event
    """
    def __init__(self, guid: Guid, events: List[Event], channels: List[Channel]):
        """
        :ivar guid: An unique global identifier
        :ivar events: A list of event that could be emitted  by the provider
        :ivar channels: A list of all channel identifier available
        """
        self.guid = guid
        self.events = events
        self.channels = channels

    def find_channel(self, identifier: int) -> Channel:
        """
        Try to find a channel from its identifier
        """
        return next((channel for channel in self.channels if channel.identifier == identifier), None)


class ManifestProvider(Provider):
    """
    Convenient class use to identify Manifest based providers
    """


class TraceLoggingProvider(Provider):
    """
     Convenient class use to identify TraceLogging providers
    """


def find_segment(name: str) -> List[Tuple[int, int]]:
    """
    Try to find the segment from name

    :ivar name: name of segment
    :ret: Start ant end address
    """
    result = []
    for seg in idautils.Segments():
        if idc.get_segm_name(seg) == name:
            result.append((idc.get_segm_start(seg), idc.get_segm_end(seg)))
    return result


def find_wevt_template(start, end) -> Stream:
    """
    This function try to retrieve the WEVT_TEMPLETE resource
    This resource start with the magic CRIM

    :ivar start: start address
    :ivar end: end address
    :ret: Stream use to parse Manifest based provider or raise an exception
    :raise: ETWBreakerWevtTemplateNotFound
    """
    resource = idc.get_bytes(start, end - start)
    result = resource.find(b"CRIM")
    if result == -1:
        raise ETWBreakerWevtTemplateNotFound()

    return Stream(resource[result:])


def find_tracelogging_meta(start, end) -> Stream:
    """
    Try to find ETW0 magic

    :ivar start: start address
    :ivar end: end address
    :ret: Stream use to parse tracelogging or None if not found
    """
    data = idc.get_bytes(start, end - start)
    result = data.find(b"ETW0")
    if result == -1:
        raise ETWBreakerTLNotFound()

    return Stream(data[result:])


def parse_tracelogging_event(stream: Stream) -> Event:
    """
    A tracelogging event is identified by its channel number_of_channel
    that are always 11. Actually we can't handle tracelogging event
    because the lonk between event and provider is made during code execution

    :ivar stream: current stream use to parse the event
    :ret: An event object for tracelogging
    """
    channel = stream.read_u8()
    if channel != 11:
        raise ETWBreakerUnexpectedToken(11, channel)
    level = stream.read_u8()
    opcode = stream.read_u8()
    keyword = stream.read_u64()
    size = stream.read_u16()
    stream.read(size - 2)
    return Event(0, 0, channel, level, opcode, 0, keyword)


def parse_tracelogging_provider(stream: Stream) -> Provider:
    """
    Create a default provider for tracelogging
    It will add a default event for this provider
    Because in traclogging all event have the event id set to 0

    :ivar stream: current stream use to parse the provider
    :ret: A provider object for tracelogging
    """
    guid = Guid(stream.read(16))
    size = stream.read_u16()
    payload = stream.read(size - 2)
    name = payload[:payload.find(b"\x00")].decode("ascii")

    return TraceLoggingProvider(guid, [Event(0, 0, 11, 0, 0, 0, 0)], [Channel(11, name)])


def parse_tracelogging(stream: Stream) -> List[Provider]:
    """
    Try to parse all tracelogging event and provider
    from an .rdata segmant

    Actually only provider are intersting. It's because the link
    between event and provider are made into the code dynamically.

    :ivar stream: current stream use to parse the event
    :ret: the list of all provider which are found
    """
    magic = stream.read(4)
    if magic != b"ETW0":
        raise ETWBreakerUnexpectedToken(b"ETW0", magic)

    stream.read(12)
    providers = []
    while True:
        type = stream.read_u8()
        if type == 6:
            parse_tracelogging_event(stream)
        elif type == 4:
            providers.append(parse_tracelogging_provider(stream))
        elif type == 0:
            # padding
            continue
        else:
            print("Unknown Trace logging type %s, expect to be the end of trace logging block"%type)
            break
    return providers


def parse_event_elements(stream: Stream) -> List[Event]:
    """
    Parse an event element
    An event is defined by :
    * unique identifier
    * a channel
    * a set of keywords
    * a level

    :ivar stream: Input stream once read the EVNT magic and the size of the payload
    :ret: List of all event parsed
    """
    number_of_event = stream.read_u32()
    stream.read(4) # padding

    events = []
    for i in range(0, number_of_event):
        event_id = stream.read_u16()
        version = stream.read_u8()
        channel = stream.read_u8()
        level = stream.read_u8()
        opcode = stream.read_u8()
        task = stream.read_u16()
        keywords = stream.read_u64()
        message_identifier = stream.read_u32()
        template_offset = stream.read_u32()
        opcode_offset = stream.read_u32()
        level_offset = stream.read_u32()
        task_offset = stream.read_u32()
        stream.read(12)
        events.append(Event(event_id, version, channel, level, opcode, task, keywords))
    return events


def parse_channel_element(stream: Stream) -> List[Channel] :
    number_of_channel = stream.read_u32()
    result = []
    for i in range(0, number_of_channel):
        unknown = stream.read_u32()
        offset = stream.read_u32()
        identifier = stream.read_u32()
        message_identifier = stream.read_u32()

        sub_stream = Stream(stream.getvalue())
        sub_stream.read(offset)
        size = sub_stream.read_u32()
        name = sub_stream.read(size-4).decode("utf-16le")
        result.append(Channel(identifier, name))
    return result


def parse_event_provider(guid: Guid, stream: Stream) -> Provider:
    """
    Parse an event provider
    An event provider is composed by a plenty of sort of element:
    * EVNT for event

    https://github.com/libyal/libfwevt/blob/master/libfwevt/fwevt_template.h

    :ivar guid: GUID of the provider
    :ivar stream: stream of the entire resource with offset set to the start of the provider
    """
    magic = stream.read(4)
    if magic != b"WEVT":
        raise ETWBreakerUnexpectedToken(b"WEVT", magic)

    size = stream.read_u32()
    message_table_id = stream.read_u32()

    number_of_element = stream.read_u32()
    number_of_unknown = stream.read_u32()

    element_descriptor = [(stream.read_u32(), stream.read_u32()) for i in range(0, number_of_element)]
    unknown = [stream.read_u32() for i in range(0, number_of_unknown)]

    events = []
    channels = []
    for offset, _ in element_descriptor:
        stream.seek(offset)
        magic = stream.read(4)
        size = stream.read_u32()

        # Event declaration
        if magic == b"EVNT":
            events = parse_event_elements(stream)
        elif magic == b"CHAN":
            channels = parse_channel_element(stream)

    return ManifestProvider(guid, events, channels)


def parse_manifest(stream: Stream) -> List[Provider]:
    """
    An ETW Manifest is a binary serialized
    It start with CRIM magic

    Then list all providers
    For each providers we can parse GUID and Provider description

    """
    magic = stream.read(4)
    if magic != b"CRIM":
        raise ETWBreakerUnexpectedToken(b"CRIM", magic)

    size = stream.read_u32()

    major_version = stream.read_u16()
    minor_version = stream.read_u16()

    number_of_provider_descriptor = stream.read_u32()

    # Read provider meta informations
    providers_descriptor = [(Guid(stream.read(16)), stream.read_u32()) for i in range(0, number_of_provider_descriptor)]

    # Parse providers
    providers = []
    for guid, offset in providers_descriptor:
        stream.seek(offset)
        providers.append(parse_event_provider(guid, stream))

    return providers


def add_breakpoint(guid: Guid, event: Event):
    """
    Add a software break point on ntdll!EtwEventWrite
    And set a condition on event id and event provider
    """
    bpt = idaapi.bpt_t()
    bpt.set_sym_bpt("ntdll_EtwEventWrite", 0)
    bpt.condition = textwrap.dedent("""
    import idaapi
    import idc

    rdx = idaapi.regval_t()
    idaapi.get_reg_val('RDX',rdx)
    event_id = int.from_bytes(idc.get_bytes(rdx.ival, 2), "little")

    rcx = idaapi.regval_t()
    idaapi.get_reg_val('RCX',rcx)
    provider_guid = idc.get_bytes((rcx.ival & 0xFFFFFFFFFFFF) + 0x20, 16)

    if event_id == %s and provider_guid == %s:
        print(f"[ETWBreaker] break on Provider {{%s}} EventId ({event_id})")
        return True
    else:
        return False
    """%(event.event_id, guid.raw, guid))
    bpt.elang = "Python"
    idaapi.add_bpt(bpt)


def delete_breakpoint(symbol: str):
    """
    Delete the breakpoint set on ntdll_EtwEventWrite
    """
    location = idaapi.bpt_location_t()
    location.set_sym_bpt(symbol)

    if idaapi.find_bpt(location, None):
        idaapi.del_bpt(location)


class ETWResultsModel(QtCore.QAbstractTableModel):
    """
    This class is QT class that help to view data from COM parsing
    """
    COL_ID = 0x00
    COL_TYPE = 0x01
    COL_GUID = 0x02
    COL_CHANNEL = 0x03
    COL_SYMBOL = 0x04


    def __init__(self, providers: List[Provider], parent=None):
        super().__init__(parent)

        self._column_headers = {
            ETWResultsModel.COL_ID : 'Event ID',
            ETWResultsModel.COL_TYPE : 'Type',
            ETWResultsModel.COL_GUID : 'GUID',
            ETWResultsModel.COL_CHANNEL : 'Channel',
            ETWResultsModel.COL_SYMBOL : 'Symbol'
        }

        self._results = []
        for provider in providers:
            self._results += [(provider, event) for event in provider.events]

        self._row_count = len(self._results)

    def flags(self, index):
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable

    def rowCount(self, index=QtCore.QModelIndex()):
        return self._row_count

    def columnCount(self, index=QtCore.QModelIndex()):
        return len(self._column_headers)

    def headerData(self, column, orientation, role=QtCore.Qt.DisplayRole):
        """
        Define the properties of the the table rows & columns.
        """
        if orientation == QtCore.Qt.Horizontal:

            # the title of the header columns has been requested
            if role == QtCore.Qt.DisplayRole:
                try:
                    return self._column_headers[column]
                except KeyError as e:
                    pass

            # the text alignment of the header has beeen requested
            elif role == QtCore.Qt.TextAlignmentRole:

                # center align all columns
                return QtCore.Qt.AlignHCenter

        # unhandled header request
        return None

    def data(self, index, role=QtCore.Qt.DisplayRole):
        """
        Define how Qt should access the underlying model data.
        """
        # data display request
        if role == QtCore.Qt.DisplayRole:

            # grab for speed
            row = index.row()
            column = index.column()

            if column == ETWResultsModel.COL_GUID:
                return "{%s}"%(self._results[row][0].guid)
            elif column == ETWResultsModel.COL_ID:
                return self._results[row][1].event_id
            elif column == ETWResultsModel.COL_CHANNEL:
                event = self._results[row][1]
                return str(self._results[row][0].find_channel(event.channel) or "Unknown channel")
            elif column == ETWResultsModel.COL_TYPE:
                return self._results[row][0].__class__.__name__
            elif column == ETWResultsModel.COL_SYMBOL:
                return self._results[row][1].find_symbol() or "No symbol"

        # font color request
        elif role == QtCore.Qt.ForegroundRole:
            return QtGui.QColor(QtCore.Qt.black)

        # unhandeled request, nothing to do
        return None


class ETWResultsForm(idaapi.PluginForm):
    """
    The Qt form use to display ETW table
    """
    def __init__(self, providers: List[Provider]):

        super().__init__()
        self.providers = providers

    def OnCreate(self, form):
        """
        Initialize the custom PyQt5 content on form creation.
        """
        # Get parent widget
        self._widget  = self.FormToPyQtWidget(form)
        self._init_ui()

    def show(self):
        """
        Make the created form visible as a tabbed view.
        """
        flags = idaapi.PluginForm.WOPN_TAB | idaapi.PluginForm.WOPN_PERSIST
        return idaapi.PluginForm.Show(self, "ETW", flags)


    def _init_ui(self):
        """
        Init ui of ETW table
        """
        self._model = ETWResultsModel(self.providers, self._widget)
        self._table = QtWidgets.QTableView()

        # set these properties so the user can arbitrarily shrink the table
        self._table.setMinimumHeight(0)
        self._table.setSizePolicy(
            QtWidgets.QSizePolicy.Ignored,
            QtWidgets.QSizePolicy.Ignored
        )

        self._table.setModel(self._model)

        # set a windbg breakpoint on double click
        self._table.doubleClicked.connect(self._ui_entry_double_click)

        # table selection should be by row, not by cell
        self._table.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)

        # more code-friendly, readable aliases
        vh = self._table.verticalHeader()
        hh = self._table.horizontalHeader()
        vh.setSectionResizeMode(QtWidgets.QHeaderView.Fixed)

        # hide the vertical header themselves as we don't need them
        vh.hide()

        # Allow multiline cells
        self._table.setWordWrap(True)
        self._table.setTextElideMode(QtCore.Qt.ElideMiddle);
        self._table.resizeColumnsToContents()
        self._table.resizeRowsToContents()

        layout = QtWidgets.QGridLayout()
        layout.addWidget(self._table)
        self._widget.setLayout(layout)

    def _ui_entry_double_click(self, index):
        """
        When user click on an event
        we send to windbg a special crafted debug command

        That will set a conditional breakpoint on ntdll!EtwEventWrite function
        with condition on function argument that match the eventid and the provider GUID
        """
        event = self._model._results[index.row()][1]
        guid = self._model._results[index.row()][0].guid
        delete_breakpoint("ntdll_EtwEventWrite")
        add_breakpoint(guid, event)


def PLUGIN_ENTRY():
    return ETWBreaker()


repository: Martyx00/CollaRE
file: plugins/ida/CollaRE.py
description: Multi-tool reverse engineering collaboration solution.
wanted_name: CollaRE
comment: CollaRE
created_at: 2021-02-14 07:44:28+00:00
pushed_at: 2024-03-27 12:04:56+00:00
url:
https://github.com/Martyx00/CollaRE/blob/a4f6c7b4c5041d9de4196290dd757c9aea439d19/plugins/ida/CollaRE.py

content:
from idautils import *
from idaapi import *
import idaapi
from idc import *
import ida_nalt
import os, json

def get_comment(ea):
    comment = ""
    repComment = get_cmt(ea, True)
    nonRepComment = get_cmt(ea, False)
    repFunComment = get_func_cmt(ea,True)
    nonRepFunComment = get_func_cmt(ea,False)
    if repComment:
        comment += repComment
    if nonRepComment:
        if comment:
            comment += "; "
        comment += nonRepComment
    if repFunComment:
        if comment:
            comment += "; "
        comment += repFunComment
    if nonRepFunComment:
        if comment:
            comment += "; "
        comment += nonRepFunComment
    return comment

def clear_comments(ea):
    set_cmt(ea,"",False)
    set_cmt(ea,"",True)
    set_func_cmt(ea,"",False)
    set_func_cmt(ea,"",True)
    

class CollaREExportAction(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
    def activate(self, ctx):
        print("[CollaRE] Exporting ...")
        if ".collare_projects" in ida_nalt.get_input_file_path():
            changes = {"function_names":{},"comments":{},"base": int(idaapi.get_imagebase())}
            for segea in Segments():
                for funcea in Functions(segea):
                    # Name
                    functionName = get_func_name(funcea)
                    if hex(funcea)[2:].upper() not in functionName:
                        changes["function_names"][int(funcea)] = {"name":functionName,"end":0}
                for ea in range(segea,get_segm_end(segea)):
                    comment = get_comment(ea)
                    if comment:
                        changes["comments"][ea] = comment 
                
                with open(os.path.join(os.path.dirname(ida_nalt.get_input_file_path()),"changes.json"),"w") as changes_file:
                    json.dump(changes,changes_file)

            print("[CollaRE] Export completed!")
            idaapi.info("CollaRE Export completed!")
        else:
            print("[CollaRE] Export failed! Not a CollaRE project!")
            idaapi.warning("This is not a CollaRE project!")
        return 1
    
    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS
    
class CollaREImportAction(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        if ".collare_projects" in ida_nalt.get_input_file_path():
            with open(os.path.join(os.path.dirname(ida_nalt.get_input_file_path()),"changes.json"),"r") as changes_file:
                changes = json.load(changes_file)
                base = changes["base"]
                if base != int(idaapi.get_imagebase()):
                    base = int(idaapi.get_imagebase()) - base
                else:
                    base = 0
                for function in changes["function_names"]:
                    # Set function names
                    function_address = int(function) + base
                    idaapi.set_name(function_address,str(changes["function_names"][function]["name"]),idaapi.SN_FORCE)
                for comment in changes["comments"]:
                    comment_address = int(comment,10) + base
                    currentComment = get_comment(comment_address)
                    clear_comments(comment_address)
                    if currentComment:
                        if currentComment in changes["comments"][comment]:
                            set_cmt(comment_address,changes["comments"][comment],False)
                        elif changes["comments"][comment] in currentComment:
                            set_cmt(comment_address,currentComment,False)
                        else:
                            set_cmt(comment_address,currentComment + "; " + changes["comments"][comment],False)
                    else:
                        set_cmt(comment_address,changes["comments"][comment],False)
            print("[*] Import completed!")
            idaapi.info("CollaRE Import completed!")
        else:
            print("[!] This is not a CollaRE project!")
            idaapi.warning("This is not a CollaRE project!")
        return 1
        
    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


class collare_t(idaapi.plugin_t):
    comment = "CollaRE"
    help = "Use this plugin to share function names and comments across differrent tool within CollaRE workspace."
    wanted_name = "CollaRE"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        export_desc = idaapi.action_desc_t(
            'collare:export',
            'CollaRE Export',
            CollaREExportAction(),
            '',
            'Export function names and comments for CollaRE projects.'
            )
        
        idaapi.register_action(export_desc)
        idaapi.attach_action_to_menu("File/Save as...", "collare:export", idaapi.SETMENU_APP)
        import_desc = idaapi.action_desc_t(
            'collare:import',
            'CollaRE Import',
            CollaREImportAction(),
            '',
            'Import function names and comments for CollaRE projects.'
            )           
        
        idaapi.register_action(import_desc)
        idaapi.attach_action_to_menu("File/Save as...", "collare:import", idaapi.SETMENU_APP)
        
    def run(self):
        pass

    def term(self):
        pass

def PLUGIN_ENTRY():
    return collare_t()


repository: Synacktiv-contrib/kernelcache-laundering
file: aarch64_pac.py
description: load iOS12 kernelcaches and PAC code in IDA
wanted_name: Aarch64 PAC
comment: ARM v8.3-A Pointer Authentication extension
created_at: 2018-10-01 08:59:35+00:00
pushed_at: 2018-10-04 09:01:04+00:00
url:
https://github.com/Synacktiv-contrib/kernelcache-laundering/blob/6fde371d3cf04e625d24f5155227bfa017e1cd11/aarch64_pac.py

content:
# AArch64 8.3-A Pointer Authentication extension
# Enhanced and Python version of xerub C++ plugin available there: https://github.com/xerub/idastuff/blob/master/arm64/aarch64_pac

# Copyright (c) 2018 Eloi Benoist-Vanderbeken - Synacktiv
# Copyright (c) 2018 xerub

# This program is free software; you can redistribute it and/or 
# modify it under the terms of the GNU General Public License version 
# 2 as published by the Free Software Foundation. 

# 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, see <http://www.gnu.org/licenses/>.

import idaapi
import ida_hexrays

PAC_NONE = 0
PAC_PACIASP = PAC_NONE + 1
PAC_PACIBSP = PAC_PACIASP + 1
PAC_AUTIASP = PAC_PACIBSP + 1
PAC_AUTIBSP = PAC_AUTIASP + 1
PAC_PACIAZ = PAC_AUTIBSP + 1
PAC_PACIBZ = PAC_PACIAZ + 1
PAC_AUTIAZ = PAC_PACIBZ + 1
PAC_AUTIBZ = PAC_AUTIAZ + 1
PAC_PACIA1716 = PAC_AUTIBZ + 1
PAC_PACIB1716 = PAC_PACIA1716 + 1
PAC_AUTIA1716 = PAC_PACIB1716 + 1
PAC_AUTIB1716 = PAC_AUTIA1716 + 1
PAC_PACIA = PAC_AUTIB1716 + 1
PAC_PACIB = PAC_PACIA + 1
PAC_PACDA = PAC_PACIB + 1
PAC_PACDB = PAC_PACDA + 1
PAC_AUTIA = PAC_PACDB + 1
PAC_AUTIB = PAC_AUTIA + 1
PAC_AUTDA = PAC_AUTIB + 1
PAC_AUTDB = PAC_AUTDA + 1
PAC_PACIZA = PAC_AUTDB + 1
PAC_PACIZB = PAC_PACIZA + 1
PAC_PACDZA = PAC_PACIZB + 1
PAC_PACDZB = PAC_PACDZA + 1
PAC_AUTIZA = PAC_PACDZB + 1
PAC_AUTIZB = PAC_AUTIZA + 1
PAC_AUTDZA = PAC_AUTIZB + 1
PAC_AUTDZB = PAC_AUTDZA + 1
PAC_PACGA = PAC_AUTDZB + 1
PAC_XPACLRI = PAC_PACGA + 1
PAC_XPACI = PAC_XPACLRI + 1
PAC_XPACD = PAC_XPACI + 1
PAC_RETAA = PAC_XPACD + 1
PAC_RETAB = PAC_RETAA + 1
PAC_BRAA = PAC_RETAB + 1
PAC_BRAB = PAC_BRAA + 1
PAC_BRAAZ = PAC_BRAB + 1
PAC_BRABZ = PAC_BRAAZ + 1
PAC_BLRAA = PAC_BRABZ + 1
PAC_BLRAB = PAC_BLRAA + 1
PAC_BLRAAZ = PAC_BLRAB + 1
PAC_BLRABZ = PAC_BLRAAZ + 1
PAC_ERETAA = PAC_BLRABZ + 1
PAC_ERETAB = PAC_ERETAA + 1
PAC_LDRAA = PAC_ERETAB + 1
PAC_LDRAB = PAC_LDRAA + 1

OP_NAMES = {
	PAC_PACIASP: "PACIASP",
	PAC_PACIBSP: "PACIBSP",
	PAC_AUTIASP: "AUTIASP",
	PAC_AUTIBSP: "AUTIBSP",
	PAC_PACIAZ: "PACIAZ",
	PAC_PACIBZ: "PACIBZ",
	PAC_AUTIAZ: "AUTIAZ",
	PAC_AUTIBZ: "AUTIBZ",
	PAC_PACIA1716: "PACIA1716",
	PAC_PACIB1716: "PACIB1716",
	PAC_AUTIA1716: "AUTIA1716",
	PAC_AUTIB1716: "AUTIB1716",
	PAC_PACIA: "PACIA",
	PAC_PACIB: "PACIB",
	PAC_PACDA: "PACDA",
	PAC_PACDB: "PACDB",
	PAC_AUTIA: "AUTIA",
	PAC_AUTIB: "AUTIB",
	PAC_AUTDA: "AUTDA",
	PAC_AUTDB: "AUTDB",
	PAC_PACIZA: "PACIZA",
	PAC_PACIZB: "PACIZB",
	PAC_PACDZA: "PACDZA",
	PAC_PACDZB: "PACDZB",
	PAC_AUTIZA: "AUTIZA",
	PAC_AUTIZB: "AUTIZB",
	PAC_AUTDZA: "AUTDZA",
	PAC_AUTDZB: "AUTDZB",
	PAC_PACGA: "PACGA",
	PAC_XPACLRI: "XPACLRI",
	PAC_XPACI: "XPACI",
	PAC_XPACD: "XPACD",
	PAC_RETAA: "RETAA",
	PAC_RETAB: "RETAB",
	PAC_BRAA: "BRAA",
	PAC_BRAB: "BRAB",
	PAC_BRAAZ: "BRAAZ",
	PAC_BRABZ: "BRABZ",
	PAC_BLRAA: "BLRAA",
	PAC_BLRAB: "BLRAB",
	PAC_BLRAAZ: "BLRAAZ",
	PAC_BLRABZ: "BLRABZ",
	PAC_ERETAA: "ERETAA",
	PAC_ERETAB: "ERETAB",
	PAC_LDRAA: "LDRAA",
	PAC_LDRAB: "LDRAB"
}

def decode_PAC(d, insn):
	if (d & 0xffffc000) == 0xdac10000: 
		m = (d >> 10) & 7
		Z = (d >> 13) & 1
		Xn = (d >> 5) & 0x1F
		Xd = d & 0x1F
		if Z == 0: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xd + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.Op3.type = idaapi.o_reg
			insn.Op3.reg = Xd + 129
			insn.Op3.dtype = idaapi.dt_qword
			insn.Op3.flags = 0
			insn.Op2.type = idaapi.o_reg
			insn.Op2.reg = Xn + 129
			insn.Op2.dtype = idaapi.dt_qword
			insn.Op2.flags = idaapi.OF_SHOW
			insn.insnpref = PAC_PACIA + m
			insn.size = 4
			return True
		elif Xn == 31: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xd + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.Op3.type = idaapi.o_reg
			insn.Op3.reg = Xd + 129
			insn.Op3.dtype = idaapi.dt_qword
			insn.Op3.flags = 0
			insn.insnpref = PAC_PACIZA + m
			insn.size = 4
			return True
	if (d & 0xfffffd1f) == 0xd503211f: 
		m = (d >> 6) & 3
		CRm = (d >> 9) & 1
		op2 = (d >> 5) & 1
		if CRm == 0: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_void
			insn.insnpref = PAC_PACIA1716 + m
		elif op2: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_void
			insn.insnpref = PAC_PACIASP + m
		else:
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_void
			insn.insnpref = PAC_PACIAZ + m
		insn.size = 4
		return True
	if (d & 0xffe0fc00) == 0x9ac03000: 
		Xm = (d >> 16) & 0x1F
		Xn = (d >> 5) & 0x1F
		Xd = d & 0x1F
		insn.itype = idaapi.ARM_hlt
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = Xd + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op2.type = idaapi.o_reg
		insn.Op2.reg = Xn + 129
		insn.Op2.dtype = idaapi.dt_qword
		insn.Op3.type = idaapi.o_reg
		insn.Op3.reg = Xm + 129
		insn.Op3.dtype = idaapi.dt_qword
		insn.insnpref = PAC_PACGA
		insn.size = 4
		return True
	if (d & 0xfffffbe0) == 0xdac143e0: 
		D = (d >> 10) & 1
		Xd = d & 0x1F
		insn.itype = idaapi.ARM_hlt
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = Xd + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op3.type = idaapi.o_reg
		insn.Op3.reg = Xd + 129
		insn.Op3.dtype = idaapi.dt_qword
		insn.Op3.flags = 0
		insn.insnpref = PAC_XPACI + D
		insn.size = 4
		return True
	if d == 0xd50320ff: 
		insn.itype = idaapi.ARM_hlt
		insn.segpref = 14
		insn.Op1.type = idaapi.o_void
		insn.insnpref = PAC_XPACLRI
		insn.size = 4
		return True
	if (d & 0xfffffbff) == 0xd65f0bff: 
		M = (d >> 10) & 1
		insn.insnpref = PAC_RETAA + M
		insn.itype = idaapi.ARM_ret
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = 30 + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op1.flags = 0
		insn.size = 4
		return True
	if (d & 0xfedff800) == 0xd61f0800: 
		is_blr = (d >> 19) & 4
		Z = (d >> 24) & 1
		M = (d >> 10) & 1
		Xn = (d >> 5) & 0x1F
		Xm = d & 0x1F
		if Z == 0 and Xm == 31: 
			insn.itype =  idaapi.ARM_blr if is_blr else idaapi.ARM_br
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xn + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.insnpref = PAC_BRAAZ + M + is_blr
			insn.size = 4
			return True
		elif Z: 
			insn.itype = idaapi.ARM_blr if is_blr else idaapi.ARM_br
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xn + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.Op2.type = idaapi.o_reg
			insn.Op2.reg = Xm + 129
			insn.Op2.dtype = idaapi.dt_qword
			insn.Op2.flags = idaapi.OF_SHOW
			insn.insnpref = PAC_BRAA + M + is_blr
			insn.size = 4
			return True
	if (d & 0xfffffbff) == 0xd69f0bff: 
		M = (d >> 10) & 1
		insn.insnpref = PAC_ERETAA + M
		insn.itype = idaapi.ARM_eret
		insn.segpref = 14
		insn.size = 4
		return True
	if (d & 0xff200400) == 0xf8200400: 
		M = (d >> 23) & 1
		imm10 = ((d & 0x400000) << 9) | ((d & 0x1ff000) << 10)
		offset = imm10 >> 19
		W = (d >> 11) & 1
		Xn = (d >> 5) & 0x1F
		Xt = d & 0x1F
		insn.itype = idaapi.ARM_ldr
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = Xt + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op2.type = idaapi.o_displ
		insn.Op2.reg = Xn + 129
		insn.Op2.dtype = idaapi.dt_qword
		insn.Op2.addr = offset
		if W != 0: 
			insn.auxpref = 0x20
		insn.insnpref = PAC_LDRAA + M
		insn.size = 4
		return True
	return False

class Aarch64PACHook(idaapi.IDP_Hooks):
	CUSTOM_INSTRUCTIONS = {idaapi.ARM_hlt, idaapi.ARM_ret, idaapi.ARM_blr, idaapi.ARM_br, idaapi.ARM_eret,idaapi.ARM_ldr}
	INDENT = 16
	def ev_ana_insn(self, outctx):
		return outctx.size if decode_PAC(idaapi.get_dword(outctx.ea), outctx) else 0

	def ev_emu_insn(self, insn):
		if insn.itype != idaapi.ARM_brk:
			return False
		return True

	def ev_out_mnem(self, outctx):
		if outctx.insn.itype in self.CUSTOM_INSTRUCTIONS:
			mnem = OP_NAMES.get(ord(outctx.insn.insnpref), None)
			if mnem is not None:
				outctx.out_custom_mnem(mnem, self.INDENT)
				return 1
		return 0

class Aarch64PACPlugin(idaapi.plugin_t):
	flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
	comment = "ARM v8.3-A Pointer Authentication extension"
	wanted_hotkey = ""
	help = "Runs transparently"
	wanted_name = "Aarch64 PAC"
	hook = None
	enabled = 1

	def init(self):
		if idaapi.ph_get_id() != idaapi.PLFM_ARM or idaapi.BADADDR <= 0xFFFFFFFF:
			return idaapi.PLUGIN_SKIP
		print "%s init"%self.comment
		self.hook = Aarch64PACHook()
		self.hook.hook()
		return idaapi.PLUGIN_KEEP

	def run():
		pass

	def term(self):
		if self.hook is not None:
			self.hook.unhook()
		print "%s unloaded"%self.comment

def PLUGIN_ENTRY():
	return Aarch64PACPlugin()


repository: Asteria-BCSD/Asteria
file: ASTExtraction/Asteria_ida_plugin.py
wanted_name: Asteria
comment: This is a plugin for the Asteria (https://github.com/Asteria-BCSD/Asteria)
created_at: 2020-06-14 02:38:15+00:00
pushed_at: 2023-10-17 08:34:58+00:00
url:
https://github.com/Asteria-BCSD/Asteria/blob/04e95084ad2d4a92e77d38004dc56f59410b8f57/ASTExtraction/Asteria_ida_plugin.py

content:
# encoding=utf-8
'''

This is a ida plugin script
## 1
To dump all/one AST feature(s) of function(s) to a database file.
## 2
To Calculate Similarity between all functions and functions in another database file.
'''

import idaapi
from idaapi import IDA_SDK_VERSION, IDAViewWrapper
import idc
import idautils
import logging
from ida_kernwin import Form, PluginForm
import os, sys
import json
import commands
import subprocess
import time

try:
    if IDA_SDK_VERSION < 690:
        # In versions prior to IDA 6.9 PySide is used...
        from PySide import QtGui

        QtWidgets = QtGui
        is_pyqt5 = False
    else:
        # ...while in IDA 6.9, they switched to PyQt5
        from PyQt5 import QtCore, QtGui, QtWidgets

        is_pyqt5 = True
except ImportError:
    pass

l = logging.getLogger("ast_generator.py")
l.setLevel(logging.ERROR)
logger = logging.getLogger('Asteria')
logger.addHandler(logging.StreamHandler())
logger.addHandler(logging.FileHandler("Asteria.log"))
logger.handlers[0].setFormatter(
    logging.Formatter("[%(filename)s][%(levelname)s] %(message)s\t(%(module)s:%(funcName)s)"))
logger.handlers[1].setFormatter(
    logging.Formatter("[%(filename)s][%(levelname)s] %(message)s\t(%(module)s:%(funcName)s)"))
logger.setLevel(logging.ERROR)


# logging.basicConfig(format='[%(filename)s][%(levelname)s] %(message)s\t(%(module)s:%(funcName)s)',
#                     filename="Asteria.log",
#                     filemode='a')

# view to show the similaity results
class CHtmlViewer(PluginForm):

    def OnCreate(self, form):
        if is_pyqt5:
            self.parent = self.FormToPyQtWidget(form)
        else:
            self.parent = self.FormToPySideWidget(form)
        self.PopulateForm()

        self.browser = None
        self.layout = None
        self.text = ""

        return 1

    def PopulateForm(self):
        self.layout = QtWidgets.QVBoxLayout()
        self.browser = QtWidgets.QTextBrowser()
        self.browser.setLineWrapMode(QtWidgets.QTextEdit.NoWrap)
        self.browser.setHtml(self.text)
        self.browser.setReadOnly(True)
        self.browser.setFontWeight(18)
        self.layout.addWidget(self.browser)
        self.parent.setLayout(self.layout)

    def ScoreToColor(self, score):
        '''
        return the color rgb value according to score
        1. score ranges 0.8~1.0
        2. color format #ff{A}{A} , A is calculated according to Score
        '''
        score = float(score)
        if score < 0.8:
            return "#ffffff"  # white

        ColorScore = int(125 - (score - 0.8) * 10 * 62)
        return "#ff%.2x%.2x" % (ColorScore, ColorScore)

    def AssembleHtml(self, jsondata):
        '''
        jsondata :dic: {'vulFuncName':{'info':['funcName','BinPath','BinName'], "rank":[[[like info], score], [[like info], score]]}}
        '''
        _html_template = """
                              <html>
                              <head>
                              <style>%(style)s</style>
                              </head>
                              <body>
                              <table class="diff_tab"  cellspacing="3px">
                              %(rows)s
                              </table>
                              </body>
                              </html>
        """
        _style = """
                 table.diff_tab {
                    font-family: Courier monospace;
                    border: solid black;
                    font-size: 15px;

                    text-align: left;
                    width: 100%;
                  }
                  td {
                    border: 2px;
                    border-bottom-style:solid;
                  }
                  """

        head_row_temp = """
                    <tr> 
                    <td class="VulFuncName" rowspan="%(rowNumber)d"> %(vulnerableFunc)s </td>
                    <td class="CandidateName"> %(candidateFunc)s </td>
                    <td class="score" bgcolor=%(bgcolor)s> %(score)s </td>
                    </tr>
        """
        other_row_temp = """
                    <tr> 
                    <td > %(candidateFunc)s </td>
                    <td bgcolor=%(bgcolor)s> %(score)s </td>
                    </tr>
        """
        table_header = """
            <thead> 
                <tr> 
                <th>Functions In Chosen File</th> <th>Functions In IDA</th>  <th>Similarity Scores</th> 
                </tr> 
            </thead>
        """
        table_content = ""
        for vulFunc in jsondata:
            candidates = jsondata[vulFunc]['rank']
            rowNumber = len(candidates)
            if rowNumber == 0:
                continue
            vulInfo = jsondata[vulFunc]['info']
            vulFuncName = vulFunc
            vulFuncBin = vulInfo[2]

            headerRow = head_row_temp % {"rowNumber": rowNumber, "vulnerableFunc": vulFuncName,
                                         "candidateFunc": candidates[0][0][0], "score": candidates[0][1],
                                         "bgcolor": self.ScoreToColor(candidates[0][1])}
            table_content += headerRow
            if rowNumber > 1:
                for candidate in candidates[1:]:
                    table_content += other_row_temp % {"candidateFunc": candidate[0][0], "score": candidate[1],
                                                       "bgcolor": self.ScoreToColor(candidate[1])}

        src = _html_template % {"style": _style, "rows": table_header + table_content}
        logger.debug(src)
        return src

    def Show(self, jsondata, title):
        self.text = self.AssembleHtml(jsondata)
        """Creates the form is not created or focuses it if it was"""
        return PluginForm.Show(self, title, options=PluginForm.FORM_PERSIST)


# function menus
class MyForm(Form):
    # Main Class for functionality of plugin
    def __init__(self):
        self.invert = False
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""STARTITEM 0
            Please choose your option.
            <##Function Feature Generation:{iButton1}> <##Function Similarity:{iButton2}>
            """,
            {
                'iButton1': F.ButtonInput(self.db_generate),
                'iButton2': F.ButtonInput(self.sim_cal)
            }
        )
        self._sqlitefilepath = idc.GetIdbPath() + ".sqlite"
        idadir = idaapi.idadir("python")
        self.mainapp_path = os.path.join(idadir, "ASTExtraction",
                                         "main_app.py")  # path to the script of the ast encoding and similarity
        if not os.path.exists(self.mainapp_path):
            logger.error("Python File {} not exists!".format(self.mainapp_path))
        self.application = os.path.join(idadir, "ASTExtraction",
                                        "application.py")  # path to the script of the ast encoding and similarity
        if not os.path.exists(self.application):
            logger.error("Python File {} not exists!".format(self.application))

    def db_generate(self, code=0):
        logger.debug("ast generating...")
        # 1. set file to save; 2. show waiting box and progress 3. invoke ida script for ast generation 4. show sucess box.
        # 1.
        logger.debug("sqlite file path {}".format(self._sqlitefilepath))

        # 2.
        ida_kernwin.show_wait_box("Processing...")

        # 3.
        # 3.1 extract asts
        try:
            idaapi.require("ASTExtraction")
            idaapi.require("ASTExtraction.ast_generator")
            idaapi.require("ASTExtraction.DbOp")
            g = ASTExtraction.ast_generator.AstGenerator()
            dbop = ASTExtraction.DbOp.DBOP(self._sqlitefilepath)
            g.run(g.get_info_of_func)
            # 3.2 save asts to database
            ida_kernwin.replace_wait_box("Processing. Saving to Database.")
            recordsNo = g.save_to(dbop)
            logger.info("%d records are inserted into database." % recordsNo)
            del dbop
        except Exception, e:
            ida_kernwin.replace_wait_box(str(e))
            logger.error("ASTExtraction import error! {}".format(e))
            time.sleep(2)
        # 4.
        ida_kernwin.replace_wait_box("AST Encoding...")

        cmd = 'python "{}" --dbpath "{}"'.format(self.application, self._sqlitefilepath)
        idaapi.msg("[Asteria] >>> AST Encoding...[{}]\n".format(cmd))

        returncode = self.invoke_system_python(cmd, hidden_window=False)
        if returncode:
            idaapi.msg("[Asteria] >>> AST Encoding failed\n")
        ida_kernwin.hide_wait_box()
        idaapi.msg("[Asteria] >>> AST Encoding Finished\n")

    def invoke_system_python(self, cmd, hidden_window=True):
        '''
        cmd : str: command to be executed
        hidden_window: bool:
        return : subprocess.returncode
        '''
        startupinfo = subprocess.STARTUPINFO()
        # to hidden the cmd window
        if hidden_window:
            if 'win32' in str(sys.platform).lower():
                startupinfo.dwFlags = subprocess.CREATE_NEW_CONSOLE | subprocess.STARTF_USESHOWWINDOW
                startupinfo.wShowWindow = subprocess.SW_HIDE

            # p = subprocess.Popen(cmd,startupinfo= startupinfo, stdout=subprocess.PIPE)#, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0)
            # while p.poll() == None:
            #     idaapi.msg(p.stdout.readline())
            # return p.returncode

        p = subprocess.Popen(cmd,
                             startupinfo=startupinfo)  # , stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0)
        # while p.poll() == None:
        #     idaapi.msg(p.stdout.readline().strip())
        # p.stdout.flush()
        return p.wait()

    def sim_cal(self, code=0):
        logger.debug("Similarity Calculation...")
        # : 1. make sure current idb has generated sqlite db file 2. Choose sqlite db file to compare with 3. do calculation 4. show similarity results.
        # 1.
        if os.path.exists(self._sqlitefilepath):
            # 2.
            sqlite_database2, _ = QtWidgets.QFileDialog.getOpenFileName()  # another sqlite db file for similarity calculation with the sqlite db file generated by current idb
            idaapi.msg("[Asteria] >>> The sqlite file to be calculated is {}\n".format(sqlite_database2))
            # 3. 3.1: TODO: should check the python environment first
            # 3.2: do the calculation and save the results to file

            result_path = self._sqlitefilepath.split(".")[0] + ".json"
            cmd = 'python "{}" --result "{}" "{}" "{}"'.format(self.mainapp_path, result_path, sqlite_database2,
                                                               self._sqlitefilepath)
            logger.info("Calculation CMD: {}".format(cmd))
            ida_kernwin.show_wait_box("Calculating... Please wait.")
            returncode = self.invoke_system_python(cmd, hidden_window=True)
            ida_kernwin.hide_wait_box()
            if returncode:
                # cmd execution fails
                idaapi.msg("[Asteria] >>> Calcualtion failed, {}.\n".format(returncode))
            # 3.3 load the results
            if not os.path.exists(result_path):
                idaapi.msg("[Asteria] >>> Result file does not exists after the calculation, please see the log.\n")
                return

            results = json.load(open(result_path, 'r'))
            # logger.info(str(results))
            # show window of result
            window = CHtmlViewer()
            title = "Similarity Scores"
            window.Show(results, title)
            self.Close(1)
            # self.Free()
        else:
            ida_kernwin.show_wait_box("Please run AST Generation first.")
            while not ida_kernwin.user_cancelled():
                pass
            ida_kernwin.hide_wait_box()
            return


class AsteriaPlugin(idaapi.plugin_t):
    """
    This is the main class of the plugin. It subclasses plugin_t as required
    by IDA. It holds the modules of plugin, which themselves provides the
    functionality of the plugin.
    """
    # plugin information
    flags = idaapi.PLUGIN_UNL
    comment = "This is a plugin for the Asteria (https://github.com/Asteria-BCSD/Asteria)"
    wanted_name = "Asteria"  # name of plugin
    wanted_hotkey = "Alt-F1"  # hot key
    help = "Please see README in https://github.com/Asteria-BCSD/Asteria"

    def init(self):
        if not self.load_plugin_decompiler():
            idaapi.msg("[Asteria]Failed to load hexray plugin.")
            return idaapi.PLUGIN_SKIP

        # To load the modules in /path/to/ida/python/Asteria

        try:
            idaapi.require("ASTExtraction")
            # idaapi.require("ASTExtraction.ast_generator")
            # idaapi.require("ASTExtraction.DbOp")
        except Exception, e:
            idaapi.msg(
                "[Asteria] Plugin initialization failed. Please read the README file and copy 'ASTExtraction' dir of the project to '/path/to/ida/python/'\n")
            logger.error(e.args)
            return idaapi.PLUGIN_SKIP

        idaapi.msg("[Asteria] >>> Asteria plugin initialization finished!\n")
        return idaapi.PLUGIN_OK  # return PLUGIN_KEEP

    def run(self, arg):
        self.form = MyForm()
        self.form.Compile()
        self.form.Execute()
        logger.debug(">>> Asteria plugin exists")

    def term(self):
        idaapi.msg("[Asteria] >>> Asteria plugin ends.\n")

    def load_plugin_decompiler(self):
        '''
        load the hexray plugins
        :return: success or not
        '''
        is_ida64 = idc.GetIdbPath().endswith(".i64")
        if not is_ida64:
            idaapi.load_plugin("hexrays")
            idaapi.load_plugin("hexarm")
        else:
            idaapi.load_plugin("hexx64")
        if not idaapi.init_hexrays_plugin():
            logger.error('[+] decompiler plugins load failed. IDAdb: %s' % idc.GetInputFilePath())
            return False
        return True


def PLUGIN_ENTRY():
    # this function returns an instance of an idapython plugin

    return AsteriaPlugin()


repository: KasperskyLab/Apihashes
file: apihashesv2.py
description: IDA Pro plugin for recognizing known hashes of API function names
wanted_name: Apihashes
comment: Resolve API hashes on code/data creation
created_at: 2022-05-12 13:50:32+00:00
pushed_at: 2022-05-12 13:51:13+00:00
url:
https://github.com/KasperskyLab/Apihashes/blob/71905ff04655a0352614782ebcfb71c54435c656/apihashesv2.py

content:
#!/usr/bin/python3
# © 2022 AO Kaspersky Lab. All Rights Reserved.
#
# Installation: copy all the files and directories in the IDA's plugins/ directory

import ida_idaapi, ida_idp, ida_ua, ida_bytes, ida_diskio, zlib, apihashesv2_search

# Hook class
class hk(ida_idp.IDB_Hooks):
    def CheckHash(self, ea, value):
        fname = apihashesv2_search.FindHash(value)
        if fname:
            print(f"[apihashes] {hex(ea)}: Found API hash for {fname}")
            ida_bytes.set_cmt(ea, fname, False)

    # This hook will check the operands of disassembled instructions
    # So, for older IDBs you may need to undefine and disassemble the code
    # again to make the plugin work
    def make_code(self, insn):
        for op in insn.ops:
            if op.type == ida_ua.o_void:
                break
            if op.type == ida_ua.o_imm and op.value != 0:
                self.CheckHash(insn.ea, op.value)

        return None

    # This hook will check the newly created data items (DWORDS, QWORDS)
    # So, for older IDBs you may need to undefine and recreate the data items
    # to force the checks
    def make_data(self, ea, flags, tid, sz):
        if sz == 4:
            opValue = ida_bytes.get_dword(ea)
        elif sz == 8:
            opValue = ida_bytes.get_qword(ea)
        else:
            return None
        self.CheckHash(ea, opValue)
        return None

class apihashes_plugin_t(ida_idaapi.plugin_t):
    flags = 0
    comment = "Resolve API hashes on code/data creation"
    help = "No help"
    wanted_name = "Apihashes"
    wanted_hotkey = ""
    

    def init(self):
        self.hk = hk()
        self.hk.hook()
        # We're looking for the database in the IDA's plugins directory
        res = apihashesv2_search.LoadHashes(ida_diskio.idadir(ida_diskio.PLG_SUBDIR) + "/apihashesv2.bin")
        print(f"[apihashes] v2 plugin loaded, {res} hashes in the database")

        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        pass


def PLUGIN_ENTRY():
    return apihashes_plugin_t()


repository: rr-/ida-images
file: rgb-ida.py
description: Image preview plugin for IDA disassembler.
url: https://github.com/rr-/ida-images/blob/b194903336e0df76d1d0b50aa82231afe5373e79/rgb-ida.py
wanted_name: Image previewer
comment: Preview memory as image
created_at: 2015-03-16 20:57:44+00:00
pushed_at: 2022-09-17 18:24:56+00:00
content:
import idaapi
import librgb
import ida_kernwin
from librgb.qt_shims import QtGui  # important for PySide legacy IDA
from librgb.qt_shims import QtWidgets

try:
    MAJOR, MINOR = map(int, idaapi.get_kernel_version().split("."))
except AttributeError:
    MAJOR, MINOR = 6, 6
USING_IDA7API = MAJOR > 6
USING_PYQT5 = USING_IDA7API or (MAJOR == 6 and MINOR >= 9)


class DockableShim(object):
    def __init__(self, title):
        self._title = title

        # IDA 7+ Widgets
        if USING_IDA7API:
            import sip

            self._form = idaapi.create_empty_widget(self._title)
            self.widget = sip.wrapinstance(int(self._form), QtWidgets.QWidget)
        # legacy IDA PluginForm's
        else:
            self._form = idaapi.create_tform(self._title, None)
            if USING_PYQT5:
                self.widget = idaapi.PluginForm.FormToPyQtWidget(self._form)
            else:
                self.widget = idaapi.PluginForm.FormToPySideWidget(self._form)

    def show(self):
        if USING_IDA7API:
            flags = (
                idaapi.PluginForm.WOPN_TAB
                | idaapi.PluginForm.WOPN_MENU
                | idaapi.PluginForm.WOPN_RESTORE
                | idaapi.PluginForm.WOPN_PERSIST
            )
            idaapi.display_widget(self._form, flags)

        # legacy IDA PluginForm's
        else:
            flags = (
                idaapi.PluginForm.FORM_TAB
                | idaapi.PluginForm.FORM_MENU
                | idaapi.PluginForm.FORM_RESTORE
                | idaapi.PluginForm.FORM_PERSIST
                | 0x80
            )  # idaapi.PluginForm.FORM_QWIDGET
            idaapi.open_tform(self._form, flags)


class ImagePreviewPlugin(idaapi.plugin_t):
    flags = 0
    wanted_name = "Image previewer"
    wanted_hotkey = "Alt + I"
    comment = "Preview memory as image"
    help = "https://github.com/rr-/ida-images"

    def init(self):
        return idaapi.PLUGIN_OK

    def term(self):
        pass

    def run(self, arg):
        class IdaWindowAdapter(librgb.GenericWindowAdapter):
            def ask_address(self, address):
                return ida_kernwin.ask_addr(address, "Please enter an address")

            def ask_file(self):
                return ida_kernwin.ask_file(1, "*.png", "Save the image as...")

        image_preview_form = DockableShim("Image preview")

        params = librgb.RendererParams()
        params.readers = [librgb.MemoryReader()]
        params.format = librgb.PixelFormats.GRAY8
        params.width = 800
        params.height = 600
        params.flip = False
        params.brightness = 50

        adapter = IdaWindowAdapter(params)
        shortcut_manager = librgb.ShortcutManager(adapter, params)
        for shortcut, func in shortcut_manager.shortcuts.items():
            adapter.define_shortcut(shortcut, image_preview_form.widget, func)

        layout = adapter.create_layout()
        image_preview_form.widget.setLayout(layout)
        adapter.draw()
        image_preview_form.show()


def PLUGIN_ENTRY():
    return ImagePreviewPlugin()


if __name__ == "__main__":
    ImagePreviewPlugin().run(0)


repository: alexhude/uEmu
file: uEmu.py
description: Tiny cute emulator plugin for IDA based on unicorn.
url: https://github.com/alexhude/uEmu/blob/c51605f25ca0d8c9e18aa7c79e2f5d1b24bbf021/uEmu.py
wanted_name: uEmu
created_at: 2017-08-08 10:37:23+00:00
pushed_at: 2024-08-13 10:50:52+00:00
content:
#
#  uEmu.py
#  Micro Emulator
#
#  Created by Alexander Hude on 26/07/17.
#  Copyright (c) 2017 Alexander Hude. All rights reserved.
#

UEMU_USE_AS_SCRIPT      = True    # Set to `False` if you want to load uEmu automatically as IDA Plugin

# === Import

import pickle
import threading
import json
import os
import collections

# IDA Python SDK
from idaapi import *
from idc import *
from idautils import *

if IDA_SDK_VERSION >= 700:
    # functions
    IDAAPI_ScreenEA     = get_screen_ea
    IDAAPI_IsCode       = is_code
    IDAAPI_DelItems     = del_items
    IDAAPI_MakeCode     = create_insn
    IDAAPI_GetFlags     = get_full_flags
    IDAAPI_SetColor     = set_color
    IDAAPI_IsLoaded     = is_loaded
    IDAAPI_HasValue     = has_value
    IDAAPI_GetBptQty    = get_bpt_qty
    IDAAPI_GetBptEA     = get_bpt_ea
    IDAAPI_GetBptAttr   = get_bpt_attr
    IDAAPI_SegStart     = get_segm_start
    IDAAPI_SegEnd       = get_segm_end
    IDAAPI_GetBytes     = get_bytes
    IDAAPI_AskYN        = ask_yn
    IDAAPI_AskFile      = ask_file
    IDAAPI_AskLong      = ask_long
    IDAAPI_NextHead     = next_head
    IDAAPI_GetDisasm    = lambda a, b: tag_remove(generate_disasm_line(a, b))
    IDAAPI_NextThat     = next_that
    IDAAPI_Jump         = jumpto
    # classes
    IDAAPI_Choose       = Choose
else:
    # functions
    IDAAPI_ScreenEA     = ScreenEA
    IDAAPI_IsCode       = isCode
    IDAAPI_DelItems     = MakeUnkn
    IDAAPI_MakeCode     = MakeCode
    IDAAPI_GetFlags     = getFlags
    IDAAPI_SetColor     = SetColor
    IDAAPI_IsLoaded     = isLoaded
    IDAAPI_HasValue     = hasValue
    IDAAPI_GetBptQty    = GetBptQty
    IDAAPI_GetBptEA     = GetBptEA
    IDAAPI_GetBptAttr   = GetBptAttr
    IDAAPI_SegStart     = SegStart
    IDAAPI_SegEnd       = SegEnd
    IDAAPI_GetBytes     = get_many_bytes
    IDAAPI_AskYN        = AskYN
    IDAAPI_AskFile      = AskFile
    IDAAPI_AskLong      = AskLong
    IDAAPI_NextHead     = NextHead
    IDAAPI_GetDisasm    = GetDisasmEx
    IDAAPI_NextThat     = nextthat
    IDAAPI_Jump         = Jump
    # classes
    IDAAPI_Choose       = Choose2

# PyQt
from PyQt5 import *
from PyQt5.QtWidgets import *

# Unicorn SDK
from unicorn import *
from unicorn.arm_const import *
from unicorn.arm64_const import *
from unicorn.mips_const import *
from unicorn.x86_const import *

# === Configuration

class UEMU_CONFIG:
    
    IDAViewColor_PC     = 0x00B3CBFF
    IDAViewColor_Reset  = 0xFFFFFFFF

    UnicornPageSize     = 0x1000

# === Helpers

class UEMU_HELPERS:

    # Menu

    MenuItem = collections.namedtuple("MenuItem", ["action", "handler", "title", "tooltip", "shortcut", "popup"])

    class IdaMenuActionHandler(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.form_type == BWN_DISASM:
                self.action_handler.handle_menu_action(self.action_type)
            return 1

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

    # IDA

    @staticmethod
    def exec_on_main(callable, dbflag):
        res = execute_sync(callable, dbflag)
        if res == -1:
            uemu_log("! <C> execute_sync(%s) failed" % (str(callable)))
            return False
        return True if res == 1 else False

    # Others

    @staticmethod
    def ALIGN_PAGE_DOWN(x):
        return x & ~(UEMU_CONFIG.UnicornPageSize - 1)

    @staticmethod
    def ALIGN_PAGE_UP(x):
        return (x + UEMU_CONFIG.UnicornPageSize - 1) & ~(UEMU_CONFIG.UnicornPageSize-1)

    @staticmethod
    def inf_is_be():
        if IDA_SDK_VERSION >= 900:
            return inf_is_be()
        elif IDA_SDK_VERSION >= 700:
            return cvar.inf.is_be()
        else:
            return cvar.inf.mf

    @staticmethod
    def get_arch():
        if ph.id == PLFM_386 and ph.flag & PR_USE64:
            return "x64"
        elif ph.id == PLFM_386 and ph.flag & PR_USE32:
            return "x86"
        elif ph.id == PLFM_ARM and ph.flag & PR_USE64:
            if UEMU_HELPERS.inf_is_be():
                return "arm64be"
            else:
                return "arm64le"
        elif ph.id == PLFM_ARM and ph.flag & PR_USE32:
            if UEMU_HELPERS.inf_is_be():
                return "armbe"
            else:
                return "armle"
        elif ph.id == PLFM_MIPS and ph.flag & PR_USE64:
            if UEMU_HELPERS.inf_is_be():
                return "mips64be"
            else:
                return "mips64le"
        elif ph.id == PLFM_MIPS and ph.flag & PR_USE32:
            if UEMU_HELPERS.inf_is_be():
                return "mipsbe"
            else:
                return "mipsle"
        else:
            return ""

    @staticmethod
    def get_pc_register(arch):
        if arch.startswith("arm64"):
            arch = "arm64"
        elif arch.startswith("arm"):
            arch = "arm"
        elif arch.startswith("mips"):
            arch = "mips"

        registers = {
            "x64" : ("rip", UC_X86_REG_RIP),
            "x86" : ("eip", UC_X86_REG_EIP),
            "arm" : ("PC", UC_ARM_REG_PC),
            "arm64" : ("PC", UC_ARM64_REG_PC),
            "mips" : ("pc", UC_MIPS_REG_PC),
        }
        return registers[arch]

    @staticmethod
    def get_stack_register(arch):
        if arch.startswith("arm64"):
            arch = "arm64"
        elif arch.startswith("arm"):
            arch = "arm"
        elif arch.startswith("mips"):
            arch = "mips"

        registers = {
            "x64" : ("rsp", UC_X86_REG_RSP),
            "x86" : ("esp", UC_X86_REG_ESP),
            "arm" : ("SP", UC_ARM_REG_SP),
            "arm64" : ("SP", UC_ARM64_REG_SP),
            "mips" : ("sp", UC_MIPS_REG_29),
        }
        return registers[arch]

    @staticmethod
    def get_register_map(arch):
        if arch.startswith("arm64"):
            arch = "arm64"
        elif arch.startswith("arm"):
            arch = "arm"
        elif arch.startswith("mips"):
            arch = "mips"

        registers = {
            "x64" : [
                [ "rax",    UC_X86_REG_RAX  ],
                [ "rbx",    UC_X86_REG_RBX  ],
                [ "rcx",    UC_X86_REG_RCX  ],
                [ "rdx",    UC_X86_REG_RDX  ],
                [ "rsi",    UC_X86_REG_RSI  ],
                [ "rdi",    UC_X86_REG_RDI  ],
                [ "rbp",    UC_X86_REG_RBP  ],
                [ "rsp",    UC_X86_REG_RSP  ],
                [ "r8",     UC_X86_REG_R8   ],
                [ "r9",     UC_X86_REG_R9   ],
                [ "r10",    UC_X86_REG_R10  ],
                [ "r11",    UC_X86_REG_R11  ],
                [ "r12",    UC_X86_REG_R12  ],
                [ "r13",    UC_X86_REG_R13  ],
                [ "r14",    UC_X86_REG_R14  ],
                [ "r15",    UC_X86_REG_R15  ],
                [ "rip",    UC_X86_REG_RIP  ],
                [ "sp",     UC_X86_REG_SP   ],
            ],
            "x86" : [
                [ "eax",    UC_X86_REG_EAX  ],
                [ "ebx",    UC_X86_REG_EBX  ],
                [ "ecx",    UC_X86_REG_ECX  ],
                [ "edx",    UC_X86_REG_EDX  ],
                [ "esi",    UC_X86_REG_ESI  ],
                [ "edi",    UC_X86_REG_EDI  ],
                [ "ebp",    UC_X86_REG_EBP  ],
                [ "esp",    UC_X86_REG_ESP  ],
                [ "eip",    UC_X86_REG_EIP  ],
                [ "sp",     UC_X86_REG_SP   ],
            ],
            "arm" : [
                [ "R0",     UC_ARM_REG_R0  ],
                [ "R1",     UC_ARM_REG_R1  ],
                [ "R2",     UC_ARM_REG_R2  ],
                [ "R3",     UC_ARM_REG_R3  ],
                [ "R4",     UC_ARM_REG_R4  ],
                [ "R5",     UC_ARM_REG_R5  ],
                [ "R6",     UC_ARM_REG_R6  ],
                [ "R7",     UC_ARM_REG_R7  ],
                [ "R8",     UC_ARM_REG_R8  ],
                [ "R9",     UC_ARM_REG_R9  ],
                [ "R10",    UC_ARM_REG_R10 ],
                [ "R11",    UC_ARM_REG_R11 ],
                [ "R12",    UC_ARM_REG_R12 ],
                [ "PC",     UC_ARM_REG_PC  ],
                [ "SP",     UC_ARM_REG_SP  ],
                [ "LR",     UC_ARM_REG_LR  ],
                [ "CPSR",   UC_ARM_REG_CPSR ]
            ],
            "arm64" : [
                [ "X0",     UC_ARM64_REG_X0  ],
                [ "X1",     UC_ARM64_REG_X1  ],
                [ "X2",     UC_ARM64_REG_X2  ],
                [ "X3",     UC_ARM64_REG_X3  ],
                [ "X4",     UC_ARM64_REG_X4  ],
                [ "X5",     UC_ARM64_REG_X5  ],
                [ "X6",     UC_ARM64_REG_X6  ],
                [ "X7",     UC_ARM64_REG_X7  ],
                [ "X8",     UC_ARM64_REG_X8  ],
                [ "X9",     UC_ARM64_REG_X9  ],
                [ "X10",    UC_ARM64_REG_X10 ],
                [ "X11",    UC_ARM64_REG_X11 ],
                [ "X12",    UC_ARM64_REG_X12 ],
                [ "X13",    UC_ARM64_REG_X13 ],
                [ "X14",    UC_ARM64_REG_X14 ],
                [ "X15",    UC_ARM64_REG_X15 ],
                [ "X16",    UC_ARM64_REG_X16 ],
                [ "X17",    UC_ARM64_REG_X17 ],
                [ "X18",    UC_ARM64_REG_X18 ],
                [ "X19",    UC_ARM64_REG_X19 ],
                [ "X20",    UC_ARM64_REG_X20 ],
                [ "X21",    UC_ARM64_REG_X21 ],
                [ "X22",    UC_ARM64_REG_X22 ],
                [ "X23",    UC_ARM64_REG_X23 ],
                [ "X24",    UC_ARM64_REG_X24 ],
                [ "X25",    UC_ARM64_REG_X25 ],
                [ "X26",    UC_ARM64_REG_X26 ],
                [ "X27",    UC_ARM64_REG_X27 ],
                [ "X28",    UC_ARM64_REG_X28 ],
                [ "PC",     UC_ARM64_REG_PC  ],
                [ "SP",     UC_ARM64_REG_SP  ],
                [ "FP",     UC_ARM64_REG_FP  ],
                [ "LR",     UC_ARM64_REG_LR  ],
                [ "NZCV",   UC_ARM64_REG_NZCV ]
            ],
            "mips" : [
                [ "zero",   UC_MIPS_REG_0   ],
                [ "at",     UC_MIPS_REG_1   ],
                [ "v0",     UC_MIPS_REG_2   ],
                [ "v1",     UC_MIPS_REG_3   ],
                [ "a0",     UC_MIPS_REG_4   ],
                [ "a1",     UC_MIPS_REG_5   ],
                [ "a2",     UC_MIPS_REG_6   ],
                [ "a3",     UC_MIPS_REG_7   ],
                [ "t0",     UC_MIPS_REG_8   ],
                [ "t1",     UC_MIPS_REG_9   ],
                [ "t2",     UC_MIPS_REG_10  ],
                [ "t3",     UC_MIPS_REG_11  ],
                [ "t4",     UC_MIPS_REG_12  ],
                [ "t5",     UC_MIPS_REG_13  ],
                [ "t6",     UC_MIPS_REG_14  ],
                [ "t7",     UC_MIPS_REG_15  ],
                [ "s0",     UC_MIPS_REG_16  ],
                [ "s1",     UC_MIPS_REG_17  ],
                [ "s2",     UC_MIPS_REG_18  ],
                [ "s3",     UC_MIPS_REG_19  ],
                [ "s4",     UC_MIPS_REG_20  ],
                [ "s5",     UC_MIPS_REG_21  ],
                [ "s6",     UC_MIPS_REG_22  ],
                [ "s7",     UC_MIPS_REG_23  ],
                [ "t8",     UC_MIPS_REG_24  ],
                [ "t9",     UC_MIPS_REG_25  ],
                [ "k0",     UC_MIPS_REG_26  ],
                [ "k1",     UC_MIPS_REG_27  ],
                [ "gp",     UC_MIPS_REG_28  ],
                [ "sp",     UC_MIPS_REG_29  ],
                [ "fp",     UC_MIPS_REG_30  ],
                [ "ra",     UC_MIPS_REG_31  ],
                [ "pc",     UC_MIPS_REG_PC  ],
            ]
        }
        return registers[arch]

    @staticmethod
    def get_register_bits(arch):
        if arch.startswith("arm64"):
            arch = "arm64"
        elif arch.startswith("arm"):
            arch = "arm"
        elif arch.startswith("mips"):
            arch = "mips"

        registers_bits = {
            "x64"   : 64,
            "x86"   : 32,
            "arm"   : 32,
            "arm64" : 64,
            "mips"  : 32
        }
        return registers_bits[arch]

    @staticmethod
    def get_register_ext_map(arch):
        if arch.startswith("arm64"):
            arch = "arm64"
        elif arch.startswith("arm"):
            arch = "arm"
        elif arch.startswith("mips"):
            arch = "mips"

        registers_ext = {
            "x64" : [
            ],
            "x86" : [
            ],
            "arm" : [
                [ "D0",     UC_ARM_REG_D0  ],
                [ "D1",     UC_ARM_REG_D1  ],
                [ "D2",     UC_ARM_REG_D2  ],
                [ "D3",     UC_ARM_REG_D3  ],
                [ "D4",     UC_ARM_REG_D4  ],
                [ "D5",     UC_ARM_REG_D5  ],
                [ "D6",     UC_ARM_REG_D6  ],
                [ "D7",     UC_ARM_REG_D7  ],
                [ "D8",     UC_ARM_REG_D8  ],
                [ "D9",     UC_ARM_REG_D9  ],
                [ "D10",    UC_ARM_REG_D10 ],
                [ "D11",    UC_ARM_REG_D11 ],
                [ "D12",    UC_ARM_REG_D12 ],
                [ "D13",    UC_ARM_REG_D13 ],
                [ "D14",    UC_ARM_REG_D14 ],
                [ "D15",    UC_ARM_REG_D15 ],
                [ "D16",    UC_ARM_REG_D16 ],
                [ "D17",    UC_ARM_REG_D17 ],
                [ "D18",    UC_ARM_REG_D18 ],
                [ "D19",    UC_ARM_REG_D19 ],
                [ "D20",    UC_ARM_REG_D20 ],
                [ "D21",    UC_ARM_REG_D21 ],
                [ "D22",    UC_ARM_REG_D22 ],
                [ "D23",    UC_ARM_REG_D23 ],
                [ "D24",    UC_ARM_REG_D24 ],
                [ "D25",    UC_ARM_REG_D25 ],
                [ "D26",    UC_ARM_REG_D26 ],
                [ "D27",    UC_ARM_REG_D27 ],
                [ "D28",    UC_ARM_REG_D28 ],
                [ "D29",    UC_ARM_REG_D29 ],
                [ "D30",    UC_ARM_REG_D30 ],
                [ "D31",    UC_ARM_REG_D31 ],
            ],
            "arm64" : [
                [ "Q0",     UC_ARM64_REG_Q0  ],
                [ "Q1",     UC_ARM64_REG_Q1  ],
                [ "Q2",     UC_ARM64_REG_Q2  ],
                [ "Q3",     UC_ARM64_REG_Q3  ],
                [ "Q4",     UC_ARM64_REG_Q4  ],
                [ "Q5",     UC_ARM64_REG_Q5  ],
                [ "Q6",     UC_ARM64_REG_Q6  ],
                [ "Q7",     UC_ARM64_REG_Q7  ],
                [ "Q8",     UC_ARM64_REG_Q8  ],
                [ "Q9",     UC_ARM64_REG_Q9  ],
                [ "Q10",    UC_ARM64_REG_Q10 ],
                [ "Q11",    UC_ARM64_REG_Q11 ],
                [ "Q12",    UC_ARM64_REG_Q12 ],
                [ "Q13",    UC_ARM64_REG_Q13 ],
                [ "Q14",    UC_ARM64_REG_Q14 ],
                [ "Q15",    UC_ARM64_REG_Q15 ],
                [ "Q16",    UC_ARM64_REG_Q16 ],
                [ "Q17",    UC_ARM64_REG_Q17 ],
                [ "Q18",    UC_ARM64_REG_Q18 ],
                [ "Q19",    UC_ARM64_REG_Q19 ],
                [ "Q20",    UC_ARM64_REG_Q20 ],
                [ "Q21",    UC_ARM64_REG_Q21 ],
                [ "Q22",    UC_ARM64_REG_Q22 ],
                [ "Q23",    UC_ARM64_REG_Q23 ],
                [ "Q24",    UC_ARM64_REG_Q24 ],
                [ "Q25",    UC_ARM64_REG_Q25 ],
                [ "Q26",    UC_ARM64_REG_Q26 ],
                [ "Q27",    UC_ARM64_REG_Q27 ],
                [ "Q28",    UC_ARM64_REG_Q28 ],
                [ "Q29",    UC_ARM64_REG_Q29 ],
                [ "Q30",    UC_ARM64_REG_Q30 ],
                [ "Q31",    UC_ARM64_REG_Q31 ],
            ],
            "mips" : [
            ]
        }
        return registers_ext[arch]

    @staticmethod
    def get_register_ext_bits(arch):
        if arch.startswith("arm64"):
            arch = "arm64"
        elif arch.startswith("arm"):
            arch = "arm"
        elif arch.startswith("mips"):
            arch = "mips"

        registers_ext_bits = {
            "x64"   : 0,
            "x86"   : 0,
            "arm"   : 64,
            "arm64" : 128,
            "mips"  : 0
        }
        return registers_ext_bits[arch]

    @staticmethod
    def get_register_ext_format(arch):
        if arch.startswith("arm64"):
            arch = "arm64"
        elif arch.startswith("arm"):
            arch = "arm"
        elif arch.startswith("mips"):
            arch = "mips"

        registers_ext_format = {
            "x64"   : "",
            "x86"   : "",
            "arm"   : "0x%.16X",
            "arm64" : "0x%.32X",
            "mips"  : 0
        }
        return registers_ext_format[arch]

    @staticmethod
    def is_thumb_ea(ea):
        def handler():
            if ph.id == PLFM_ARM and not ph.flag & PR_USE64:
                if IDA_SDK_VERSION >= 700:
                    t = get_sreg(ea, "T")  # get T flag
                else:
                    t = get_segreg(ea, 20) # get T flag
                return t is not BADSEL and t is not 0
            else:
                return 0
        return UEMU_HELPERS.exec_on_main(handler, MFF_READ)

    @staticmethod
    def trim_spaces(string):
        return ' '.join(str(string).split())

    @staticmethod
    def bytes_to_str(bytes):
        return " ".join("{:02X}".format(ord(c) if type(c) is str else c) for c in bytes)

    @staticmethod
    def is_alt_pressed():
        if QtWidgets.QApplication.keyboardModifiers() & QtCore.Qt.AltModifier:
            return True
        else:
            return False

    class InitedCallable(object):
        def __call__(self, flags):
            return IDAAPI_HasValue(flags)

    class UninitedCallable(object):
        def __call__(self, flags):
            return not IDAAPI_HasValue(flags)

# === Log

def uemu_log(entry, name="uEmu"):
    msg("[" + name + "]: " + entry + "\n")

# === uEmuInitView

class uEmuInitView(object):
    def __init__(self, owner):
        super(uEmuInitView, self).__init__()
        self.owner = owner

# === uEmuCpuContextView

class uEmuCpuContextView(simplecustviewer_t):

    def __init__(self, owner, extended):
        super(uEmuCpuContextView, self).__init__()
        self.hooks = None
        self.owner = owner
        self.extended = extended
        self.lastAddress = None
        self.lastContext = {}
        self.lastArch = None
        self.columns = None

    def Create(self, title):

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

        if IDA_SDK_VERSION >= 700:
            self.menu_cols1 = 1
            self.menu_cols2 = 2
            self.menu_cols3 = 3
            self.menu_update = 4

            class Hooks(UI_Hooks):

                class PopupActionHandler(action_handler_t):
                    def __init__(self, owner, menu_id):
                        action_handler_t.__init__(self)
                        self.owner = owner
                        self.menu_id = menu_id

                    def activate(self, ctx):
                        self.owner.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, "1 Column",       self.PopupActionHandler(self.form, self.form.menu_cols1),   None, None, -1))
                        attach_dynamic_action_to_popup(widget, popup, action_desc_t(None, "2 Columns",      self.PopupActionHandler(self.form, self.form.menu_cols2),   None, None, -1))
                        attach_dynamic_action_to_popup(widget, popup, action_desc_t(None, "3 Columns",      self.PopupActionHandler(self.form, self.form.menu_cols3),   None, None, -1))
                        attach_action_to_popup(widget, popup, "-", None)
                        attach_dynamic_action_to_popup(widget, popup, action_desc_t(None, "Change Context", self.PopupActionHandler(self.form, self.form.menu_update),  None, None, -1))
                        attach_action_to_popup(widget, popup, "-", None)

            if self.hooks is None:
                self.hooks = Hooks(self)
                self.hooks.hook()
        else:
            self.menu_cols1 = self.AddPopupMenu("1 Column")
            self.menu_cols2 = self.AddPopupMenu("2 Columns")
            self.menu_cols3 = self.AddPopupMenu("3 Columns")
            self.menu_sep = self.AddPopupMenu("")
            self.menu_update = self.AddPopupMenu("Change Context")
        return True

    def OnPopupMenu(self, menu_id):
        if menu_id == self.menu_cols1:
            self.columns = 1
        elif menu_id == self.menu_cols2:
            self.columns = 2
        elif menu_id == self.menu_cols3:
            self.columns = 3
        elif menu_id == self.menu_update:
            self.owner.change_cpu_context()
        else:
            # Unhandled
            return False
        if self.lastArch is not None and self.lastContext is not None and self.lastAddress is not None:
            self.SetContent(self.lastAddress, None)
        return True

    def SetContent(self, address, context):

        arch = UEMU_HELPERS.get_arch()
        if arch == "":
            return

        self.ClearLines()

        if self.extended:
            hdr_title = COLSTR("  CPU Extended context at [ ", SCOLOR_AUTOCMT)
        else:
            hdr_title = COLSTR("  CPU context at [ ", SCOLOR_AUTOCMT)
        hdr_title += COLSTR("0x%X: " % address, SCOLOR_DREF)
        hdr_title += COLSTR(UEMU_HELPERS.trim_spaces(IDAAPI_GetDisasm(address, 0)), SCOLOR_INSN)
        hdr_title += COLSTR(" ]", SCOLOR_AUTOCMT)
        self.AddLine(hdr_title)
        self.AddLine("")

        if context is None and self.lastContext == 0:
            self.Refresh()
            return

        if self.columns is None:
            cols = self.owner.get_context_columns()
        else:
            cols = self.columns
        if self.extended:
            regList = UEMU_HELPERS.get_register_ext_map(arch)
        else:
            regList = UEMU_HELPERS.get_register_map(arch)
        reg_cnt = len(regList)
        lines = int(reg_cnt // cols) if reg_cnt % cols == 0 else (reg_cnt // cols) + 1
        line = ""
        for i in range(lines):
            if i != 0:
                self.AddLine(line)
                line = ""

            for j in range(i, reg_cnt, lines):
                reg_label = regList[j][0]
                reg_key = regList[j][1]

                line = line + COLSTR(" %4s: " % str(reg_label), SCOLOR_REG)

                if context is not None:
                    currentValue = context.reg_read(reg_key)
                else:
                    currentValue = self.lastContext[reg_label]

                if self.extended:
                    value_format = UEMU_HELPERS.get_register_ext_format(arch)
                else:
                    if ph.flag & PR_USE64:
                        value_format = "0x%.16X"
                    else:
                        value_format = "0x%.8X"

                if reg_label in self.lastContext:
                    if self.lastContext[reg_label] != currentValue:
                        line += COLSTR(str(value_format % currentValue), SCOLOR_VOIDOP)
                    else:
                        line += COLSTR(str(value_format % currentValue), SCOLOR_NUMBER)
                else:
                    line += COLSTR(str(value_format % currentValue), SCOLOR_NUMBER)

                self.lastContext[reg_label] = currentValue

                line = line.ljust(35 * (int(j/lines) + 1))

        self.AddLine(line)

        self.Refresh()
        self.lastArch = arch
        self.lastAddress = address

    def OnClose(self):
        if self.hooks:
            self.hooks.unhook()
            self.hooks = None
        if self.extended:
            self.owner.ext_context_view_closed()
        else:
            self.owner.context_view_closed()

# === uEmuMemoryView

class uEmuMemoryRangeDialog(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 uEmuMemoryView(simplecustviewer_t):
    def __init__(self, owner, address, size):
        super(uEmuMemoryView, self).__init__()
        self.owner = owner
        self.viewid = address
        self.address = address
        self.size = size
        self.lastContent = []

    def Create(self, title):

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

        return True

    def SetContent(self, context):

        self.ClearLines()
        
        if context is None:
            return

        try:
            memory = context.mem_read(self.address, self.size)
        except UcError:
            return

        size = len(memory)

        hdr_title = COLSTR("  Memory at [ ", SCOLOR_AUTOCMT)
        hdr_title += COLSTR("0x%X: %d byte(s)" % (self.address, size), SCOLOR_DREF)
        hdr_title += COLSTR(" ]", SCOLOR_AUTOCMT)
        self.AddLine(str(hdr_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.address
        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.owner.memory_view_closed(self.viewid)

# === uEmuStackView

class uEmuStackView(simplecustviewer_t):

    stack_prelines = 12
    stack_postlines = 52

    def __init__(self, owner):
        super(uEmuStackView, self).__init__()
        self.owner = owner
        arch = UEMU_HELPERS.get_arch()
        _, self.uc_reg_sp = UEMU_HELPERS.get_stack_register(arch)

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

    def SetContent(self, context):

        self.ClearLines()
        if context is None:
            return

        sp = context.reg_read(self.uc_reg_sp)

        self.AddLine('')
        self.AddLine(COLSTR('  Stack at 0x%X' % sp, SCOLOR_AUTOCMT))
        self.AddLine('')

        arch = UEMU_HELPERS.get_arch()
        reg_bit_size = UEMU_HELPERS.get_register_bits(arch)
        reg_byte_size = reg_bit_size // 8
        value_format = '% .16X' if reg_bit_size == 64 else '% .8X'

        for i in range(-self.stack_prelines, self.stack_postlines):
            clr = SCOLOR_DREF if i < 0 else SCOLOR_INSN
            cur_addr = (sp + i * reg_byte_size)
            line = ('  ' + value_format + ': ') % cur_addr
            try:
                value = context.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.owner.stack_view_closed()

# === uEmuControlView

class uEmuControlView(PluginForm):
    def __init__(self, owner):
        self.owner = owner
        PluginForm.__init__(self)

    def OnCreate(self, form):
        self.parent = self.FormToPyQtWidget(form)
        self.PopulateForm()

    def PopulateForm(self):
        btnStart = QPushButton("Start")
        btnRun = QPushButton("Run")
        btnStep = QPushButton("Step")
        btnStop = QPushButton("Stop")

        btnStart.clicked.connect(self.OnEmuStart)
        btnRun.clicked.connect(self.OnEmuRun)
        btnStep.clicked.connect(self.OnEmuStep)
        btnStop.clicked.connect(self.OnEmuStop)

        hbox = QHBoxLayout()
        hbox.setAlignment(QtCore.Qt.AlignCenter)
        hbox.addWidget(btnStart)
        hbox.addWidget(btnRun)
        hbox.addWidget(btnStep)
        hbox.addWidget(btnStop)

        self.parent.setLayout(hbox)

    def OnEmuStart(self, code=0):
        self.owner.emu_start()

    def OnEmuRun(self, code=0):
        self.owner.emu_run()

    def OnEmuStep(self, code=0):
        self.owner.emu_step()

    def OnEmuStop(self, code=0):
        self.owner.emu_stop()

    def OnClose(self, form):
        self.owner.contol_view_closed()

# === uEmuMappeduMemoryView

class uEmuMappeduMemoryView(IDAAPI_Choose):

    def __init__(self, owner, memory, flags=0, width=None, height=None, embedded=False):
        IDAAPI_Choose.__init__(
            self,
            "uEmu Mapped Memory",
            [ ["Start", 20], ["End", 20], ["Permissions", 10] ],
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = memory
        self.icon = -1
        self.selcount = 0
        self.popup_names = [ "", "Dump To File", "Show", "" ]
        self.owner = owner

    def OnClose(self):
        pass

    def OnDeleteLine(self, n): # Save JSON
        filePath = IDAAPI_AskFile(1, "*.bin", "Dump memory")
        if filePath is not None:
            with open(filePath, 'wb') as outfile:
                address = self.items[n][0]
                size = self.items[n][1] - self.items[n][0] + 1
                outfile.seek(0, 0)
                outfile.write(self.owner.unicornEngine.get_mapped_bytes(address, size))
                outfile.close()
        return n

    def OnEditLine(self, n):
        address = self.items[n][0]
        size = self.items[n][1] - self.items[n][0] + 1
        self.owner.show_memory(address, size)

    def OnGetLine(self, n):
        return [
            "0x%X" % self.items[n][0],
            "0x%X" % self.items[n][1],
            "%c%c%c" % (
                "r" if self.items[n][2] & UC_PROT_READ else "-",
                "w" if self.items[n][2] & UC_PROT_WRITE else "-",
                "x" if self.items[n][2] & UC_PROT_EXEC else "-")
        ]

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

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

# === Settings

class uEmuSettingsDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:chk_followpc}
BUTTON YES* Save
BUTTON CANCEL Cancel
uEmu Settings
<Follow PC:{chk_followpc}>
<Convert to Code automatically:{chk_forcecode}>
<Trace instructions:{chk_trace}>
<Lazy mapping:{chk_lazymapping}>{emu_group}>
""", {
        'emu_group': Form.ChkGroupControl(("chk_followpc", "chk_forcecode", "chk_trace", "chk_lazymapping")),
        })

# === uEmuUnicornEngine

class uEmuRegisterValueDialog(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("Enter hex value for [ " + regName + " ]"),
        'reg_val': Form.NumericInput(tp=Form.FT_HEX, swidth=20)
        })

class uEmuRegisterValueLHDialog(Form):
    def __init__(self, regName):
        Form.__init__(self, r"""STARTITEM {id:reg_valh}
BUTTON YES* Save
BUTTON CANCEL Cancel
Register Value
{reg_label}
<##High\::{reg_valh}>
<##Low\: :{reg_vall}>
""", {
        'reg_label': Form.StringLabel("Enter hex value for [ " + regName + " ]"),
        'reg_valh': Form.NumericInput(tp=Form.FT_HEX, swidth=20),
        'reg_vall': Form.NumericInput(tp=Form.FT_HEX, swidth=20)
        })

class uEmuMapBinaryFileDialog(Form):
    def __init__(self, address):
        Form.__init__(self, r"""STARTITEM {id:file_name}
BUTTON YES* Map
BUTTON CANCEL Cancel
Map Binary File
{form_change_cb}
<#Select file to open#File\::{file_name}>
{note_label}
<##Address\:    :{mem_addr}>
<##Data Offset\::{mem_offset}>
<##Data Size\:  :{mem_size}> {total_label}

""", {
        'file_name': Form.FileInput(swidth=50, open=True),
        'note_label': Form.StringLabel("NOTE: Region MUST cover 0x%X" % (address)),
        'mem_addr': Form.NumericInput(swidth=20, tp=Form.FT_HEX),
        'mem_offset': Form.NumericInput(swidth=20, tp=Form.FT_HEX),
        'mem_size': Form.NumericInput(swidth=20, tp=Form.FT_HEX),
        'total_label': Form.StringLabel(""),
        'form_change_cb': Form.FormChangeCb(self.OnFormChange)
        })

    def OnFormChange(self, fid):
        if fid == self.file_name.id:
            filepath = self.GetControlValue(self.file_name)
            if os.path.isfile(filepath):
                size = os.path.getsize(filepath)
                self.SetControlValue(self.mem_size, size)
                self.SetControlValue(self.total_label, "Total: %d (0x%X)" % (size, size))
        return 1


class uEmuContextInitDialog(IDAAPI_Choose):

    def __init__(self, regs, flags=0, width=None, height=None, embedded=False):
        IDAAPI_Choose.__init__(
            self,
            "uEmu CPU Context Edit",
            [ ["Register", 10], ["Value", 30] ],
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = regs
        self.icon = -1
        self.selcount = 0
        self.popup_names = [ "Load CPU Context", "Save CPU Context", "Edit Register Value", "" ]

    def OnClose(self):
        pass

    def OnInsertLine(self, n=0): # load JSON
        filePath = IDAAPI_AskFile(0, "*.json", "Load CPU Context")
        if filePath is not None:
            with open(filePath, 'r') as infile:
                self.items = json.load(infile)
                infile.close()
                self.Refresh()

    def OnDeleteLine(self, n): # Save JSON
        filePath = IDAAPI_AskFile(1, "*.json", "Save CPU Context")
        if filePath is not None:
            with open(filePath, 'w') as outfile:
                json.dump(self.items, outfile)
                outfile.close()
        return n

    def OnEditLine(self, n):
        bits = self.items[n][2]
        ext = self.items[n][3]
        if ext:
            regDlg = uEmuRegisterValueLHDialog(self.items[n][0])
            regDlg.Compile()
            regDlg.reg_valh.value = int(self.items[n][1], 0) >> (bits >> 1)
            regDlg.reg_vall.value = int(self.items[n][1], 0)
            ok = regDlg.Execute()
            if ok == 1:
                newValue = (regDlg.reg_valh.value << (bits >> 1)) | regDlg.reg_vall.value
                self.items[n][1] = "0x%X" % newValue
        else:
            regDlg = uEmuRegisterValueDialog(self.items[n][0])
            regDlg.Compile()
            regDlg.reg_val.value = int(self.items[n][1], 0)
            ok = regDlg.Execute()
            if ok == 1:
                newValue = regDlg.reg_val.value
                self.items[n][1] = "0x%X" % newValue
        self.Refresh()

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

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

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

class uEmuUnicornEngine(object):

    mu = None
    pc = BADADDR

    emuActive       = False
    emuRunning      = False
    emuThread       = None

    kStepCount_Run  = 0
    emuStepCount    = 1

    fix_context     = None
    extended        = False

    def __init__(self, owner):
        super(uEmuUnicornEngine, self).__init__()
        self.owner = owner

        uc_setup = {
            "x64"       : [ UC_X86_REG_RIP,     UC_ARCH_X86,    UC_MODE_64                              ],
            "x86"       : [ UC_X86_REG_EIP,     UC_ARCH_X86,    UC_MODE_32                              ],
            "arm64be"   : [ UC_ARM64_REG_PC,    UC_ARCH_ARM64,  UC_MODE_ARM     | UC_MODE_BIG_ENDIAN    ],
            "arm64le"   : [ UC_ARM64_REG_PC,    UC_ARCH_ARM64,  UC_MODE_ARM     | UC_MODE_LITTLE_ENDIAN ],
            "armbe"     : [ UC_ARM_REG_PC,      UC_ARCH_ARM,    UC_MODE_ARM     | UC_MODE_BIG_ENDIAN    ],
            "armle"     : [ UC_ARM_REG_PC,      UC_ARCH_ARM,    UC_MODE_ARM     | UC_MODE_LITTLE_ENDIAN ],
            "mips64be"  : [ UC_MIPS_REG_PC,     UC_ARCH_MIPS,   UC_MODE_MIPS64  | UC_MODE_BIG_ENDIAN    ],
            "mips64le"  : [ UC_MIPS_REG_PC,     UC_ARCH_MIPS,   UC_MODE_MIPS64  | UC_MODE_LITTLE_ENDIAN ],
            "mipsbe"    : [ UC_MIPS_REG_PC,     UC_ARCH_MIPS,   UC_MODE_MIPS32  | UC_MODE_BIG_ENDIAN    ],
            "mipsle"    : [ UC_MIPS_REG_PC,     UC_ARCH_MIPS,   UC_MODE_MIPS32  | UC_MODE_LITTLE_ENDIAN ],
        }
        arch = UEMU_HELPERS.get_arch()
        if arch == "":
            uemu_log("CPU is not supported")    
            return

        self.uc_reg_pc, self.uc_arch, self.uc_mode = uc_setup[arch]
        uemu_log("Unicorn version [ %s ]" % (unicorn.__version__))
        uemu_log("CPU arch set to [ %s ]" % (arch))

    def is_active(self):
        return self.emuActive

    def is_running(self):
        return self.emuRunning

    def get_context(self):
        # Looks like unicorn context is not serializable in python
        # return self.mu.context_save()

        uc_context = {}
        reg_map = UEMU_HELPERS.get_register_map(UEMU_HELPERS.get_arch())
        reg_ext_map = UEMU_HELPERS.get_register_ext_map(UEMU_HELPERS.get_arch())
        uc_context["cpu"] = [ [ row[1], self.mu.reg_read(row[1]) ] for row in reg_map ]
        uc_context["cpu_ext"] = [ [ row[1], self.mu.reg_read(row[1]) ] for row in reg_ext_map ]
        uc_context["mem"] = [ [ memStart, memEnd, memPerm, self.mu.mem_read(memStart, memEnd - memStart + 1) ] for (memStart, memEnd, memPerm) in self.mu.mem_regions() ]
       
        return uc_context

    def set_context(self, context):
        # Looks like unicorn context is not serializable in python
        # self.mu.context_restore(context)

        for reg in context["cpu"]:
            self.mu.reg_write(reg[0], reg[1])

        for reg in context["cpu_ext"]:
            self.mu.reg_write(reg[0], reg[1])

        for mem in context["mem"]:
            try:
                memStart = mem[0]
                memEnd = mem[1]
                memPerm = mem[2]
                memData = mem[3]
                uemu_log("  map [%X:%X]" % (memStart, memEnd))
                self.mu.mem_map(memStart, memEnd - memStart + 1, memPerm)
                self.mu.mem_write(memStart, str(memData))
            except UcError as e:
                uemu_log("! <U> %s" % e)

        self.pc = self.mu.reg_read(self.uc_reg_pc)
        IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_PC)

        if not self.emuActive:
            if self.owner.trace_inst():
                self.trace_log()

            self.emuActive = True

    def is_memory_mapped(self, address):
        try:
            self.mu.mem_read(address, 1)
            return True
        except UcError as e:
            return False

    def jump_to_pc(self):
        IDAAPI_Jump(self.pc)

    def set_pc(self, pc):
        IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_Reset)
        self.pc = pc
        self.mu.reg_write(self.uc_reg_pc, self.pc)
        IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_PC)

    def map_memory(self, address, size):
        # - size is unsigned and must be != 0
        # - starting address must be aligned to 4KB
        # - map size must be multiple of the page size (4KB)
        # - permissions can only contain READ or WRITE perms
        try:
            memStart = address
            memEnd = address + size
            memStartAligned = UEMU_HELPERS.ALIGN_PAGE_DOWN(memStart)
            memEndAligned = UEMU_HELPERS.ALIGN_PAGE_UP(memEnd)
            uemu_log("  map [%X:%X] -> [%X:%X]" % (memStart, memEnd - 1, memStartAligned, memEndAligned - 1))
            self.mu.mem_map(memStartAligned, memEndAligned - memStartAligned)
        except UcError as e:
            uemu_log("! <U> %s" % e)

    def map_empty(self, address, size):
        self.map_memory(address, size)
        self.mu.mem_write(UEMU_HELPERS.ALIGN_PAGE_DOWN(address), b"\x00" * UEMU_CONFIG.UnicornPageSize)

    def map_binary(self, address, size):
        binMapDlg = uEmuMapBinaryFileDialog(address)
        binMapDlg.Compile()
        
        ok = binMapDlg.Execute()
        if ok != 1:
            return False

        bin_name = binMapDlg.file_name.value
        with open(bin_name, 'rb') as file:
            file.seek(0, 2)
            file_size = file.tell()

            mem_addr = binMapDlg.mem_addr.value
            mem_offset = binMapDlg.mem_offset.value
            mem_size = binMapDlg.mem_size.value

            # check if file is big enough
            if mem_offset + mem_size > file_size:
                file.close()
                return False

            # validate range
            if address < UEMU_HELPERS.ALIGN_PAGE_DOWN(mem_addr):
                return False
            if (address + size) > (mem_addr + mem_size):
                return False

            # read data from file
            file.seek(mem_offset, 0)
            data = file.read(mem_size)
            file.close()

            try:
                self.map_memory(mem_addr, mem_size)
                self.mu.mem_write(mem_addr, data)
            except UcError as e:
                return False

            return True

        return False

    def copy_inited_data(self, start, end):
        if IDAAPI_IsLoaded(start):
            find_inited = False
        else:
            find_inited = True
        ptr = IDAAPI_NextThat(start, end, UEMU_HELPERS.InitedCallable() if find_inited else UEMU_HELPERS.UninitedCallable())
        while ptr < end:
            if not find_inited:
                # found uninitialized
                tmp_end = end if ptr > end else ptr
                uemu_log("  cpy [%X:%X]" % (start, tmp_end - 1))
                self.mu.mem_write(start, IDAAPI_GetBytes(start, tmp_end - start))
            else:
                # found initialized
                uemu_log("  skp [%X:%X]" % (start, ptr - 1))
                start = ptr

            find_inited = not find_inited
            ptr = IDAAPI_NextThat(ptr, end, UEMU_HELPERS.InitedCallable() if find_inited else UEMU_HELPERS.UninitedCallable())
        # process tail if any
        if not find_inited:
            tmp_end = end if ptr > end else ptr
            uemu_log("  cpy [%X:%X]" % (start, tmp_end - 1))
            self.mu.mem_write(start, IDAAPI_GetBytes(start, tmp_end - start))

    def fetch_segments(self):
        uemu_log("Fetching segments...")
        for segEA in Segments():
            segStart = IDAAPI_SegStart(segEA)
            segEnd = IDAAPI_SegEnd(segEA)
            found = False
            overlap_start = False
            overlap_end = False
            tmp = 0
            for (memStart, memEnd, memPerm) in self.mu.mem_regions():
                memEnd = memEnd + 1
                if segStart >= memStart and segEnd <= memEnd:
                    found = True
                    break
                if segStart >= memStart and segStart < memEnd:
                    overlap_start = True
                    tmp = memEnd
                    break
                if segEnd >= memStart and segEnd < memEnd:
                    overlap_end = True
                    tmp = memStart
                    break

            if found:               # already mapped
                continue
            elif overlap_start:     # partial overlap (start)
                self.map_memory(tmp, segEnd - tmp)
                self.copy_inited_data(tmp, segEnd)
            elif overlap_end:       # patrial overlap (end)
                self.map_memory(segStart, tmp - segStart)
                self.copy_inited_data(segStart, tmp)
            else:                   # not found
                self.map_memory(segStart, segEnd - segStart)
                self.copy_inited_data(segStart, segEnd)

    def get_mapped_memory(self):
        return [ [ memStart, memEnd, memPerm ] for (memStart, memEnd, memPerm) in self.mu.mem_regions() ]

    def get_mapped_bytes(self, address, size):
        return self.mu.mem_read(address, size)

    def trace_log(self):
        if self.owner.ext_hooks.trace_log is not None:
            if self.owner.ext_hooks.trace_log(self.pc):
                return

        bytes = IDAAPI_GetBytes(self.pc, IDAAPI_NextHead(self.pc) - self.pc)
        bytes_hex = UEMU_HELPERS.bytes_to_str(bytes)
        uemu_log("* TRACE<I> 0x%X | %-16s | %s" % (self.pc, bytes_hex, IDAAPI_GetDisasm(self.pc, 0)))

    def hook_mem_access(self, uc, access, address, size, value, user_data):
        def bpt_sync_check():
            return 1 if self.is_breakpoint_reached(address) else 0
        
        if UEMU_HELPERS.exec_on_main(bpt_sync_check, MFF_READ):
            # vvv Workaround to fix issue when register are still updated even if emu_stop is called
            self.fix_context = self.mu.context_save()
            # ^^^
            uc.emu_stop()
            def result_handler():
                uemu_log("Memory breakpoint [0x%X] reached from 0x%X : %s" % (address, self.pc, UEMU_HELPERS.trim_spaces(IDAAPI_GetDisasm(self.pc, 0))))

                IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_PC)
                self.owner.update_context(self.pc, self.mu)

                if self.owner.follow_pc():
                    self.jump_to_pc()

                self.emuStepCount = 1
                return 1

            UEMU_HELPERS.exec_on_main(result_handler, MFF_WRITE)

    def hook_mem_invalid(self, uc, access, address, size, value, user_data):
        def result_handler():
            uemu_log("! <M> Missing memory at 0x%x, data size = %u, data value = 0x%x" %(address, size, value))

            if self.owner.lazy_mapping() and IDAAPI_IsLoaded(address):
                page_start = UEMU_HELPERS.ALIGN_PAGE_DOWN(address)
                page_end = UEMU_HELPERS.ALIGN_PAGE_UP(address + size)
                self.map_memory(page_start, page_end - page_start)
                self.copy_inited_data(page_start, page_end)
                return True

            while True:
                ok = IDAAPI_AskYN(1, "Memory [%X] is not mapped!\nDo you want to map it?\n   YES - Load Binary\n   NO - Fill page with zeroes\n   Cancel - Stop Emulation" % (address))
                if ok == 0:
                    self.map_empty(address, size)
                    break
                elif ok == 1:
                    if self.map_binary(address, size):
                        break
                else:
                    self.interrupt()
                    break
            return 1

        UEMU_HELPERS.exec_on_main(result_handler, MFF_WRITE)
        return 1

    def init_cpu_context(self, pc):
        self.extended = False
        
        # enable ARMv7 VFP
        if UEMU_HELPERS.get_arch() in ["armle", "armbe"]:
            if IDAAPI_AskYN(1, "Enable VFP instruction emulation?") == 1:
                tmp_val = self.mu.reg_read(UC_ARM_REG_C1_C0_2)
                tmp_val = tmp_val | (0xf << 20)
                self.mu.reg_write(UC_ARM_REG_C1_C0_2, tmp_val)
                enable_vfp = 0x40000000
                self.mu.reg_write(UC_ARM_REG_FPEXC, enable_vfp)
                self.extended = True
                uemu_log("VFP enabled")
        # enable ARMv8 FP/SIMD
        if UEMU_HELPERS.get_arch() in ["arm64le", "arm64be"]:
            if IDAAPI_AskYN(1, "Enable FP/SIMD instruction emulation?") == 1:
                cpacr = self.mu.reg_read(UC_ARM64_REG_CPACR_EL1)
                cpacr = cpacr | (0x3 << 20)
                self.mu.reg_write(UC_ARM64_REG_CPACR_EL1, cpacr)
                self.extended = True
                uemu_log("FP/SIMD enabled")

        reg_map = UEMU_HELPERS.get_register_map(UEMU_HELPERS.get_arch())
        regs = [ [ row[0], "0x0", UEMU_HELPERS.get_register_bits(UEMU_HELPERS.get_arch()), False ] for row in reg_map ]
        regs_len = len(regs)
        if self.extended:
            reg_ext_map = UEMU_HELPERS.get_register_ext_map(UEMU_HELPERS.get_arch())
            regs = regs + [ [ row[0], "0x0", UEMU_HELPERS.get_register_ext_bits(UEMU_HELPERS.get_arch()), True ] for row in reg_ext_map ]

        if self.owner.ext_hooks.init_context is not None:
            if self.owner.ext_hooks.init_context():
                return True

        cpuContext = uEmuContextInitDialog(regs)
        ok = cpuContext.show()
        if ok:
            for idx, val in enumerate(cpuContext.items[0:regs_len-1]):
                self.mu.reg_write(reg_map[idx][1], int(val[1], 0))
            self.mu.reg_write(self.uc_reg_pc, pc)
            if self.extended:
                for idx, val in enumerate(cpuContext.items[regs_len:]):
                    self.mu.reg_write(reg_ext_map[idx][1], int(val[1], 0))
            return True
        else:
            return False

    def set_cpu_context(self):
        reg_map = UEMU_HELPERS.get_register_map(UEMU_HELPERS.get_arch())
        regs = [ [ row[0], "0x%X" % self.mu.reg_read(row[1]), UEMU_HELPERS.get_register_bits(UEMU_HELPERS.get_arch()), False ] for row in reg_map ]
        regs_len = len(regs)
        if self.extended:
            reg_ext_map = UEMU_HELPERS.get_register_ext_map(UEMU_HELPERS.get_arch())
            regs = regs + [ [ row[0], "0x%X" % self.mu.reg_read(row[1]), UEMU_HELPERS.get_register_ext_bits(UEMU_HELPERS.get_arch()), True ] for row in reg_ext_map ]

        cpuContext = uEmuContextInitDialog(regs)
        ok = cpuContext.show()
        if ok:
            for idx, val in enumerate(cpuContext.items[0:regs_len-1]):
                if reg_map[idx][0] == UEMU_HELPERS.get_pc_register(UEMU_HELPERS.get_arch())[0]:
                    if self.pc != int(val[1], 0):
                        self.set_pc(int(val[1], 0))
                else:
                    self.mu.reg_write(reg_map[idx][1], int(val[1], 0))
            if self.extended:
                for idx, val in enumerate(cpuContext.items[regs_len:]):
                    self.mu.reg_write(reg_ext_map[idx][1], int(val[1], 0))
            return True
        else:
            return False

    def run_from(self, address):
        self.mu = Uc(self.uc_arch, self.uc_mode)
        self.mu.hook_add(UC_HOOK_MEM_READ_UNMAPPED | UC_HOOK_MEM_WRITE_UNMAPPED | UC_HOOK_MEM_FETCH_UNMAPPED, self.hook_mem_invalid)
        self.mu.hook_add(UC_HOOK_MEM_READ | UC_HOOK_MEM_WRITE, self.hook_mem_access)
        self.pc = address

        try:
            if self.init_cpu_context(address) == False:
                return

            if not self.owner.lazy_mapping():
                uemu_log("Mapping segments...")

                # get all segments and merge neighbours
                lastAddress = BADADDR
                for segEA in Segments():
                    segStart = IDAAPI_SegStart(segEA)
                    segEnd = IDAAPI_SegEnd(segEA)
                    endAligned = UEMU_HELPERS.ALIGN_PAGE_UP(segEnd)

                    # merge with previous if
                    # - we have mapped some segments already
                    # - aligned old segment is overlapping new segment
                    # otherwise map new

                    uemu_log("* seg [%X:%X]" % (segStart, segEnd))
                    if lastAddress != BADADDR and lastAddress > segStart:
                        if lastAddress < segEnd:
                            self.map_memory(lastAddress, endAligned - lastAddress)
                    else:
                        self.map_memory(segStart, endAligned - segStart)

                    # copy initialized bytes
                    self.copy_inited_data(segStart, segEnd)

                    lastAddress = endAligned

            if self.owner.trace_inst():
                self.trace_log()

            IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_PC)
            self.emuActive = True

        except UcError as e:
            uemu_log("! <U> %s" % e)

    def step_thread_main(self):
        try:
            emu_skip = False
            if self.owner.ext_hooks.emu_step is not None:
                emu_skip = self.owner.ext_hooks.emu_step(self.pc)

            if emu_skip == False:
                self.mu.emu_start(self.pc | 1 if UEMU_HELPERS.is_thumb_ea(self.pc) else self.pc, -1, count=1)

            # vvv Workaround to fix issue when registers are still updated even if emu_stop is called
            if self.fix_context is not None:
                self.mu.context_restore(self.fix_context)
                self.fix_context = None
            # ^^^
            def result_handler():
                self.pc = self.mu.reg_read(self.uc_reg_pc)
                if self.owner.trace_inst():
                    self.trace_log()

                if not IDAAPI_IsCode(IDAAPI_GetFlags(self.pc)) and self.owner.force_code():
                    uemu_log("Creating code at 0x%X" % (self.pc))
                    IDAAPI_DelItems(self.pc, DELIT_SIMPLE)
                    IDAAPI_MakeCode(self.pc)

                if self.emuStepCount != 1:
                    if not self.is_breakpoint_reached(self.pc):
                        if self.emuStepCount == self.kStepCount_Run:
                            self.step(self.kStepCount_Run)
                        else:
                            self.step(self.emuStepCount - 1)
                        return 1
                    else:
                        uemu_log("Breakpoint reached at 0x%X : %s" % (self.pc, UEMU_HELPERS.trim_spaces(IDAAPI_GetDisasm(self.pc, 0))))

                if not self.emuRunning:
                    uemu_log("Emulation interrupted at 0x%X : %s" % (self.pc, UEMU_HELPERS.trim_spaces(IDAAPI_GetDisasm(self.pc, 0))))

                # emulation stopped - update UI
                IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_PC)

                self.owner.update_context(self.pc, self.mu)

                if self.owner.follow_pc():
                    self.jump_to_pc()

                self.emuStepCount = 1
                self.emuRunning = False
                return 1

            UEMU_HELPERS.exec_on_main(result_handler, MFF_WRITE)

        except UcError as e:

            def result_handler():
                if self.emuRunning:
                    disasm = UEMU_HELPERS.trim_spaces(IDAAPI_GetDisasm(self.pc, 0))
                    if IDAAPI_AskYN(1, "Unhandled exception:\n   %s\nat:\n   0x%X: %s\n\nDo you want to skip this instruction?" % (e, self.pc, disasm)) != 1:
                        IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_PC)

                        if self.owner.follow_pc():
                            self.jump_to_pc()

                        self.emuRunning = False
                        return 1

                    self.pc = IDAAPI_NextHead(self.pc)
                    uemu_log("! <U> Unable to emulate [ %s ] - SKIP to 0x%X" % (disasm, self.pc))

                if self.emuStepCount != 1:
                    if not self.is_breakpoint_reached(self.pc):
                        if self.emuStepCount == self.kStepCount_Run:
                            self.step(self.kStepCount_Run)
                        else:
                            self.step(self.emuStepCount - 1)
                        return 1
                    else:
                        uemu_log("Breakpoint reached at 0x%X : %s" % (self.pc, UEMU_HELPERS.trim_spaces(IDAAPI_GetDisasm(self.pc, 0))))

                if not self.emuRunning:
                    uemu_log("Emulation interrupted at 0x%X : %s" % (self.pc, UEMU_HELPERS.trim_spaces(IDAAPI_GetDisasm(self.pc, 0))))

                # emulation stopped - update UI
                IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_PC)

                self.owner.update_context(self.pc, self.mu)

                if self.owner.follow_pc():
                    self.jump_to_pc()

                self.emuStepCount = 1
                self.emuRunning = False
                return 1

            UEMU_HELPERS.exec_on_main(result_handler, MFF_WRITE)

    def step(self, count=1):
        self.emuStepCount = count
        IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_Reset)

        self.emuThread = threading.Thread(target=self.step_thread_main)
        self.emuRunning = True
        self.emuThread.start()

    def run(self):
        self.step(self.kStepCount_Run)

    def interrupt(self):
        IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_Reset)
        self.mu.emu_stop()

        self.emuStepCount = 1
        self.emuRunning = False

    def is_breakpoint_reached(self, address):
        for idx in range(IDAAPI_GetBptQty()):
            if IDAAPI_GetBptEA(idx) == address and IDAAPI_GetBptAttr(address, BPTATTR_FLAGS) & BPT_ENABLED:
                return True
        return False

    def reset(self):
        self.interrupt()
        # unmap all regions
        for (start, end, perm) in self.mu.mem_regions():
            uemu_log("  unmap [%X:%X]" % (start, end))
            self.mu.mem_unmap(start, end - start + 1)

        IDAAPI_SetColor(self.pc, CIC_ITEM, UEMU_CONFIG.IDAViewColor_Reset)
        self.pc = BADADDR
        self.mu = None
        self.emuActive = False

# === uEmuPlugin

class uEmuExtensionHooks:
    init_context = None
    trace_log = None
    emu_step = None

class uEmuPlugin(plugin_t, UI_Hooks):

    popup_menu_hook = None

    flags = PLUGIN_HIDE
    comment = ""

    help = "Tiny cute emulator"
    wanted_name = "uEmu"
    plugin_name = "uEmu"
    wanted_hotkey = ""

    ext_hooks = None

    # --- PLUGIN DATA

    unicornEngine = None

    settings = {
        "follow_pc"     : False,
        "force_code"    : False,
        "trace_inst"    : False,
        "lazy_mapping"  : False,
    }

    controlView = None
    emuInitView = None
    cpuContextView = None
    cpuExtContextView = None
    stackView = None
    memoryViews = {}

    # --- PLUGIN LIFECYCLE

    def __init__(self, name = "uEmu", hooks = uEmuExtensionHooks()):
        super(uEmuPlugin, self).__init__()
        self.plugin_name = name
        self.wanted_name = name

        self.ext_hooks = hooks

        uemu_log("Init plugin %s" % (self.plugin_name))

    def init(self):
        self.hook_ui_actions()
        return PLUGIN_KEEP

    def run(self, arg):
        uemu_log("Run plugin " + self.plugin_name + " with (%d)" % arg)

    def run(self, arg = 0):
        uemu_log("Run plugin " + self.plugin_name)
        self.register_menu_actions()
        self.attach_main_menu_actions()
        self.unicornEngine = uEmuUnicornEngine(self)

    def term(self): # asynchronous unload (external, UI_Hook::term)
        self.unicornEngine = None
        self.unhook_ui_actions()
        self.detach_main_menu_actions()
        self.unregister_menu_actions()
        uemu_log("Unload plugin " + self.plugin_name)

    def get_context_columns(self):
        return 2

    def unload_plugin(self): # synchronous unload (internal, Main Menu)
        if self.unicornEngine.is_active():
            self.unicornEngine.reset()
        self.close_windows()
        self.memoryViews.clear()
        self.detach_main_menu_actions()
        self.unregister_menu_actions()

    def do_nothing(self):
        pass

    def ready_to_run(self):
        uemu_log("UI ready. Run plugin")
        self.register_menu_actions()
        self.attach_main_menu_actions()
        self.unicornEngine = uEmuUnicornEngine(self)

    def follow_pc(self):
        return self.settings["follow_pc"]

    def force_code(self):
        return self.settings["force_code"]

    def trace_inst(self):
        return self.settings["trace_inst"]

    def lazy_mapping(self):
        return self.settings["lazy_mapping"]

    def load_project(self):
        filePath = IDAAPI_AskFile(0, "*.emu", "Open uEmu project")
        if filePath is None:
            return
        with open(filePath, 'rb') as file:
            settings = pickle.load(file)
            self.unicornEngine.set_context(pickle.load(file))
            file.close()

            uemu_log("Project loaded from %s" % filePath)

    def save_project(self):
        filePath = IDAAPI_AskFile(1, "*.emu", "Save FridaLink project")
        if filePath is None:
            return
        with open(filePath, 'wb') as file:
            pickle.dump(self.settings, file, pickle.HIGHEST_PROTOCOL)
            pickle.dump(self.unicornEngine.get_context(), file, pickle.HIGHEST_PROTOCOL)
            file.close()
            uemu_log("Project saved to %s" % filePath)

    # --- MAIN MENU

    MENU_ITEMS = []

    def 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 handle_menu_action(self, action):
        [x.handler() for x in self.MENU_ITEMS if x.action == action]

    def register_menu_actions(self):
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":start",             self.emu_start,             "Start",                      "Start emulation",           None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":run",               self.emu_run,               "Run",                        "Run",                       None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":step",              self.emu_step,              "Step",                       "Step to next instruction",  "SHIFT+CTRL+ALT+S",     True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":stop",              self.emu_stop,              "Stop",                       "Stop emulation",            None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":reset",             self.emu_reset,             "Reset",                      "Reset emulation",           None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem("-",                                     self.do_nothing,            "",                           None,                        None,                   True    ))
        self.add_custom_menu()
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":jmp_pc",            self.jump_to_pc,            "Jump to PC",                 "Jump to PC",                "SHIFT+CTRL+ALT+J",     True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":set_pc",            self.set_pc,                "Set PC",                     "Set PC",                    None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":cng_cpu",           self.change_cpu_context,    "Change CPU Context",         "Change CPU Context",        None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem("-",                                     self.do_nothing,            "",                           None,                        None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":ctl_view",          self.show_controls,         "Show Controls",              "Show Control Window",       None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":cpu_view",          self.show_cpu_context,      "Show CPU Context",           "Show CPU Registers",        None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":cpu_ext_view",      self.show_cpu_ext_context,  "Show CPU Extended Context",  "Show Extended Registers",   None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":stack_view",        self.show_stack_view,       "Show Stack View",            "Show Stack View",           None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":mem_view",          self.show_memory,           "Show Memory Range",          "Show Memory Range",         None,                   True    ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":mem_map",           self.show_mapped,           "Show Mapped Memory",         "Show Mapped Memory",        None,                   False   ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":fetch_segs",        self.fetch_segments,        "Fetch Segments",             "Fetch Segments",            None,                   False   ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem("-",                                     self.do_nothing,            "",                           None,                        None,                   False   ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":load_prj",          self.load_project,          "Load Project",               "Load Project",              None,                   False   ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":save_prj",          self.save_project,          "Save Project",               "Save Project",              None,                   False   ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":settings",          self.show_settings,         "Settings",                   "Settings",                  None,                   False   ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem("-",                                     self.do_nothing,            "",                           None,                        None,                   False   ))
        self.MENU_ITEMS.append(UEMU_HELPERS.MenuItem(self.plugin_name + ":unload",            self.unload_plugin,         "Unload Plugin",              "Unload Plugin",             None,                   False   ))

        for item in self.MENU_ITEMS:
            if item.action == "-":
                continue
            self.register_new_action(item.action, item.title, UEMU_HELPERS.IdaMenuActionHandler(self, item.action), item.shortcut, item.tooltip,  -1)

    def unregister_menu_actions(self):
        for item in self.MENU_ITEMS:
            if item.action == "-":
                continue
            unregister_action(item.action)

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

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

    def add_custom_menu(self): # used by extensions
        pass

    # --- POPUP MENU

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

    def 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.MENU_ITEMS:
                if item.popup:
                    attach_action_to_popup(widget, popup_handle, item.action, self.plugin_name + "/")
    
    # IDA 6.x
    def finish_populating_tform_popup(self, form, popup_handle):
        if get_tform_type(form) == BWN_DISASM:
            for item in self.MENU_ITEMS:
                if item.popup:
                    attach_action_to_popup(form, popup_handle, item.action, self.plugin_name + "/")

    # --- DELEGATES

    def contol_view_closed(self):
        self.controlView = None

    def context_view_closed(self):
        self.cpuContextView = None

    def ext_context_view_closed(self):
        self.cpuExtContextView = None

    def memory_view_closed(self, viewid):
        del self.memoryViews[viewid]

    def stack_view_closed(self):
        self.stackView = None

    def update_context(self, address, context):
        # Update CPU Context Views
        if self.cpuContextView is not None:
            self.cpuContextView.SetContent(address, context)
        if self.cpuExtContextView is not None:
            self.cpuExtContextView.SetContent(address, context)

        # Update Stack View
        if self.stackView is not None:
            self.stackView.SetContent(context)

        # Update Memory Views
        for viewid in self.memoryViews:
            self.memoryViews[viewid].SetContent(context)

    # --- METHODS

    def emu_start(self):
        if self.unicornEngine.is_active():
            uemu_log("Emulator is already active")
            return

        if self.unicornEngine.is_running():
            uemu_log("Emulator is running")
            return

        self.unicornEngine.run_from(IDAAPI_ScreenEA())
        uemu_log("Emulation started")

    def emu_run(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.unicornEngine.is_running():
            uemu_log("Emulator is running")
            return

        self.unicornEngine.run()

    def emu_step(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.unicornEngine.is_running():
            uemu_log("Emulator is running")
            return
        count = 1
        if UEMU_HELPERS.is_alt_pressed():
            count = IDAAPI_AskLong(1, "Enter Step Count")
            if count is None or count < 0:
                return
            if count == 0:
                count = 1

        self.unicornEngine.step(count)

    def emu_stop(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if not self.unicornEngine.is_running():
            uemu_log("Emulator is not running")
            return

        self.unicornEngine.interrupt()

    def emu_reset(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return
        
        self.unicornEngine.reset()

        self.close_windows()

        self.memoryViews.clear()
        uemu_log("Emulation reset")

    def jump_to_pc(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.unicornEngine.is_running():
            uemu_log("Emulator is running")
            return

        self.unicornEngine.jump_to_pc()

    def set_pc(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.unicornEngine.is_running():
            uemu_log("Emulator is running")
            return
        
        self.unicornEngine.set_pc(IDAAPI_ScreenEA())

    def change_cpu_context(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.unicornEngine.is_running():
            uemu_log("Emulator is running")
            return

        self.unicornEngine.set_cpu_context()
        if self.cpuContextView is not None:
            self.cpuContextView.SetContent(self.unicornEngine.pc, self.unicornEngine.mu)
        if self.cpuExtContextView is not None:
            self.cpuExtContextView.SetContent(self.unicornEngine.pc, self.unicornEngine.mu)

    def show_controls(self):
        if self.controlView is None:
            self.controlView = uEmuControlView(self)
            self.controlView.Show("uEmu Control")

    def show_cpu_context(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.cpuContextView is None:
            self.cpuContextView = uEmuCpuContextView(self, False)
            self.cpuContextView.Create("uEmu CPU Context")
            self.cpuContextView.SetContent(self.unicornEngine.pc, self.unicornEngine.mu)
            self.cpuContextView.Show()
            self.cpuContextView.Refresh()

    def show_cpu_ext_context(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if UEMU_HELPERS.get_arch() not in ["armle", "armbe", "arm64le", "arm64be"]:
            return

        if self.cpuExtContextView is None:
            self.cpuExtContextView = uEmuCpuContextView(self, True)
            self.cpuExtContextView.Create("uEmu CPU Extended Context")
            self.cpuExtContextView.SetContent(self.unicornEngine.pc, self.unicornEngine.mu)
            self.cpuExtContextView.Show()
            self.cpuExtContextView.Refresh()

    def show_memory(self, address = 0, size = 16):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        memRangeDlg = uEmuMemoryRangeDialog()
        memRangeDlg.Compile()
        memRangeDlg.mem_addr.value = address
        memRangeDlg.mem_size.value = size
        ok = memRangeDlg.Execute()
        if ok == 1:
            mem_addr = memRangeDlg.mem_addr.value
            mem_size = memRangeDlg.mem_size.value
            mem_cmnt = memRangeDlg.mem_cmnt.value

            if mem_addr not in self.memoryViews:
                if not self.unicornEngine.is_memory_mapped(mem_addr):
                    ok = IDAAPI_AskYN(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.unicornEngine.map_empty(mem_addr, mem_size)
                    elif ok == 1:
                        if not self.unicornEngine.map_binary(mem_addr, mem_size):
                            return
                    else:
                        return

                self.memoryViews[mem_addr] = uEmuMemoryView(self, mem_addr, mem_size)
                self.memoryViews[mem_addr].Create("uEmu Memory [ " + mem_cmnt + " ]")
                self.memoryViews[mem_addr].SetContent(self.unicornEngine.mu)
            self.memoryViews[mem_addr].Show()
            self.memoryViews[mem_addr].Refresh()

    def show_stack_view(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.stackView is None:
            self.stackView = uEmuStackView(self)
            self.stackView.Create("uEmu Stack View")
            self.stackView.SetContent(self.unicornEngine.mu)
            self.stackView.Show()
            self.stackView.Refresh()

    def show_mapped(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        mappeduEmuMemoryView = uEmuMappeduMemoryView(self, self.unicornEngine.get_mapped_memory())
        mappeduEmuMemoryView.show()

    def fetch_segments(self):
        if not self.unicornEngine.is_active():
            uemu_log("Emulator is not active")
            return

        if self.unicornEngine.is_running():
            uemu_log("Emulator is running")
            return

        self.unicornEngine.fetch_segments()

    def show_settings(self):
        kSettingsMask_FollowPC  = 0x1
        kSettingsMask_ForceCode = 0x2
        kSettingsMask_TraceInst = 0x4
        kSettingsMask_LazyMapping = 0x8

        settingsDlg = uEmuSettingsDialog()
        settingsDlg.Compile()
        settingsDlg.emu_group.value = 0

        if self.settings["follow_pc"]: settingsDlg.emu_group.value = settingsDlg.emu_group.value | kSettingsMask_FollowPC
        if self.settings["force_code"]: settingsDlg.emu_group.value = settingsDlg.emu_group.value | kSettingsMask_ForceCode
        if self.settings["trace_inst"]: settingsDlg.emu_group.value = settingsDlg.emu_group.value | kSettingsMask_TraceInst
        if self.settings["lazy_mapping"]: settingsDlg.emu_group.value = settingsDlg.emu_group.value | kSettingsMask_LazyMapping

        ok = settingsDlg.Execute()
        if ok == 1:
            self.settings["follow_pc"] = True if settingsDlg.emu_group.value & kSettingsMask_FollowPC else False
            self.settings["force_code"] = True if settingsDlg.emu_group.value & kSettingsMask_ForceCode else False
            self.settings["trace_inst"] = True if settingsDlg.emu_group.value & kSettingsMask_TraceInst else False
            self.settings["lazy_mapping"] = True if settingsDlg.emu_group.value & kSettingsMask_LazyMapping else False

    def close_windows(self):
        if self.cpuContextView is not None:
            self.cpuContextView.Close()

        if self.cpuExtContextView is not None:
            self.cpuExtContextView.Close()

        if self.stackView is not None:
            self.stackView.Close()

        for viewid in list(self.memoryViews):
            self.memoryViews[viewid].Close()

def PLUGIN_ENTRY():
    return uEmuPlugin()

if UEMU_USE_AS_SCRIPT:
    if __name__ == '__main__':
        uEmu = uEmuPlugin()
        uEmu.init()
        uEmu.run()


