# plugins

repository: computerline1z/IDAPatcher6.1
file: idapatcher.py
description: IDA Patcher for v6.1
wanted_name: IDA Patcher
comment: Enhances manipulation and application of patched bytes.
created_at: 2014-03-31 04:44:29+00:00
pushed_at: 2014-03-31 12:04:13+00:00
url:
https://github.com/computerline1z/IDAPatcher6.1/blob/b7cc364b58b8deadf7b9c60d05218ab7f75227c8/idapatcher.py

content:
#!/usr/bin/env python
#
# IDA Patcher is a plugin for Hex-Ray's IDA Pro disassembler designed to 
# enhance IDA's ability to patch binary files and memory. The plugin is 
# useful for tasks related to malware analysis, exploit development as well
# as bug patching. IDA Patcher blends into the standard IDA user interface
# through the addition of a subview and several menu items. 
#
# VERSION 0.1.1
#
# Copyright (C) 2013 Peter Kacherginsky
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met: 
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer. 
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution. 
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import idaapi
import idautils
import idc

from idaapi import Form, Choose2, plugin_t

import os
import shutil
import struct
import binascii

IDAPATCHER_VERSION = "0.1.1"

#--------------------------------------------------------------------------
# Forms
#--------------------------------------------------------------------------
class PatchRestoreForm(Form):
    """
    Form to aid in restoring patched bytes to their original values.
    """
    def __init__(self, addr_str, fpos_str, patch_str, org_str):
        Form.__init__(self, 
r"""BUTTON YES* Restore
BUTTON CANCEL Cancel
Restore patch bytes

Address        {strAddr}
File offset    {strFpos}
<:{strOrg}>
""", {
        'strAddr': Form.StringLabel(addr_str),
        'strFpos': Form.StringLabel(fpos_str),
        'strOrg': Form.MultiLineTextControl(text=org_str, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_READONLY),

        })

        self.Compile()

#--------------------------------------------------------------------------
class PatchEditForm(Form):
    """
    Form to edit patched bytes.
    """
    def __init__(self, addr_str, fpos_str, patch_str, org_str):
        Form.__init__(self,
r"""Edit patch bytes

Address        {strAddr}
File offset    {strFpos}
<:{strPatch}>
""", {
        'strAddr':  Form.StringLabel(addr_str),
        'strFpos':  Form.StringLabel(fpos_str),
        'strPatch': Form.MultiLineTextControl(text=patch_str, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT),
        })

        self.Compile()

#--------------------------------------------------------------------------
class PatchApplyForm(Form):
    """
    Form to prompt for target file, backup file, and the address
    range to save patched bytes.
    """
    def __init__(self, start_ea, end_ea, org_file, bkp_file):
        Form.__init__(self,
r"""Apply patches to input file

{FormChangeCb}
<##Start EA   :{intStartEA}>
<##End EA     :{intEndEA}>
<##Input file :{orgFile}>
<##Backup file:{bkpFile}>

<##Create backup:{rBackup}>
<##Restore original bytes:{rRestore}>{cGroup1}>
""", {
        'intStartEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=start_ea),
        'intEndEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=end_ea),
        'orgFile': Form.FileInput(swidth=50, open=True, value=org_file),
        'bkpFile': Form.FileInput(swidth=50, open=True, value=bkp_file),
        'cGroup1': Form.ChkGroupControl(("rBackup", "rRestore")),
        'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
        })

        self.Compile()

    def OnFormChange(self, fid):
        # Set initial state
        if fid == -1:
            self.EnableField(self.bkpFile, False)

        # Toggle backup checkbox
        elif fid == self.rBackup.id:
            self.rBackup.checked = not self.rBackup.checked
            self.EnableField(self.bkpFile, self.rBackup.checked)

        # Toggle restore checkbox
        elif fid == self.rRestore.id:
            self.rRestore.checked = not self.rRestore.checked

        return 1

#--------------------------------------------------------------------------
class PatchFillForm(Form):
    """
    Form to fill a range of addresses with a specified byte value.
    """
    def __init__(self, start_ea, end_ea, fill_value):

        Form.__init__(self,
r"""BUTTON YES* Fill
Fill bytes

<##Start EA   :{intStartEA}>
<##End EA     :{intEndEA}>
<##Value      :{intPatch}>
""", {
        'intStartEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=start_ea),
        'intEndEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=end_ea),
        'intPatch': Form.NumericInput(swidth=40,tp=Form.FT_HEX,value=fill_value),
        })

        self.Compile()

class DataImportForm(Form):
    """
    Form to import data of various types into selected area.
    """
    def __init__(self, start_ea, end_ea):
        Form.__init__(self,
r"""BUTTON YES* Import
Import data

{FormChangeCb}
<##Start EA   :{intStartEA}>
<##End EA     :{intEndEA}>

Import type:                    Patching options:
<hex string:{rHex}><##Trim to selection:{cSize}>{cGroup}>
<string literal:{rString}>
<binary file:{rFile}>{rGroup}>

<:{strPatch}>
<##Import BIN file:{impFile}>
""", {        
        'intStartEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=start_ea),
        'intEndEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=end_ea),

        'cGroup': Form.ChkGroupControl(("cSize",)),
        'rGroup': Form.RadGroupControl(("rHex", "rString", "rFile")),

        'strPatch': Form.MultiLineTextControl(swidth=80, flags=Form.MultiLineTextControl.TXTF_FIXEDFONT),
        'impFile': Form.FileInput(swidth=50, open=True),

        'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
        })

        self.Compile()

    def OnFormChange(self, fid):
        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.strPatch)
            self.EnableField(self.strPatch, True)
            self.EnableField(self.impFile, False)

        # Form OK pressed
        elif fid == -2:
            pass

        # Form from text box
        elif fid == self.rHex.id or fid == self.rString.id:
            self.SetFocusedField(self.strPatch)
            self.EnableField(self.strPatch, True)
            self.EnableField(self.impFile, False)

        # Form import from file
        elif fid == self.rFile.id:
            self.SetFocusedField(self.rFile)
            self.EnableField(self.impFile, True)
            self.EnableField(self.strPatch, False)

        return 1

#--------------------------------------------------------------------------
# Chooser
#--------------------------------------------------------------------------
class PatchView(Choose2):
    """
    Chooser class to display and manage patched bytes in the database.
    """
    def __init__(self):
        Choose2.__init__(self,
                         "Patches",
                         [ ["Address",  10 | Choose2.CHCOL_HEX], 
                           ["Name",     18 | Choose2.CHCOL_PLAIN], 
                           ["Size",      4 | Choose2.CHCOL_DEC], 
                           ["Modified", 10 | Choose2.CHCOL_HEX],
                           ["Original", 10 | Choose2.CHCOL_HEX], 
                           ["Comment",  30 | Choose2.CHCOL_PLAIN]
                         ],
                         flags = Choose2.CH_MULTI_EDIT)

        self.popup_names = ["Insert", "Delete", "Edit", "Refresh"]
        
        self.icon = 47

        # Items for display and corresponding data
        # NOTE: Could become desynchronized, so to avoid this
        #       refresh the view after each change.
        self.items = []
        self.items_data  = []

        # Initialize/Refresh the view
        self.refreshitems()

        # Data members
        self.patch_file = None
        self.restore = False

        # Command callbacks
        self.cmd_apply_patches = None
        self.cmd_restore_bytes = None

    def show(self):
        # Attempt to open the view
        if self.Show() < 0: return False

        # Add extra context menu commands
        # NOTE: Make sure you check for duplicates.
        if self.cmd_apply_patches == None:
            self.cmd_apply_patches = self.AddCommand("Apply patches to input file...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_NO_SELECTION, icon=27)
        if self.cmd_restore_bytes == None:
            self.cmd_restore_bytes = self.AddCommand("Restore original byte(s)...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=139)

        return True

    # Patch byte visitor callback to apply the patches
    # NOTE: Only bytes with fpos > -1 can be applied.
    def apply_patch_byte(self, ea, fpos, org_val, patch_val):
        if fpos != -1:
            self.patch_file.seek(fpos)

            if self.restore:
                self.patch_file.write(struct.pack('B', org_val))
            else:
                self.patch_file.write(struct.pack('B', patch_val))

        return 0

    # Patch byte visitor callback to collect and aggregate bytes
    def get_patch_byte(self, ea, fpos, org_val, patch_val):

        # Aggregate contiguous bytes (base ea + length)
        # NOTE: Looking at the last item [-1] is sufficient
        #       since we are dealing with sorted data.
        if len(self.items_data) and (ea - self.items_data[-1][0] == self.items_data[-1][2]):

            # Increment length
            self.items_data[-1][2] += 1
            self.items[-1][2] = str(self.items_data[-1][2])

            # Append patched bytes
            self.items_data[-1][3].append(patch_val)
            self.items[-1][3] = " ".join(["%02X" % x for x in self.items_data[-1][3]])

            # Append original bytes
            self.items_data[-1][4].append(org_val)
            self.items[-1][4] =  " ".join(["%02X" % x for x in self.items_data[-1][4]])


        # Add new patch byte to the list
        else:

            name = SegName(ea)

            if GetFunctionName(ea) or Name(ea):
                name += ": %s" % GetFunctionName(ea) or Name(ea)


            comment = Comment(ea) or RptCmt(ea) or ""
            # DATA STORAGE FORMAT:      address, function / fpos, len,    patched byte(s), original byte(s), comments
            self.items.append(     ["%08X" % ea,            name, "1", "%02X" % patch_val, "%02X" % org_val, comment])
            self.items_data.append([         ea,            fpos,   1,        [patch_val],        [org_val], None]   )

        return 0

    def refreshitems(self):
        self.items_data = []
        self.items = []
        idaapi.visit_patched_bytes(0, idaapi.BADADDR, self.get_patch_byte)

    def OnCommand(self, n, cmd_id):

        # Apply patches to a file
        if cmd_id == self.cmd_apply_patches:

            # Set initial start/end EA values
            start_ea = 0x0
            end_ea = idaapi.cvar.inf.maxEA

            # Set initial output file values
            org_file = GetInputFilePath()
            bkp_file = "%s.bak" % org_file

            # Create the form
            f = PatchApplyForm(start_ea, end_ea, org_file, bkp_file)

            # Execute the form
            ok = f.Execute()
            if ok == 1:
                # Get restore checkbox
                self.restore = f.rRestore.checked

                # Get updated ea max/min
                start_ea = f.intStartEA.value
                end_ea = f.intEndEA.value

                # Get updated file path
                new_org_file = f.orgFile.value

                # Backup the file before replacing
                if f.rBackup.checked:
                    bkp_file = f.bkpFile.value
                    shutil.copyfile(org_file, bkp_file)

                # Apply patches
                try:
                    self.patch_file = open(new_org_file,'rb+')
                except Exception, e:
                    idaapi.warning("Cannot update file '%s'" % new_org_file)
                else:
                    r = idaapi.visit_patched_bytes(start_ea, end_ea, self.apply_patch_byte)
                    self.patch_file.close()
                    self.restore = False

                    # Update db input file, so we are working
                    # with a patched version.
                    #if not org_file == new_org_file:
                    #    idaapi.set_root_filename(new_org_file)
                    #    org_file = new_org_file

            # Dispose the form
            f.Free()

        # Restore selected byte(s)
        elif cmd_id == self.cmd_restore_bytes:

            # Empty list

            if n == -2 or n ==-3:
                return 1

            elif not len(self.items) > 0:
                idaapi.warning("There are no patches to restore.")
                return 1

            # Nothing selected
            elif n == -1:
                idaapi.warning("Please select bytes to restore.")
                return 1

            ea = self.items_data[n][0]
            fpos =  self.items_data[n][1]
            buf = self.items_data[n][4]

            addr_str = "%#x" % ea
            fpos_str = "%#x" % fpos if fpos != -1 else "N/A"  
            patch_str = self.items[n][3]
            org_str = self.items[n][4]

            # Create the form
            f = PatchRestoreForm(addr_str, fpos_str, patch_str, org_str)

            # Execute the form
            ok = f.Execute()
            if ok == 1:

                # Restore original bytes
                idaapi.put_many_bytes(ea, struct.pack("B"*len(buf), *buf))

                # Refresh all IDA views
                self.refreshitems()
                idaapi.refresh_idaview_anyway()

            # Dispose the form
            f.Free()

        return 1

    def OnClose(self):
        self.cmd_apply_patches = None
        self.cmd_restore_bytes = None

    def OnEditLine(self, n):

        # Empty list
        if n == -1:
            return

        # Multiselect START_SEL/END_SEL protocol
        if n == -2 or n ==-3:
            return

        ea = self.items_data[n][0]
        fpos =  self.items_data[n][1]
        patch_buf = self.items_data[n][3]
        orig_buf = self.items_data[n][4]

        addr_str = "%#x" % ea
        fpos_str = "%#x" % fpos if fpos != -1 else "N/A"     
        patch_str = self.items[n][3]
        org_str = self.items[n][4]    

        # Create the form
        f = PatchEditForm(addr_str, fpos_str, patch_str, org_str)

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            # Convert hex bytes to binary
            buf = f.strPatch.value
            buf = buf.replace(' ','')       # remove spaces
            buf = buf.replace('\\x','')     # remove '\x' prefixes
            buf = buf.replace('0x','')      # remove '0x' prefixes
            try:
                buf = binascii.unhexlify(buf)   # convert to bytes
            except Exception, e:
                idaapi.warning("Invalid input: %s" % e)
                f.Free()
                return

            # Restore original bytes first
            idaapi.put_many_bytes(ea, struct.pack("B"*len(orig_buf), *orig_buf))

            # Now apply newly patched bytes
            idaapi.patch_many_bytes(ea, buf)

            # Refresh all IDA views
            self.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()

    def OnSelectLine(self, n):
        idaapi.jumpto(self.items_data[n][0])

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

    def OnGetIcon(self, n):

        # Empty list
        if not len(self.items) > 0:
            return -1

        if self.items_data[n][1] == -1:
            return 138
        else:
            return 137

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

    def OnRefresh(self, n):
        self.refreshitems()
        return n

    def OnActivate(self):
        self.refreshitems()

#--------------------------------------------------------------------------
# Manager
#--------------------------------------------------------------------------
class PatchManager():
    """ Class that manages GUI forms and patching methods of the plugin. """
    
    def __init__(self): 
        self.addmenu_item_ctxs = list()
        self.patch_view = PatchView()

    #--------------------------------------------------------------------------
    # Menu Items
    #--------------------------------------------------------------------------
    def add_menu_item_helper(self, menupath, name, hotkey, flags, pyfunc, args):

        # add menu item and report on errors
        addmenu_item_ctx = idaapi.add_menu_item(menupath, name, hotkey, flags, pyfunc, args)
        if addmenu_item_ctx is None:
            return 1
        else:
            self.addmenu_item_ctxs.append(addmenu_item_ctx)
            return 0

    def add_menu_items(self):

        if self.add_menu_item_helper("View/Open subviews/Problems", "Patches", "", 1, self.show_patches_view, None): return 1
        if self.add_menu_item_helper("Edit/Patch program/", "Edit selection...", "", 0, self.show_edit_form, None):  return 1
        if self.add_menu_item_helper("Edit/Patch program/", "Fill selection...", "", 0, self.show_fill_form, None):  return 1
        if self.add_menu_item_helper("Edit/Export data...", "Import data...", "Shift-I", 1, self.show_import_form, None):   return 1

        return 0

    def del_menu_items(self):
        for addmenu_item_ctx in self.addmenu_item_ctxs:
            idaapi.del_menu_item(addmenu_item_ctx)

    #--------------------------------------------------------------------------
    # View Callbacks
    #--------------------------------------------------------------------------   
    
    # Patches View
    def show_patches_view(self):
        self.patch_view.show()

    # Patches Edit Dialog
    def show_edit_form(self):
        selection, start_ea, end_ea = idaapi.read_selection()
        
        if not selection:
            start_ea = idaapi.get_screen_ea()
            end_ea = start_ea + 1

        buf_len = end_ea - start_ea
        buf = get_many_bytes(start_ea, buf_len) or "\xFF"*buf_len
        buf_str = " ".join(["%02X" % ord(x) for x in buf])

        fpos = idaapi.get_fileregion_offset(start_ea)

        addr_str = "%#X" % start_ea
        fpos_str = "%#x" % fpos if fpos != -1 else "N/A" 

        f = PatchEditForm(addr_str, fpos_str, buf_str, buf_str)

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            # Convert hex bytes to binary
            buf = f.strPatch.value
            buf = buf.replace(' ','')       # remove spaces
            buf = buf.replace('\\x','')     # remove '\x' prefixes
            buf = buf.replace('0x','')      # remove '0x' prefixes
            try:
                buf = binascii.unhexlify(buf)   # convert to bytes
            except Exception, e:
                idaapi.warning("Invalid input: %s" % e)
                f.Free()
                return

            # Now apply newly patched bytes
            idaapi.patch_many_bytes(start_ea, buf)

            # Refresh all IDA views
            self.patch_view.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()

    # Fill range with a value form
    def show_fill_form(self):
        selection, start_ea, end_ea = idaapi.read_selection()
        
        if not selection:
            start_ea = idaapi.get_screen_ea()
            end_ea = start_ea + 1
        
        # Default fill value
        fill_value = 0x00

        # Create the form
        f = PatchFillForm(start_ea, end_ea, fill_value)

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            # Get updated values
            start_ea = f.intStartEA.value
            end_ea = f.intEndEA.value
            fill_value = f.intPatch.value

            # Now apply newly patched bytes
            # NOTE: fill_value is expected to be one byte
            #       so if a user provides a larger patch_byte()
            #       will trim the value as expected.


            for ea in range(start_ea, end_ea):
                idaapi.patch_byte(ea, fill_value)

            # Refresh all IDA views
            self.patch_view.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()

    # Import data form
    def show_import_form(self):
        selection, start_ea, end_ea = idaapi.read_selection()

        if not selection:
            start_ea = idaapi.get_screen_ea()
            end_ea = start_ea + 1

        # Create the form
        f = DataImportForm(start_ea, end_ea);

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            start_ea = f.intStartEA.value
            end_ea = f.intEndEA.value

            if f.rFile.selected:
                imp_file = f.impFile.value

                try:
                    f_imp_file = open(imp_file,'rb+')
                except Exception, e:
                    idaapi.warning("File I/O error({0}): {1}".format(e.errno, e.strerror))
                    return
                else:
                    buf = f_imp_file.read()
                    f_imp_file.close()

            else:

                buf = f.strPatch.value

                # Hex values, unlike string literal, needs additional processing
                if f.rHex.selected:
                    buf = buf.replace(' ','')       # remove spaces
                    buf = buf.replace('\\x','')     # remove '\x' prefixes
                    buf = buf.replace('0x','')      # remove '0x' prefixes
                    try:
                        buf = binascii.unhexlify(buf)   # convert to bytes
                    except Exception, e:
                        idaapi.warning("Invalid input: %s" % e)
                        f.Free()
                        return

            if not len(buf):
                idaapi.warning("There was nothing to import.")
                return

            # Trim to selection if needed:
            if f.cSize.checked:
                buf_size = end_ea - start_ea
                buf = buf[0:buf_size]

            # Now apply newly patched bytes
            idaapi.patch_many_bytes(start_ea, buf)

            # Refresh all IDA views
            self.patch_view.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()
       
#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class idapatcher_t(plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = "Enhances manipulation and application of patched bytes."
    help = "Enhances manipulation and application of patched bytes."
    wanted_name = "IDA Patcher"
    wanted_hotkey = ""

    def init(self):  
        global idapatcher_manager

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

            idapatcher_manager = PatchManager()
            if idapatcher_manager.add_menu_items():
                print "[!] Failed to initialize IDA Patcher."
                idapatcher_manager.del_menu_items()
                del idapatcher_manager
                return idaapi.PLUGIN_SKIP
            else:
                print("[*] IDA Patcher v%s (c) Peter Kacherginsky <iphelix@thesprawl.org>" % IDAPATCHER_VERSION)
            
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        global idapatcher_manager
        idapatcher_manager.show_patches_view()

    # NOTE: term() gets called immediately after run()
    #       as the plugin is expected to complete execution
    #       so just pass as this plugin is expected to run
    #       for the duration of IDA session.
    def term(self):
        pass
        

def PLUGIN_ENTRY():
    return idapatcher_t()

#--------------------------------------------------------------------------
# Script / Testing
#--------------------------------------------------------------------------
def idapatcher_main():
    global idapatcher_manager

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

    idapatcher_manager = PatchManager()
    idapatcher_manager.add_menu_items()
    idapatcher_manager.show_patches_view()

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


repository: BinSigma/BinSourcerer
file: BinSourcerer/IdaProTextExtractor_HRC.py
description: BinSourcerer
wanted_name: BinSourcerer Features Extractor
comment: Extract all the available features
created_at: 2014-02-26 21:59:45+00:00
pushed_at: 2022-03-01 15:52:39+00:00
url:
https://github.com/BinSigma/BinSourcerer/blob/5958abc1d0091d60638291ebc1118990b4fa97dc/BinSourcerer/IdaProTextExtractor_HRC.py

content:
#!/usr/bin/env python
# -*-coding: utf8 -*-
#
################################################################################
#
#  Copyright 2014 Concordia University
# 
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
################################################################################
# Description:
#   This plugin extracts certain featuers from IDA Pro disassembly and
#   generates a feature file, which is used in an online code search process.
#   This file is part of the BinSourcerer framework.
#
# Status: Debug
#
################################################################################




DIRECTION_FORW = 0x03

import idautils
import os

from idaapi import *
import hashlib
import idc

#Offline Analyzer (OA), Windows function names and a brief description.
WinMalFuncInSight = {
    'accept':'The program may listen for incoming connections on a socket.',
    'AdjustTokenPrivileges':'The program may modify specific access privileges.',
    'AttachThreadInput':'The program may contain keylogging or spyware functionality.',
    'bind':'The program may listen for incoming connections.',
    'BitBlt':'The program may capture screenshots.',
    'CallNextHookEx':'The code calls the next hook in the chain set by SetWindowsHookEx.',
    'CertOpenSystemStore':'The program may access the certificates stored on the local system.',
    'CheckRemoteDebuggerPresent':'The code may contain an anti-debugging technique.',
    'CoCreateInstance':'The code creates a COM object.',
    'connect':'The program may connect to a remote command-and-control server.',
    'ConnectNamedPipe':'The code may create a server pipe for interprocess communication',
    'ControlService':'The code may start, stop, modify, or send a signal to a running service.',
    'CreateFile':'The program may create a new file or open an existing file.',
    'CreateFileMapping':'The code may read and modify PE files by creating handles to file mappings.',
    'CreateMutex':'Only one instance of the program runs on a system.',
    'CreateProcess':'The code may create and launch a new process.',
    'CreateRemoteThread':'The program may inject code into a different process.',
    'CreateService':'The program may create an autorun service or load kernel drivers.',
    'CreateToolhelp32Snapshot':'The code may create a snapshot of processes, heaps, threads, and modules.',
    'CryptAcquireContext':'The code may initialize the use of Windows encryption.',
    'DeviceIoControl':'The code may send control messages from user space to kernel space.',
    'DllCanUnloadNow':'The program implements a COM server. (Exported function)',
    'DllGetClassObject':'The program implements a COM server. (Exported function)',
    'DllInstall':'The program implements a COM server. (Exported function)',
    'DllRegisterServer':'The program implements a COM server. (Exported function)',
    'DllUnregisterServer':'The program implements a COM server. (Exported function)',
    'EnableExecuteProtectionSupport':'The code may modify the Data Execution Protection (DEP) settings of the host.',
    'EnumProcesses':'The code may enumerate through processes to find a process to inject into.',
    'EnumProcessModules':'The code may enumerate the loaded modules for injection.',
    'FindFirstFile':'The program may search through a directory and enumerate the filesystem.',
    'FindNextFile':'The program may search through a directory and enumerate the filesystem.',
    'FindResource':'The code may search for a resource in an executable or loaded DLL',
    'FindWindow':'The code searches for an open window on the desktop. (It may also contain anti-debugging)',
    'FtpPutFile':'The program may upload files to a remote FTP server.',
    'GetAdaptersInfo':'The code may obtain information about the MAC, network adapters on the system.',
    'GetAsyncKeyState':'The code may contain keylogger functionality.',
    'GetDC':'The code returns a handle to a device context for a window or the whole screen. (Screen Capture)',
    'GetForegroundWindow':'The code returns a handle to the window currently in the foreground of the desktop. (Keylogger)',
    'gethostbyname':'The code may perform a DNS lookup on a particular hostname. (C&C)',
    'gethostname':'The code may retrieve the hostname of the computer. (Backdoor)',
    'GetKeyState':'The code may obtain the status of a particular key on the keyboard. (Keylogger)',
    'GetModuleFilename':'The code may return the filename of a module that is loaded in the current process.',
    'GetModuleHandle':'The code may obtain a handle to an already loaded module. (Inject Code)',
    'GetProcAddress':'The code may retrieve the address of a function in a DLL loaded into memory.',
    'GetStartupInfo':'The code may get details about how the current process was configured to run',
    'GetSystemDefaultLangId':'The code may return the default language settings for the system.',
    'GetTempPath':'The code may return the temporary file path.',
    'GetThreadContext':'The code may read the context of a thread. (register values and current state)',
    'GetTickCount':'The code may retrieve the number of milliseconds since bootup. (Anti-debugging)',
    'GetVersionEx':'The code may return information about which version of Windows is currently running.',
    'GetWindowsDirectory':'The code may get the file path to the Windows directory.',
    'inet_addr':'The code may process IP addresses for network connectivity.',
    'InternetOpen':'The code may initialize the high-level Internet access functions from WinINet',
    'InternetOpenUrl':'The program may open a specific URL for a connection using FTP, HTTP, or HTTPS.',
    'InternetReadFile':'The program may read data from a previously opened URL.',
    'InternetWriteFile':'The program may write data to a previously opened URL.',
    'IsDebuggerPresent':'The code may check to see if the current process is being debugged. (Anti-debugging)',
    'IsNTAdmin':'The code checks if the user has administrator privileges.',
    'IsWoW64Process':'The code may check if it is running on a 64-bit operating system.',
    'LdrLoadDll':'The code uses a low-level function to load a DLL into a process and it may attempt to be stealthy.',
    'LoadLibrary':'The code loads a DLL into a process that may not have been loaded when the program started.',
    'LoadResource':'The code may load a resource from a PE file into memory.',
    'LsaEnumerateLogonSessions':'The code may enumerate through logon sessions on the current system.',
    'MapViewOfFile':'The code maps a file into memory and makes the contents of the file accessible via memory addresses.',
    'MapVirtualKey':'The code may translate a virtual-key code into a character value. (Keylogger)',
    'MmGetSystemRoutineAddress':'The code may retrieve the address of a function from another module. (Kernel code)',
    'Module32First':'The code may enumerate through modules loaded into a process. (Injector)',
    'Module32Next':'The code mya enumerate through modules loaded into a process. (Injector)',
    'NetScheduleJobAdd':'The code may submit a request for a program to be run at a specified date and time in the future.',
    'NetShareEnum':'The code may enumerate network shares.',
    'NtQueryDirectoryFile':'The code may get information about files in a directory. (Rootkit)',
    'NtQueryInformationProcess':'The code may return various information about a specified process. (Anti-debugging)',
    'NtSetInformationProcess':'The code may change the privilege level of a program or to bypass Data Execution Prevention (DEP).',
    'OleInitialize':'The code initialize the COM library.',
    'OpenMutex':'The code may open a handle to a mutual exclusion object to ensure that only a single instance is running.',
    'OpenProcess':'The code may open a handle to another process to read/write or to inject code.',
    'OpenSCManager':'The code may open a handle to the service control manager.',
    'OutputDebugString':'The code may output a string to a debugger if one is attached. (Anti-debugging)',
    'PeekNamedPipe':'The code may copy data from a named pipe without removing data from the pipe. (Reverse shell)',
    'Process32First':'The code may enumerate through processes to find a process to inject into.',
    'Process32Next':'The code may enumerate through processes to find a process to inject into.',
    'QueryPerformanceCounter':'The code may retrieve the value of the hardware-based performance counter. (Anti-debugging)',
    'QueueUserAPC':'The code may execute code for a different thread and inject code into another process.',
    'ReadProcessMemory':'The code may read the memory of a remote process.',
    'recv':'The code may receive data from a remote command-and-control server.',
    'RegisterHotKey':'The code may register a handler to be notified anytime a user enters a particular key combination.',
    'RegOpenKey':'The code may open a handle to a registry key for reading and editing.',
    'ResumeThread':'The code may resume a previously suspended thread for injection.',
    'RtlCreateRegistryKey':'The code may create a registry from kernel-mode code.',
    'RtlWriteRegistryValue':'The code may write a value to the registry from kernel-mode code.',
    'SamIConnect':'The code may connect to the Security Account Manager in order to make future calls that access credential information. (Hash-dumping)',
    'SamIGetPrivateData':'The code may query the private information about a specific user from the Security Account Manager (SAM) database. (Hash-dumping)',
    'SamQueryInformationUse':'The code may query information about a specific user in the Security Account Manager (SAM) database. (Hash-dumping)',
    'send':'The code may send data to a remote command-and-control server.',
    'SetFileTime':'The code may modify the creation/access/last modified time of a file.',
    'SetThreadContext':'The code may modify the context of a given thread for code injection.',
    'SetWindowsHookEx':'The code may set a hook function to be called whenever a certain event is called. (Keylogger)',
    'SfcTerminateWatcherThread':'The code may disable Windows file protection and modify the otherwise protected files.',
    'ShellExecute':'The code may execute another program and create a new process.',
    'StartServiceCtrlDispatcher':'The code may run as a service.',
    'SuspendThread':'The code may suspend a thread to modify it by performing code injection.',
    'system':'The code may run another program provided by some C runtime libraries. (wrapper function to CreateProcess.',
    'Thread32First':'The code may iterate through the threads of a process to find an appropriate thread to inject into.',
    'Thread32Next':'The code may iterate through the threads of a process to find an appropriate thread to inject into.',
    'Toolhelp32ReadProcessMemory':'The code may read the memory of a remote process.',
    'URLDownloadToFile':'The code may download a file from a web server and save it to disk.',
    'VirtualAllocEx':'The code may allocate memory in a remote process as part of process injection.',
    'VirtualProtectEx':'The code may change the protection on a region of memory from read-only to executable.',
    'WideCharToMultiByte':'The code may convert a Unicode string into an ASCII string.',
    'WinExec':'The code may execute another program by creating a new process.',
    'WlxLoggedOnSAS':'The code may perform Graphical Identification and Authentication (GINA) replacement.',
    'Wow64DisableWow64FsRedirection':'The code may disable file redirection that occurs in 32-bit files loaded on a 64-bit system.',
    'WriteProcessMemory':'The code may write data to a remote process as part of process injection.',
    'WSAStartup':'The code may initialize low-level network functionality.'
     }

# These lists are for renaming functions if they contain certain API names. Based on the script by Alexander Hanel.

# resource manipulation
launcher = ['LCH','FindResource','LoadResource','SizeofResource']
# injection attacks
process_injection = ['PSJ','VirtualAllocEx','WriteProcessMemory']
DLL_injection = ['DLJ','CreateToolhelp32Snapshot','Process32First','Process32Next', 'OpenProcess', 'CreateRemoteThread', 'WriteProcessMemory']
direct_injection = ['DRJ', 'VirtualAllocEx','WriteProcessMemory','CreateRemoteThread']
process_replacement = ['PRP', 'ZwUnmapViewOfSection','VirtualAllocEx','WriteProcessMemory','SetThreadContext','ResumeThread']
hook_injection = ['HKJ', 'SetWindowsHookEx','CallNextHookEx','UnhookWindowsHookEx']
APC_injection = ['ACJ', 'CreateRemoteThread','WaitForSingleObjectEx','WaitForMultipleObjectsEx','Sleep']
APC_userspace = ['AUJ', 'QueueUserAPC','WaitForSingleObjectEx','CreateToolhelp32Snapshot','Process32First','Process32Next','Thread32Next',\
                 'Thread32First','NtQuerySystemInformation','ZwQuerySystemInformation']
APC_kernelspace = ['AKJ', 'KeInitializeApc','KeInsertQueueApc']
# networking code
windows_networking = ['WNT','WSAStartup','getaddrinfo','socket','connect','send','recv','WSAGetLastError','InternetOpen','InternetConnect',\
                      'InternetOpenURL','InternetReadFile','InternetWriteFile','HTTPOpenRequest','HTTPQueryInfo','HTTPSendRequest',\
                      'URLDownloadToFile','CoInitialize','CoCreateInstance','Navigate','gethostbyname']
anti_debugging = ['ADB','IsDebuggerPresent','CheckRemoteDebuggerPresent','NtQueryInfromationProcess','OutputDebugString','FindWindow',\
                  'QueryPerformanceCounter','GetTickCount']
# anti-vm x86 instructions -> e.g. GetMnem(i) == "cpuid"
anti_vmx86 = ['AVM','sidt','sgdt','sldt','smsw','str','in','cpuid']

reg = [ 'REG', 'RegCloseKey' ,'RegConnectRegistryA' ,'RegConnectRegistryW' ,'RegCreateKeyA' ,'RegCreateKeyExA' ,'RegCreateKeyExW',\
        'RegCreateKeyW' ,'RegDeleteKeyA' ,'RegDeleteKeyW' ,'RegDeleteValueA' ,'RegDeleteValueW' ,'RegDisablePredefinedCache' ,\
        'RegDisablePredefinedCacheEx' ,'RegEnumKeyA' ,'RegEnumKeyExA' ,'RegEnumKeyExW' ,'RegEnumKeyW' ,'RegEnumValueA' ,\
        'RegEnumValueW' ,'RegFlushKey' ,'RegGetKeySecurity' ,'RegLoadKeyA' ,'RegLoadKeyW' ,'RegNotifyChangeKeyValue' ,\
        'RegOpenCurrentUser' ,'RegOpenKeyA' ,'RegOpenKeyExA' ,'RegOpenKeyExW' ,'RegOpenKeyW' ,'RegOpenUserClassesRoot' ,\
        'RegOverridePredefKey' ,'RegQueryInfoKeyA' ,'RegQueryInfoKeyW' ,'RegQueryMultipleValuesA' ,'RegQueryMultipleValuesW' ,\
        'RegQueryValueA' ,'RegQueryValueExA' ,'RegQueryValueExW' ,'RegQueryValueW' ,'RegReplaceKeyA' ,'RegReplaceKeyW' ,\
        'RegRestoreKeyA' ,'RegRestoreKeyW' ,'RegSaveKeyA' ,'RegSaveKeyExA' ,'RegSaveKeyExW' ,'RegSaveKeyW' ,'RegSetKeySecurity' ,\
        'RegSetValueA' ,'RegSetValueExA' ,'RegSetValueExW' ,'RegSetValueW' ,'RegUnLoadKeyA' ,'RegUnLoadKeyW', 'SHDeleteEmptyKeyA' ,\
        'SHDeleteEmptyKeyW' ,'SHDeleteKeyA' ,'SHDeleteKeyW' ,'SHOpenRegStream2A' ,'SHOpenRegStream2W' ,'SHOpenRegStreamA' ,\
        'SHOpenRegStreamW' ,'SHQueryInfoKeyA' ,'SHQueryInfoKeyW' ,'SHQueryValueExA' ,'SHQueryValueExW' ,'SHRegCloseUSKey' ,\
        'SHRegCreateUSKeyA' ,'SHRegCreateUSKeyW' ,'SHRegDeleteEmptyUSKeyA' ,'SHRegDeleteEmptyUSKeyW' ,'SHRegDeleteUSValueA' ,\
        'SHRegDeleteUSValueW' ,'SHRegDuplicateHKey' ,'SHRegEnumUSKeyA' ,'SHRegEnumUSKeyW' ,'SHRegEnumUSValueA' ,'SHRegEnumUSValueW'\
        ,'SHRegGetBoolUSValueA' ,'SHRegGetBoolUSValueW' ,'SHRegGetPathA' ,'SHRegGetPathW' ,'SHRegGetUSValueA' ,'SHRegGetUSValueW' ,\
        'SHRegGetValueA' ,'SHRegGetValueW' ,'SHRegOpenUSKeyA' ,'SHRegOpenUSKeyW' ,'SHRegQueryInfoUSKeyA' ,'SHRegQueryInfoUSKeyW' ,\
        'SHRegQueryUSValueA' ,'SHRegQueryUSValueW' ,'SHRegSetPathA' ,'SHRegSetPathW' ,'SHRegSetUSValueA' ,'SHRegSetUSValueW' ,\
        'SHRegWriteUSValueA' ,'SHRegWriteUSValueW' ,'SHDeleteOrphanKeyA' ,'SHDeleteOrphanKeyW' ,'SHDeleteValueA' ,'SHDeleteValueW' ,\
        'SHEnumKeyExA' ,'SHEnumKeyExW' ,'SHEnumValueA' ,'SHEnumValueW' ,'SHGetValueA' ,'SHGetValueW' ,'SHOpenRegStream2A' ,\
        'SHOpenRegStream2W' ,'SHOpenRegStreamA' ,'SHOpenRegStreamW' ,'SHQueryInfoKeyA' ,'SHQueryInfoKeyW' ,'SHQueryValueExA' ,\
        'SHQueryValueExW' ,'SHRegCloseUSKey' ,'SHRegCreateUSKeyA' ,'SHRegCreateUSKeyW' ,'SHRegDeleteEmptyUSKeyA' ,\
        'SHRegDeleteEmptyUSKeyW' ,'SHRegDeleteUSValueA' ,'SHRegDeleteUSValueW' ,'SHRegDuplicateHKey' ,'SHRegEnumUSKeyA' ,\
        'SHRegEnumUSKeyW' ,'SHRegEnumUSValueA' ,'SHRegEnumUSValueW' ,'SHRegGetBoolUSValueA' ,'SHRegGetBoolUSValueW' ,'SHRegGetPathA' ,\
        'SHRegGetPathW' ,'SHRegGetUSValueA' ,'SHRegGetUSValueW' ,'SHRegGetValueA' ,'SHRegGetValueW' ,'SHRegOpenUSKeyA' ,'SHRegOpenUSKeyW' ,\
        'SHRegQueryInfoUSKeyA' ,'SHRegQueryInfoUSKeyW' ,'SHRegQueryUSValueA' ,'SHRegQueryUSValueW' ,'SHRegSetPathA' ,'SHRegSetPathW' ,\
        'SHRegSetUSValueA' ,'SHRegSetUSValueW' ,'SHRegWriteUSValueA' ,'SHRegWriteUSValueW']
winsock = [ 'NET', 'FreeAddrInfoW', 'GetAddrInfoW', 'GetNameInfoW', 'WEP', 'WPUCompleteOverlappedRequest', 'WSAAccept', \
            'WSAAddressToStringA', 'WSAAddressToStringW', 'WSAAsyncGetHostByAddr', 'WSAAsyncGetHostByName', 'WSAAsyncGetProtoByName',\
            'WSAAsyncGetProtoByNumber', 'WSAAsyncGetServByName', 'WSAAsyncGetServByPort', 'WSAAsyncSelect', 'WSACancelAsyncRequest',\
            'WSACancelBlockingCall', 'WSACleanup', 'WSACloseEvent', 'WSAConnect', 'WSACreateEvent', 'WSADuplicateSocketA',\
            'WSADuplicateSocketW', 'WSAEnumNameSpaceProvidersA', 'WSAEnumNameSpaceProvidersW', 'WSAEnumNetworkEvents', 'WSAEnumProtocolsA',\
            'WSAEnumProtocolsW', 'WSAEventSelect', 'WSAGetLastError', 'WSAGetOverlappedResult', 'WSAGetQOSByName', \
            'WSAGetServiceClassInfoA', 'WSAGetServiceClassInfoW', 'WSAGetServiceClassNameByClassIdA', 'WSAGetServiceClassNameByClassIdW',\
            'WSAHtonl', 'WSAHtons', 'WSAInstallServiceClassA', 'WSAInstallServiceClassW', 'WSAIoctl', 'WSAIsBlocking', 'WSAJoinLeaf', \
            'WSALookupServiceBeginA', 'WSALookupServiceBeginW', 'WSALookupServiceEnd', 'WSALookupServiceNextA', 'WSALookupServiceNextW', \
            'WSANSPIoctl', 'WSANtohl', 'WSANtohs', 'WSAProviderConfigChange', 'WSARecv', 'WSARecvDisconnect', 'WSARecvFrom', \
            'WSARemoveServiceClass', 'WSAResetEvent', 'WSASend', 'WSASendDisconnect', 'WSASendTo', 'WSASetBlockingHook', 'WSASetEvent',\
            'WSASetLastError', 'WSASetServiceA', 'WSASetServiceW', 'WSASocketA', 'WSASocketW', 'WSAStartup', 'WSAStringToAddressA', \
            'WSAStringToAddressW', 'WSAUnhookBlockingHook', 'WSAWaitForMultipleEvents', 'WSApSetPostRoutine', 'WSCDeinstallProvider', \
            'WSCEnableNSProvider', 'WSCEnumProtocols', 'WSCGetProviderPath', 'WSCInstallNameSpace', 'WSCInstallProvider', 'WSCUnInstallNameSpace',\
            'WSCUpdateProvider', 'WSCWriteNameSpaceOrder', 'WSCWriteProviderOrder', '__WSAFDIsSet', 'accept', 'bind', 'closesocket', 'connect', \
            'freeaddrinfo', 'getaddrinfo', 'gethostbyaddr', 'gethostbyname', 'gethostname', 'getnameinfo', 'getpeername', 'getprotobyname', \
            'getprotobynumber', 'getservbyname', 'getservbyport', 'getsockname', 'getsockopt', 'htonl', 'htons', 'inet_addr', 'inet_ntoa', \
            'ioctlsocket', 'listen', 'ntohl', 'ntohs', 'recv', 'recvfrom', 'select', 'send', 'sendto', 'setsockopt', 'shutdown', 'socket']
WinINet = [ 'NET', 'CreateMD5SSOHash', 'DetectAutoProxyUrl', 'DllInstall', 'ForceNexusLookup', 'ForceNexusLookupExW', 'InternetAlgIdToStringA',\
            'InternetAlgIdToStringW', 'InternetAttemptConnect', 'InternetAutodial', 'InternetAutodialCallback', 'InternetAutodialHangup',\
            'InternetCanonicalizeUrlA', 'InternetCanonicalizeUrlW', 'InternetCheckConnectionA', 'InternetCheckConnectionW', \
            'InternetClearAllPerSiteCookieDecisions', 'InternetCloseHandle', 'InternetCombineUrlA', 'InternetCombineUrlW', \
            'InternetConfirmZoneCrossing', 'InternetConfirmZoneCrossingA', 'InternetConfirmZoneCrossingW', 'InternetConnectA',\
            'InternetConnectW', 'InternetCrackUrlA', 'InternetCrackUrlW', 'InternetCreateUrlA', 'InternetCreateUrlW', 'InternetDial',\
            'InternetDialA', 'InternetDialW', 'InternetEnumPerSiteCookieDecisionA', 'InternetEnumPerSiteCookieDecisionW', 'InternetErrorDlg',\
            'InternetFindNextFileA', 'InternetFindNextFileW', 'InternetFortezzaCommand', 'InternetGetCertByURL', 'InternetGetCertByURLA',\
            'InternetGetConnectedState', 'InternetGetConnectedStateEx', 'InternetGetConnectedStateExA', 'InternetGetConnectedStateExW',\
            'InternetGetCookieA', 'InternetGetCookieExA', 'InternetGetCookieExW', 'InternetGetCookieW', 'InternetGetLastResponseInfoA', \
            'InternetGetLastResponseInfoW', 'InternetGetPerSiteCookieDecisionA', 'InternetGetPerSiteCookieDecisionW', 'InternetGoOnline',\
            'InternetGoOnlineA', 'InternetGoOnlineW', 'InternetHangUp', 'InternetInitializeAutoProxyDll', 'InternetLockRequestFile',\
            'InternetOpenA', 'InternetOpenUrlA', 'InternetOpenUrlW', 'InternetOpenW', 'InternetQueryDataAvailable', 'InternetQueryFortezzaStatus',\
            'InternetQueryOptionA', 'InternetQueryOptionW', 'InternetReadFile', 'InternetReadFileExA', 'InternetReadFileExW', \
            'InternetSecurityProtocolToStringA', 'InternetSecurityProtocolToStringW', 'InternetSetCookieA', 'InternetSetCookieExA', \
            'InternetSetCookieExW', 'InternetSetCookieW', 'InternetSetDialState', 'InternetSetDialStateA', 'InternetSetDialStateW',\
            'InternetSetFilePointer', 'InternetSetOptionA', 'InternetSetOptionExA', 'InternetSetOptionExW', 'InternetSetOptionW', \
            'InternetSetPerSiteCookieDecisionA', 'InternetSetPerSiteCookieDecisionW', 'InternetSetStatusCallback', 'InternetSetStatusCallbackA',\
            'InternetSetStatusCallbackW', 'InternetShowSecurityInfoByURL', 'InternetShowSecurityInfoByURLA', 'InternetShowSecurityInfoByURLW', \
            'InternetTimeFromSystemTime', 'InternetTimeFromSystemTimeA', 'InternetTimeFromSystemTimeW', 'InternetTimeToSystemTime',\
            'InternetTimeToSystemTimeA', 'InternetTimeToSystemTimeW', 'InternetUnlockRequestFile', 'InternetWriteFile', 'InternetWriteFileExA',\
            'InternetWriteFileExW', 'IsHostInProxyBypassList', 'ParseX509EncodedCertificateForListBoxEntry', 'PrivacyGetZonePreferenceW', \
            'PrivacySetZonePreferenceW', 'ResumeSuspendedDownload', 'ShowCertificate', 'ShowClientAuthCerts', 'ShowSecurityInfo', \
            'ShowX509EncodedCertificate','UrlZonesDetach', '_GetFileExtensionFromUrl'] 
cache = [ 'NET','CommitUrlCacheEntryA', 'CommitUrlCacheEntryW', 'CreateUrlCacheContainerA', 'CreateUrlCacheContainerW', 'CreateUrlCacheEntryA',\
          'CreateUrlCacheEntryW', 'CreateUrlCacheGroup', 'DeleteIE3Cache', 'DeleteUrlCacheContainerA', 'DeleteUrlCacheContainerW', \
          'DeleteUrlCacheEntry', 'DeleteUrlCacheEntryA', 'DeleteUrlCacheEntryW', 'DeleteUrlCacheGroup', 'FindCloseUrlCache', 'FindFirstUrlCacheContainerA',\
          'FindFirstUrlCacheContainerW', 'FindFirstUrlCacheEntryA', 'FindFirstUrlCacheEntryExA', 'FindFirstUrlCacheEntryExW', 'FindFirstUrlCacheEntryW', \
          'FindFirstUrlCacheGroup', 'FindNextUrlCacheContainerA', 'FindNextUrlCacheContainerW', 'FindNextUrlCacheEntryA', 'FindNextUrlCacheEntryExA',\
          'FindNextUrlCacheEntryExW', 'FindNextUrlCacheEntryW', 'FindNextUrlCacheGroup', 'FreeUrlCacheSpaceA', 'FreeUrlCacheSpaceW', 'GetUrlCacheConfigInfoA', \
          'GetUrlCacheConfigInfoW', 'GetUrlCacheEntryInfoA', 'GetUrlCacheEntryInfoExA', 'GetUrlCacheEntryInfoExW', 'GetUrlCacheEntryInfoW', \
          'GetUrlCacheGroupAttributeA', 'GetUrlCacheGroupAttributeW', 'GetUrlCacheHeaderData', 'IncrementUrlCacheHeaderData', 'IsUrlCacheEntryExpiredA',\
          'IsUrlCacheEntryExpiredW', 'LoadUrlCacheContent', 'ReadUrlCacheEntryStream', 'RegisterUrlCacheNotification', 'RetrieveUrlCacheEntryFileA', \
          'RetrieveUrlCacheEntryFileW', 'RetrieveUrlCacheEntryStreamA', 'RetrieveUrlCacheEntryStreamW', 'RunOnceUrlCache', 'SetUrlCacheConfigInfoA',\
          'SetUrlCacheConfigInfoW', 'SetUrlCacheEntryGroup', 'SetUrlCacheEntryGroupA', 'SetUrlCacheEntryGroupW', 'SetUrlCacheEntryInfoA', 'SetUrlCacheEntryInfoW',\
          'SetUrlCacheGroupAttributeA', 'SetUrlCacheGroupAttributeW', 'SetUrlCacheHeaderData', 'UnlockUrlCacheEntryFile', 'UnlockUrlCacheEntryFileA', \
          'UnlockUrlCacheEntryFileW', 'UnlockUrlCacheEntryStream', 'UpdateUrlCacheContentPath']
ftp = [ 'FTP','FtpCommandA' ,'FtpCommandW' ,'FtpCreateDirectoryA' ,'FtpCreateDirectoryW' ,'FtpDeleteFileA' ,'FtpDeleteFileW' ,'FtpFindFirstFileA' ,\
        'FtpFindFirstFileW' ,'FtpGetCurrentDirectoryA' ,'FtpGetCurrentDirectoryW' ,'FtpGetFileA' ,'FtpGetFileEx' ,'FtpGetFileSize' ,'FtpGetFileW' ,\
        'FtpOpenFileA' ,'FtpOpenFileW' ,'FtpPutFileA' ,'FtpPutFileEx' ,'FtpPutFileW' ,'FtpRemoveDirectoryA' ,'FtpRemoveDirectoryW' ,'FtpRenameFileA' ,\
        'FtpRenameFileW' ,'FtpSetCurrentDirectoryA' ,'FtpSetCurrentDirectoryW']
gopher = [ 'GPR', 'GopherCreateLocatorA', 'GopherCreateLocatorW', 'GopherFindFirstFileA', 'GopherFindFirstFileW', 'GopherGetAttributeA', \
           'GopherGetAttributeW', 'GopherGetLocatorTypeA', 'GopherGetLocatorTypeW', 'GopherOpenFileA', 'GopherOpenFileW'] 
url = ['URL', 'UrlApplySchemeA' ,'UrlApplySchemeW' ,'UrlCanonicalizeA' ,'UrlCanonicalizeW' ,'UrlCombineA' ,'UrlCombineW' ,'UrlCompareA' ,\
       'UrlCompareW' ,'UrlCreateFromPathA' ,'UrlCreateFromPathW' ,'UrlEscapeA' ,'UrlEscapeW' ,'UrlGetLocationA' ,'UrlGetLocationW' ,'UrlGetPartA'\
       ,'UrlGetPartW' ,'UrlHashA' ,'UrlHashW' ,'UrlIsA' ,'UrlIsNoHistoryA' ,'UrlIsNoHistoryW' ,'UrlIsOpaqueA' ,'UrlIsOpaqueW' ,'UrlIsW' ,'UrlUnescapeA' ,'UrlUnescapeW']
dir = ['DIR','CreateDirectoryA', 'CreateDirectoryExA', 'CreateDirectoryExW', 'CreateDirectoryW', 'GetCurrentDirectoryA', 'GetCurrentDirectoryW',\
       'GetDllDirectoryA', 'GetDllDirectoryW', 'GetSystemDirectoryA', 'GetSystemDirectoryW', 'GetSystemWindowsDirectoryA', 'GetSystemWindowsDirectoryW',\
       'GetSystemWow64DirectoryA', 'GetSystemWow64DirectoryW', 'GetVDMCurrentDirectories', 'GetWindowsDirectoryA', 'GetWindowsDirectoryW', \
       'ReadDirectoryChangesW', 'RemoveDirectoryA', 'RemoveDirectoryW', 'SetCurrentDirectoryA', 'SetCurrentDirectoryW', 'SetDllDirectoryA',\
       'SetDllDirectoryW', 'SetVDMCurrentDirectories', 'SHCreateDirectory', 'SHCreateDirectoryExA', 'SHCreateDirectoryExW']
mutex = ['MTX','CreateMutexA', 'CreateMutexW', 'OpenMutexA', 'OpenMutexW', 'ReleaseMutex']
pipe = [ 'PIP', 'CallNamedPipeA', 'CallNamedPipeW', 'ConnectNamedPipe', 'CreateNamedPipeA', 'CreateNamedPipeW', 'CreatePipe', 'DisconnectNamedPipe',\
         'GetNamedPipeHandleStateA', 'GetNamedPipeHandleStateW', 'GetNamedPipeInfo', 'PeekNamedPipe', 'SetNamedPipeHandleState', 'TransactNamedPipe',\
         'WaitNamedPipeA', 'WaitNamedPipeW']
http = [ 'HTP', 'HttpAddRequestHeadersA', 'HttpAddRequestHeadersW', 'HttpCheckDavCompliance', 'HttpEndRequestA', 'HttpEndRequestW',\
         'HttpOpenRequestA', 'HttpOpenRequestW', 'HttpQueryInfoA', 'HttpQueryInfoW', 'HttpSendRequestA', 'HttpSendRequestExA', \
         'HttpSendRequestExW', 'HttpSendRequestW' ] 
enum = [ 'ENP', 'CreateToolhelp32Snapshot', 'Process32First' ,'Process32FirstW' ,'Process32Next' ,'Process32NextW']
hash = ['HAS', 'CryptCreateHash' ,'CryptDestroyHash' ,'CryptDuplicateHash' ,'CryptGetHashParam' ,'CryptHashData' ,'CryptHashSessionKey' ,\
        'CryptSetHashParam' ,'CryptSignHashA' ,'CryptSignHashW', 'FreeEncryptionCertificateHashList']
crypt = ['CRY', 'CryptAcquireContextA' ,'CryptAcquireContextW' ,'CryptContextAddRef' ,'CryptDecrypt' ,'CryptDeriveKey' ,'CryptDestroyKey' ,\
         'CryptDuplicateKey' ,'CryptEncrypt' ,'CryptEnumProviderTypesA' ,'CryptEnumProviderTypesW' ,'CryptEnumProvidersA' ,'CryptEnumProvidersW'\
         ,'CryptExportKey' ,'CryptGenKey' ,'CryptGenRandom' ,'CryptGetDefaultProviderA' ,'CryptGetDefaultProviderW' ,'CryptGetKeyParam' ,\
         'CryptGetProvParam' ,'CryptGetUserKey' ,'CryptImportKey' ,'CryptReleaseContext' ,'CryptSetKeyParam' ,'CryptSetProvParam' ,\
         'CryptSetProviderA' ,'CryptSetProviderExA' ,'CryptSetProviderExW' ,'CryptSetProviderW' ,'CryptVerifySignatureA' ,'CryptVerifySignatureW' ,\
         'DecryptFileA' ,'DecryptFileW', 'EncryptFileA' ,'EncryptFileW' ,'EncryptedFileKeyInfo' ,'EncryptionDisable', 'WriteEncryptedFileRaw', \
         'OpenEncryptedFileRawA' ,'OpenEncryptedFileRawW', 'DuplicateEncryptionInfoFile', 'SetUserFileEncryptionKey', 'ReadEncryptedFileRaw', \
         'RemoveUsersFromEncryptedFile', 'FileEncryptionStatusA', 'FileEncryptionStatusW', 'FreeEncryptedFileKeyInfo', 'CloseEncryptedFileRaw', \
         'AddUsersToEncryptedFile', 'QueryRecoveryAgentsOnEncryptedFile', 'QueryUsersOnEncryptedFile', 'ChainWlxLogoffEvent' ,'CryptAcquireContextU' ,\
         'CryptBinaryToStringA' ,'CryptBinaryToStringW' ,'CryptCloseAsyncHandle' ,'CryptCreateAsyncHandle' ,'CryptDecodeMessage' ,'CryptDecodeObject' ,\
         'CryptDecodeObjectEx' ,'CryptDecryptAndVerifyMessageSignature' ,'CryptDecryptMessage' ,'CryptEncodeObject' ,'CryptEncodeObjectEx' ,\
         'CryptEncryptMessage' ,'CryptEnumKeyIdentifierProperties' ,'CryptEnumOIDFunction' ,'CryptEnumOIDInfo' ,'CryptEnumProvidersU' ,'CryptExportPKCS8' ,\
         'CryptExportPublicKeyInfo' ,'CryptExportPublicKeyInfoEx' ,'CryptFindLocalizedName' ,'CryptFindOIDInfo' ,'CryptFormatObject' ,\
         'CryptFreeOIDFunctionAddress' ,'CryptGetAsyncParam' ,'CryptGetDefaultOIDDllList' ,'CryptGetDefaultOIDFunctionAddress' ,\
         'CryptGetKeyIdentifierProperty' ,'CryptGetMessageCertificates' ,'CryptGetMessageSignerCount' ,'CryptGetOIDFunctionAddress' ,\
         'CryptGetOIDFunctionValue' ,'CryptHashCertificate' ,'CryptHashMessage' ,'CryptHashPublicKeyInfo' ,'CryptHashToBeSigned' ,\
         'CryptImportPKCS8' ,'CryptImportPublicKeyInfo' ,'CryptImportPublicKeyInfoEx' ,'CryptInitOIDFunctionSet' ,'CryptInstallDefaultContext' ,\
         'CryptInstallOIDFunctionAddress' ,'CryptLoadSip' ,'CryptMemAlloc' ,'CryptMemFree' ,'CryptMemRealloc' ,'CryptMsgCalculateEncodedLength' ,\
         'CryptMsgClose' ,'CryptMsgControl' ,'CryptMsgCountersign' ,'CryptMsgCountersignEncoded' ,'CryptMsgDuplicate' ,'CryptMsgEncodeAndSignCTL' ,\
         'CryptMsgGetAndVerifySigner' ,'CryptMsgGetParam' ,'CryptMsgOpenToDecode' ,'CryptMsgOpenToEncode' ,'CryptMsgSignCTL' ,'CryptMsgUpdate' ,\
         'CryptMsgVerifyCountersignatureEncoded' ,'CryptMsgVerifyCountersignatureEncodedEx' ,'CryptProtectData' ,'CryptQueryObject' ,\
         'CryptRegisterDefaultOIDFunction' ,'CryptRegisterOIDFunction' ,'CryptRegisterOIDInfo' ,'CryptSIPAddProvider' ,\
         'CryptSIPCreateIndirectData' ,'CryptSIPGetSignedDataMsg' ,'CryptSIPLoad' ,'CryptSIPPutSignedDataMsg' ,'CryptSIPRemoveProvider' ,\
         'CryptSIPRemoveSignedDataMsg' ,'CryptSIPRetrieveSubjectGuid' ,'CryptSIPRetrieveSubjectGuidForCatalogFile' ,'CryptSIPVerifyIndirectData' ,\
         'CryptSetAsyncParam' ,'CryptSetKeyIdentifierProperty' ,'CryptSetOIDFunctionValue' ,'CryptSetProviderU' ,'CryptSignAndEncodeCertificate' ,\
         'CryptSignAndEncryptMessage' ,'CryptSignCertificate' ,'CryptSignHashU' ,'CryptSignMessage' ,'CryptSignMessageWithKey' ,\
         'CryptStringToBinaryA' ,'CryptStringToBinaryW' ,'CryptUninstallDefaultContext' ,'CryptUnprotectData' ,'CryptUnregisterDefaultOIDFunction' ,\
         'CryptUnregisterOIDFunction' ,'CryptUnregisterOIDInfo' ,'CryptVerifyCertificateSignature' ,'CryptVerifyCertificateSignatureEx' ,\
         'CryptVerifyDetachedMessageHash' ,'CryptVerifyDetachedMessageSignature' ,'CryptVerifyMessageHash' ,'CryptVerifyMessageSignature' ,\
         'CryptVerifyMessageSignatureWithKey' ,'CryptVerifySignatureU' ,'I_CertProtectFunction' ,'I_CertSrvProtectFunction' ,'I_CertSyncStore' ,\
         'I_CertUpdateStore' ,'I_CryptAddRefLruEntry' ,'I_CryptAddSmartCardCertToStore' ,'I_CryptAllocTls' ,'I_CryptCreateLruCache' ,\
         'I_CryptCreateLruEntry' ,'I_CryptDetachTls' ,'I_CryptDisableLruOfEntries' ,'I_CryptEnableLruOfEntries' ,'I_CryptEnumMatchingLruEntries' ,\
         'I_CryptFindLruEntry' ,'I_CryptFindLruEntryData' ,'I_CryptFindSmartCardCertInStore' ,'I_CryptFlushLruCache' ,'I_CryptFreeLruCache' ,\
         'I_CryptFreeTls' ,'I_CryptGetAsn1Decoder' ,'I_CryptGetAsn1Encoder' ,'I_CryptGetDefaultCryptProv' ,'I_CryptGetDefaultCryptProvForEncrypt' ,\
         'I_CryptGetFileVersion' ,'I_CryptGetLruEntryData' ,'I_CryptGetLruEntryIdentifier' ,'I_CryptGetOssGlobal' ,'I_CryptGetTls' ,'I_CryptInsertLruEntry' ,\
         'I_CryptInstallAsn1Module' ,'I_CryptInstallOssGlobal' ,'I_CryptReadTrustedPublisherDWORDValueFromRegistry' ,'I_CryptRegisterSmartCardStore' ,\
         'I_CryptReleaseLruEntry' ,'I_CryptRemoveLruEntry' ,'I_CryptSetTls' ,'I_CryptTouchLruEntry' ,'I_CryptUninstallAsn1Module' ,\
         'I_CryptUninstallOssGlobal' ,'I_CryptUnregisterSmartCardStore' ,'I_CryptWalkAllLruCacheEntries']
service = ['SRV', 'ChangeServiceConfig2A' ,'ChangeServiceConfig2W' ,'ChangeServiceConfigA' ,'ChangeServiceConfigW' ,'CloseServiceHandle' ,\
           'ControlService' ,'CreateServiceA' ,'CreateServiceW' ,'DeleteService' ,'EnumDependentServicesA' ,'EnumDependentServicesW' ,\
           'EnumServiceGroupW' ,'EnumServicesStatusA' ,'EnumServicesStatusExA' ,'EnumServicesStatusExW' ,'EnumServicesStatusW' ,\
           'GetServiceDisplayNameA' ,'GetServiceDisplayNameW' ,'GetServiceKeyNameA' ,'GetServiceKeyNameW' ,'I_ScPnPGetServiceName' ,\
           'I_ScSetServiceBitsA' ,'I_ScSetServiceBitsW' ,'LockServiceDatabase' ,'OpenServiceA' ,'OpenServiceW' ,'PrivilegedServiceAuditAlarmA' ,\
           'PrivilegedServiceAuditAlarmW' ,'QueryServiceConfig2A' ,'QueryServiceConfig2W' ,'QueryServiceConfigA' ,'QueryServiceConfigW' ,\
           'QueryServiceLockStatusA' ,'QueryServiceLockStatusW' ,'QueryServiceObjectSecurity' ,'QueryServiceStatus' ,'QueryServiceStatusEx' ,\
           'RegisterServiceCtrlHandlerA' ,'RegisterServiceCtrlHandlerExA' ,'RegisterServiceCtrlHandlerExW' ,'RegisterServiceCtrlHandlerW' ,\
           'SetServiceBits' ,'SetServiceObjectSecurity' ,'SetServiceStatus' ,'StartServiceA' ,'StartServiceCtrlDispatcherA' ,'StartServiceCtrlDispatcherW' ,\
           'StartServiceW' ,'UnlockServiceDatabase' ,'WdmWmiServiceMain']
file = ['FIL', 'CompareFileTime' ,'CopyFileA' ,'CopyFileExA' ,'CopyFileExW' ,'CopyFileW' ,'CopyLZFile' ,'CreateFileA' ,'CreateFileMappingA' ,\
        'CreateFileMappingW' ,'CreateFileW' ,'DeleteFileA' ,'DeleteFileW' ,'DosDateTimeToFileTime' ,'FileTimeToDosDateTime' ,\
        'FileTimeToLocalFileTime' ,'FileTimeToLocalFileTime' ,'FileTimeToSystemTime' ,'FlushFileBuffers' ,'FlushViewOfFile' ,\
        'GetCPFileNameFromRegistry' ,'GetCompressedFileSizeA' ,'GetCompressedFileSizeW' ,'GetFileAttributesA' ,'GetFileAttributesExA' ,\
        'GetFileAttributesExW' ,'GetFileAttributesW' ,'GetFileInformationByHandle' ,'GetFileSize' ,'GetFileSizeEx' ,'GetFileTime' ,\
        'GetFileType' ,'GetSystemTimeAsFileTime' ,'GetTempFileNameA' ,'GetTempFileNameW' ,'LZCloseFile' ,'LZCreateFileW' ,'LZOpenFileA',\
        'LZOpenFileW' ,'LocalFileTimeToFileTime' ,'LocalFileTimeToFileTime' ,'LockFile' ,'LockFileEx' ,'MapViewOfFile' ,'MapViewOfFileEx' ,\
        'MoveFileA' ,'MoveFileExA' ,'MoveFileExW' ,'MoveFileW' ,'MoveFileWithProgressA' ,'MoveFileWithProgressW' ,'OpenDataFile' ,'OpenFile' ,\
        'OpenFileMappingA' ,'OpenFileMappingW' ,'OpenProfileUserMapping' ,'PrivCopyFileExW' ,'PrivMoveFileIdentityW' ,'ReadFile' ,'ReadFileEx' ,\
        'ReplaceFile' ,'ReplaceFileA' ,'ReplaceFileW' ,'SetEndOfFile' ,'SetFileAttributesA' ,'SetFileAttributesW' ,'SetFilePointer' ,\
        'SetFilePointerEx' ,'SetFileShortNameA' ,'SetFileShortNameW' ,'SetFileTime' ,'SetFileValidData' ,'SystemTimeToFileTime' ,\
        'UnlockFile' ,'UnlockFileEx' ,'UnmapViewOfFile' ,'WriteFile' ,'WriteFileEx' ,'WriteFileGather' ,'GetFileSecurityA' ,\
        'GetFileSecurityW' ,'SetFileSecurityA' ,'SetFileSecurityW', 'CreateFileU']
os_info = [ 'OSI', 'GetComputerNameA' ,'GetComputerNameExA' ,'GetComputerNameExW' ,'GetComputerNameW' ,'GetDiskFreeSpaceA' ,\
            'GetDiskFreeSpaceExA' ,'GetDiskFreeSpaceExW' ,'GetDiskFreeSpaceW' ,'GetDriveTypeA' ,'GetDriveTypeW', 'GetVersion' ,\
            'GetVersionExA' ,'GetVersionExW', 'GetSystemInfo', 'GetSystemMetrics', 'CheckTokenMembership']
cert = ['CER','CertAddCRLContextToStore' ,'CertAddCRLLinkToStore' ,'CertAddCTLContextToStore' ,'CertAddCTLLinkToStore' ,\
        'CertAddCertificateContextToStore' ,'CertAddCertificateLinkToStore' ,'CertAddEncodedCRLToStore' ,'CertAddEncodedCertificateToStore' ,\
        'CertAddEncodedCertificateToSystemStoreA' ,'CertAddEncodedCertificateToSystemStoreW' ,'CertAddEnhancedKeyUsageIdentifier' ,\
        'CertAddSerializedElementToStore' ,'CertAddStoreToCollection' ,'CertAlgIdToOID' ,'CertCloseStore' ,'CertCompareCertificate' ,\
        'CertCompareCertificateName' ,'CertCompareIntegerBlob' ,'CertComparePublicKeyInfo' ,'CertControlStore' ,'CertCreateCTLContext' ,\
        'CertCreateCTLEntryFromCertificateContextProperties' ,'CertCreateCertificateChainEngine' ,'CertCreateCertificateContext' ,'CertCreateContext',\
        'CertCreateSelfSignCertificate' ,'CertDeleteCTLFromStore' ,'CertDeleteCertificateFromStore' ,'CertDuplicateCTLContext' ,\
        'CertDuplicateCertificateChain' ,'CertDuplicateCertificateContext' ,'CertDuplicateStore' ,'CertEnumCRLContextProperties' ,\
        'CertEnumCRLsInStore' ,'CertEnumCTLContextProperties' ,'CertEnumCTLsInStore' ,'CertEnumCertificateContextProperties' ,\
        'CertEnumCertificatesInStore' ,'CertEnumPhysicalStore' ,'CertEnumSubjectInSortedCTL' ,'CertEnumSystemStore' ,\
        'CertEnumSystemStoreLocation' ,'CertFindAttribute' ,'CertFindCRLInStore' ,'CertFindCertificateInCRL' ,'CertFindCertificateInStore',\
        'CertFindChainInStore' ,'CertFindExtension' ,'CertFindRDNAttr' ,'CertFindSubjectInCTL' ,'CertFindSubjectInSortedCTL' ,\
        'CertFreeCRLContext' ,'CertFreeCertificateChain' ,'CertFreeCertificateChainEngine' ,'CertFreeCertificateContext' ,'CertGetCRLContextProperty' ,\
        'CertGetCRLFromStore' ,'CertGetCTLContextProperty' ,'CertGetCertificateChain' ,'CertGetCertificateContextProperty' ,'CertGetEnhancedKeyUsage' ,\
        'CertGetIssuerCertificateFromStore' ,'CertGetNameStringA' ,'CertGetNameStringW' ,'CertGetPublicKeyLength' ,'CertGetStoreProperty' ,\
        'CertGetSubjectCertificateFromStore' ,'CertGetValidUsages' ,'CertIsRDNAttrsInCertificateName' ,'CertIsValidCRLForCertificate' ,\
        'CertNameToStrA' ,'CertNameToStrW' ,'CertOIDToAlgId' ,'CertOpenStore' ,'CertOpenSystemStoreA' ,'CertOpenSystemStoreW' ,'CertRDNValueToStrA',\
        'CertRDNValueToStrW' ,'CertRegisterPhysicalStore' ,'CertRegisterSystemStore' ,'CertRemoveEnhancedKeyUsageIdentifier' ,\
        'CertRemoveStoreFromCollection' ,'CertResyncCertificateChainEngine' ,'CertSaveStore' ,'CertSerializeCRLStoreElement' ,\
        'CertSerializeCertificateStoreElement' ,'CertSetCRLContextProperty' ,'CertSetCertificateContextPropertiesFromCTLEntry' ,\
        'CertSetCertificateContextProperty' ,'CertSetEnhancedKeyUsage' ,'CertSetStoreProperty' ,'CertStrToNameA' ,'CertStrToNameW' ,\
        'CertUnregisterPhysicalStore' ,'CertUnregisterSystemStore' ,'CertVerifyCRLRevocation' ,'CertVerifyCRLTimeValidity' ,\
        'CertVerifyCTLUsage' ,'CertVerifyCertificateChainPolicy' ,'CertVerifyCertificateChainPolicy' ,'CertVerifyRevocation' ,\
        'CertVerifySubjectCertificateContext','CertVerifyTimeValidity' ,'CertVerifyValidityNesting' ,'CloseCertPerformanceData' ,\
        'CollectCertPerformanceData' ,'CryptAcquireCertificatePrivateKey' ,'CryptFindCertificateKeyProvInfo' ,'CryptGetMessageCertificates' ,\
        'CryptHashCertificate' ,'CryptSignAndEncodeCertificate' ,'CryptSignCertificate' ,'CryptVerifyCertificateSignature' ,\
        'CryptVerifyCertificateSignatureEx' ,'I_CertProtectFunction' ,'I_CertSrvProtectFunction' ,'I_CertSyncStore' ,'I_CertUpdateStore' ,\
        'I_CryptAddSmartCardCertToStore' ,'I_CryptFindSmartCardCertInStore' ,'OpenCertPerformanceData' ,'PFXExportCertStore' ,'PFXExportCertStoreEx' ,'PFXImportCertStore'] 
file_s = ['SRC', 'FindFirstFileW', 'FindNextFileW', 'FindClose']
modify = ['MOD', 'WriteProcessMemory', 'ReadProcessMemory'] 
virtual = ['VIR', 'VirtualAlloc' ,'VirtualAllocEx' ,'VirtualBufferExceptionHandler' ,'VirtualFree' ,'VirtualFreeEx' ,'VirtualLock' ,
           'VirtualProtect' ,'VirtualProtectEx' ,'VirtualQuery' ,'VirtualQueryEx' ,'VirtualUnlock']
critical_section = [ 'CRT', 'DeleteCriticalSection' ,'EnterCriticalSection' ,'InitializeCriticalSection' ,
                     'InitializeCriticalSectionAndSpinCount' ,'LeaveCriticalSection' ,'SetCriticalSectionSpinCount' ,'TryEnterCriticalSection']

api_matrix = [ reg, winsock, WinINet, cache, ftp, gopher, dir, mutex, pipe, http, enum, hash, crypt, service, file, cert, os_info, file_s, modify, virtual, critical_section,\
               launcher, process_injection, DLL_injection, direct_injection, process_replacement, hook_injection, APC_injection, APC_injection, APC_userspace, APC_kernelspace,\
               windows_networking, anti_debugging]

def getAllStrings():
    #We want to use a specific string configuration
    stringList = idautils.Strings(default_setup = False)
    # Lets use C et UNICODE string
    stringList.setup(strtypes=Strings.STR_C | Strings.STR_UNICODE, ignore_instructions = True, display_only_existing_strings = True)

    parsedStringList = []
    #Get all strings
    for string in stringList:                
        if string:
            parsedStringList.append((str(string), string.ea))

    return parsedStringList
#-----------------------------------------------------------------------
# Configuration agent
# The two next function are needed so the configuration manager will be
# able to provide configuration for this module.
# configurationNeed is called first. If no config for this 
# utility exists, user is prompted for utility configuration. If
# configuration exists, previous configs will be used. In all cases,
# configuration step is ended when configurationProvision
# is called with a configuration list as arg.
#-----------------------------------------------------------------------
'''
IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH = ""
def configurationNeed():
    return ["IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH", "IDA_PRO_TEXT_EXTRACTOR_EXTRACTED_FILE"]
            
def configurationProvision(utilityConfig=[]):
    global IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH
    global IDA_PRO_TEXT_EXTRACTOR_EXTRACTED_FILE
    IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH = utilityConfig[0]
    IDA_PRO_TEXT_EXTRACTOR_EXTRACTED_FILE = utilityConfig[1]
'''

def sanitize_name(name):
    name = name.strip()
    if name.endswith("A"):
        name = name[:-1]
    elif name.endswith("W"):
        name = name[:-1]
    while name.startswith("_"):
        name = name[1:]
    return name

def getFuncDeclaration(func):
    if func != idc.BADADDR:
        ft = idc.GetType(func)
        if ft:
            return ft
        else:
            #print "-Type info not available!"
            return ""
    else:
        return ""    
  
def prefix_lines(lines):
    if not lines or len(lines) == 0:
        return lines
    result = [RESULT_PREFIX+line for line in lines.splitlines()]
    return "\n".join(result)
    
def prefix_lines_oa(lines):
    if not lines or len(lines) == 0:
        return lines
    result = [RESULT_PREFIX_OA+line for line in lines.splitlines()]
    return "\n".join(result)
    
def remove_prefixed_lines(lines):
    if not lines or len(lines) == 0:
        return lines
    result = [line for line in lines.splitlines() \
              if not line.startswith(RESULT_PREFIX)]
    return "\n".join(result)
    
def remove_prefixed_lines_oa(lines):
    if not lines or len(lines) == 0:
        return lines
    result = [line for line in lines.splitlines() \
              if not line.startswith(RESULT_PREFIX_OA)]
    return "\n".join(result)

  
def get_all_function_eas():
    result = []
    for fnum in xrange(idaapi.get_func_qty()):
        func = idaapi.getn_func(fnum)
        f_ea = func.startEA
        if f_ea != idaapi.BADADDR:
            result.append(f_ea)
    return result
    
def get_import_segments():
    result = []
    for seg_ea in idautils.Segments():
        seg = idaapi.getseg(seg_ea)
        if seg.type == idaapi.SEG_XTRN:
            result.append( (seg.startEA, seg.endEA) )
    return result
    
def u2signed(num):
    if num < 0x80000000:
        return num
    else:
        return num - 0x100000000
    
class FunctionX(object):
    def __init__(self, start_ea):
        self.start_ea = start_ea
        
    def get_end_address(self):
        return "0x%x" % idc.GetFunctionAttr(self.start_ea, idc.FUNCATTR_END)
    def get_number_of_instructions(self):
        icount = 0
        for i in idautils.FuncItems(self.start_ea): icount = icount + 1
        return icount
    def get_size_of_argument_bytes(self):
        try:
            fframe = idc.GetFrame(self.start_ea)
            if fframe is not None:
                fret = idc.GetMemberOffset(fframe, " r")
            if fret is not -1:
                firstArg = fret + 4
            args = idc.GetStrucSize(fframe) - firstArg
            return args
        except:
            pass     
            return
    def get_number_of_arguments(self):
        try:
            args = self.get_size_of_argument_bytes()
            out = args/4
        except:
            out = 0
            pass
        return out
    def get_size_of_local_variables(self):
        try:
            locals = idc.GetFunctionAttr(self.start_ea, idc.FUNCATTR_FRSIZE)
            return locals
        except:
            pass
            return
    def get_function_flags(self):
        try:
            return idc.GetFunctionFlags(self.start_ea)
        except:
            pass
            return
    def get_name(self):
        return idaapi.get_func_name(self.start_ea)
    def get_demangled_name(self):
        mangled = idaapi.get_func_name(self.start_ea)
        demangled = idc.Demangle(mangled, idc.GetLongPrm(idc.INF_SHORT_DN))
        if  demangled == None:
             return "" #mangled
        else:
            #print " --Demangling..."
            return demangled
    def get_instruction_eas(self):
        result = []
        for chunk in self.__get_chunks():
            ins_ea = idc.FindCode(chunk.startEA-1, DIRECTION_FORW)
            while ins_ea <= chunk.endEA and ins_ea != idc.BADADDR:
                result.append(ins_ea)
                ins_ea = idc.FindCode(ins_ea, DIRECTION_FORW)
        return result    
    def get_instructions(self):
        return [Instruction(iea, ) for iea in self.get_instruction_eas()]
    def get_code_refs_from(self):
        result = []
        for iea in self.get_instruction_eas():
            inst = Instruction(iea, )
            result += inst.get_code_refs_from()
        return list(set(result))
    def get_data_refs_from(self):
        result = []
        for iea in self.get_instruction_eas():
            inst = Instruction(iea, )
            result += inst.get_data_refs_from()
        return list(set(result))
    def set_comment(self, comment, repeatable = False):
        func = self.__get_func_t()
        idaapi.set_func_cmt(func, comment, repeatable)
    def get_comment(self, repeatable = False):
        func = self.__get_func_t()
        return idaapi.get_func_cmt(func, repeatable)
    def __get_chunks(self):
        func = self.__get_func_t()
        result = [func]
        ft_iter = idaapi.func_tail_iterator_t(func, self.start_ea)
        if ft_iter.first():
            result.append(ft_iter.chunk())
        while ft_iter.next():
            result.append(ft_iter.chunk())
        return result
    def __get_func_t(self):
        result = idaapi.get_func(self.start_ea)
        if not result:
            raise RuntimeError, \
                  "Cannot retrieve function information @ address %s" % \
                  self.start_ea
        return result
    def __repr__(self):
        return self.get_name()

class Instruction(object):
    def __init__(self, inst_ea):
        self.iea = inst_ea
        
    def get_code_refs_from(self):
        return idautils.CodeRefsFrom(self.iea, False)
    def get_data_refs_from(self):
        return idautils.DataRefsFrom(self.iea, False)
    def get_operands(self):
        result = []
        ins = self.__get_insn_t()
        if not ins:
            print >> sys.stderr, \
                  "Cannot retrieve operand information @ address %s" % \
                  hex(self.iea)
            return result
        for i in xrange(6):
            oper = ins[i]
            if oper.type != idaapi.o_void:
                result.append( Operand(i, self.iea, ) )
        return result
    def __get_insn_t(self):
        ins = idautils.DecodeInstruction(self.iea)
        if ins:
	    return ins
        return None
    def __repr__(self):        
        return idaapi.tag_remove( idaapi.generate_disasm_line(self.iea) )
        
class Operand(object):
    def __init__(self, oper_num, inst_ea):
        self.onum = oper_num
        self.iea = inst_ea
        
        
    def get_immediate(self):
        oper = self.__get_op_t()
        if not oper or oper.type != idaapi.o_imm:
            return None        
        value = oper.value
        if value in idautils.DataRefsFrom(self.iea):
            return None
        return value
    def get_string(self):
        mem = self.get_memory_address()
        if not mem:
            return None
        flags = idaapi.getFlags(mem)
        if not idaapi.isASCII(flags):
            return None
        tinfo = idaapi.opinfo_t()
        idaapi.get_opinfo(mem, 0, flags, tinfo)
        slen = idaapi.get_max_ascii_length(mem, tinfo.strtype)
        return idaapi.get_ascii_contents2(mem, slen, tinfo.strtype)
        
    def get_symbolicName(self):
        #Check if constant:
        value = self.get_immediate()
        if value is None:
            return None
        
        #Now lets try to see if it uses symbolic names
        operandCounter = 0
        operandText = GetOpnd(self.iea, self.onum)
        if (operandText != ""):
                if operandText[-1] == "h":
                    operandTextTemp = operandText[:-1]
                try:
                    result = int(operandTextTemp, 16)
                    return None
                except:
                    return operandText
        #End of symbolic names chases...
        
    def get_memory_address(self):
        oper = self.__get_op_t()
        if oper.type == idaapi.o_mem:
            return oper.addr
        elif oper.type == idaapi.o_imm and self.iea != idc.BADADDR:
            ref = oper.value
            if ref in idautils.DataRefsFrom(self.iea):
                return ref
        elif (oper.type == idaapi.o_displ or oper.type == idaapi.o_phrase) \
                 and not self.is_stackref():
            return oper.addr
        return None
    def is_stackref(self):
        oper = self.__get_op_t()
        if not oper.type in [idaapi.o_displ, idaapi.o_phrase]:
            return False
        offset = u2signed(oper.addr)
        return ( idaapi.get_stkvar(oper, offset) != None )
    def __get_op_t(self):
        ins = idautils.DecodeInstruction(self.iea)
        if not ins:
            return None
        else:
            return ins.Operands[self.onum]
             

 
#-----------------------------------------------------------------------
# IdaProExtractor 
# This class implements IDA Pro extraction capabilities. It also uses 
# a UI part (IdaProExtractorUI) in order to get 
# configuration selection from the User
#-----------------------------------------------------------------------
class IdaProTextExtractorScript():
    
    def __init__(self, sList=[], flg=0):
        self._sList = sList #func, eas for selected functions
        self._flg = flg
    #
    # This method is in charge of features extraction coordination from an IdaPro Input source
    # it returns an array containing all source features built with the following format:
    # [(functionName, [("searchTerm", "type") ..]), ..]
    # returns None if it encounters an error
    # Input file singleFile should be built like:
    # (fileName, [func_ea, ..])
    #
    def pluginExtract(self, _flg, _sList):
        if _flg:
            return self.getFunctionFeatures(self.getSelectedFunctionNamesAndEas(_sList))
        else:
            return self.getFunctionFeatures(self.getAllFunctionNamesAndEas())

    #
    # This method is used so plugin user could have the chance
    # to select functions to be extracted in a list of all the functions
    # that are inside a specific 
    #
   
    def getAllFunctionNamesAndEas(self):
        result = []
        allFuncEAS = get_all_function_eas()
        for ea in allFuncEAS:
            result.append((idaapi.get_func_name(ea), ea))
        return result


    def getSelectedFunctionNamesAndEas(self, _sList):
        result = []
        allFuncEAS = _sList
        for ea in allFuncEAS:
            result.append((idaapi.get_func_name(ea), ea))
        return result
 
    #
    # This method aims at returning all extracted features for a 
    # specific range of function. Its returns a features list.
    #
    def getFunctionFeatures(self, funcEaList=None):
        result = []
        if funcEaList:
            for cidx, currFunc in enumerate(funcEaList):
			
                tempFunc = FunctionX(currFunc[1])
                currFuncFeatures = self.extractFeatures(tempFunc)
                result.append(((tempFunc.get_name(),currFunc[1]), currFuncFeatures))
        #Check for strings!
        idbStringList = getAllStrings()
        
        for singleResult in result:
            for searchTerm in singleResult[1]:
                for singleString in idbStringList:
                    if searchTerm[1] in singleString[0]:
                        idbStringList.remove(singleString)
        
        alreadyInList = [] #This is used as a buffer to prevent duplicate strings
        for string in idbStringList:
            #Prevent duplicate strings
            if string[0] not in alreadyInList:
                #hash = hashlib.md5(string).hexdigest()
                result.append((("Orphan String " + str(string[0]), string[1]),[('s', str(string[0]))]))
                alreadyInList.append(string[0])
        #print len(idbStringList)
        #print result[0]        
        return result

    
    def extractFeatures(self, func):
        #print " --Extracting..."
        IDATA_SEGMENTS1 = get_import_segments()
        joint=[]
        name=[]
        demangle=[]
        constants=[]
        strlist=[]
        implist=[]
        impset = set()
        plist=[]
        rtype=[]
        arglist=[]
        numinst=[]
        numargs=[]
        argbytes=[]
        varbytes=[]
        fnflags=[]
        tagmal=[]
        tagapi=[]
        outcall=[]
        calllist=[]
        fname = func.get_name()
        dename = func.get_demangled_name()
        name.append(('n',fname))
        if dename:
            demangle.append(('ne',dename)) #ne: demangled C++ names
        for inst in func.get_instructions():
            for oper in inst.get_operands():
                val = oper.get_immediate()
                if val:
                    constants.append(('c', "%x" % val))
                    #Symbolic name can't be used in const we have to put it in strings
                    symbolicName = oper.get_symbolicName()
                    if symbolicName is not None and len(symbolicName) > 0:
                        try:
                            transition = int(symbolicName)
                        except:
                            strlist.append(('s',"%s" % symbolicName))
                    
                soper = oper.get_string()
                if soper:
                    if len(soper)>0:
                        #soper = replace_all(soper,BlackDict) 
                        strlist.append(('s',"%s" % soper))
                coderefs = inst.get_code_refs_from()
                for ref in coderefs:
                    if len([start for (start, end) in IDATA_SEGMENTS1 \
                            if start <= ref <= end]) > 0:
                        iname = idaapi.get_name(inst.iea, ref)
                        if iname:
                            impset.add( sanitize_name(iname) )
        for item in impset:
            implist.append( ( 'i', sanitize_name(item) ) )
        addr = func.start_ea 
        fend = idc.GetFunctionAttr(addr, idc.FUNCATTR_END)
        tt = getFuncDeclaration(addr)
        plist.append( ( 'p', tt ) ) #p: prototype
        params =  tt[tt.find("(")+1:tt.find(")")]
        lparams = params.split(',')
        sl = [(par.strip()) for par in lparams]
        for arg in sl:
            if arg:
                arglist.append( ( 'a', arg ) ) #a: arguments
        rt = tt[:tt.find(' ')]
        if rt:
            rtype.append( ( 'r', rt ) ) #r: return type    
        numinst.append( ( 'm', str(func.get_number_of_instructions()) ) ) #m: number of instructions
        g=func.get_number_of_arguments()
        if g:
            numargs.append( ( 'g', str(g) ) ) #g: number of arguments
        bbb=func.get_size_of_argument_bytes()
        if bbb:
            argbytes.append( ( 'b', str(bbb) ) ) #b: size of argument bytes
        varbytes.append( ( 'l', str(func.get_size_of_local_variables()) ) ) #l: size of local vars bytes
        fnflags.append( ( 'f', str(func.get_function_flags()) ) ) #f: function flags
        for elmntz in implist:
            if elmntz[1] in WinMalFuncInSight:
                tagmal.append(('d','MAL'))
        ad=0
        for api_row in api_matrix:
            l = api_row[0]
            apis = api_row[1:]
            for api in apis:
                ref_addrs = idautils.CodeRefsTo(idc.LocByName(api),0)
                for ref in ref_addrs:
                    func_name = idc.GetFunctionName(ref)
                    func_addr = idc.LocByName(func_name)
                    ad=addr 
                    if func_addr == ad:
                        tagapi.append(('t',l))
                    
                    
        bb=func.get_code_refs_from()
        outcall.append(('o', str( len(bb) ) )) # o: number of code refs from this function
        items = idautils.FuncItems(func.start_ea)
        for i in items:
            for xref in idautils.XrefsFrom(i, 0):
                if xref.type == idc.fl_CN or xref.type == idc.fl_CF:
                    dstf = idc.Name(xref.to)
                    calllist.append(('k', dstf )) # o: number of code refs from this function
                    demangled = idc.Demangle(dstf, idc.GetLongPrm(idc.INF_SHORT_DN))
                    if  demangled != None:
                        calllist.append(('ke', demangled )) #ke: demangled C++ call
        joint = name+demangle+constants+strlist+implist+rtype+arglist+numinst+numargs+argbytes+varbytes+fnflags+tagmal+tagapi+outcall+calllist
        z=fend - ad
        if z:
            joint.append(('z', str(z))) #function size (bytes)
        joint.append(('cx',str(len(constants))))
        joint.append(('sx',str(len(strlist))))
        joint.append(('mx',str(numinst[0][1])))
        if g:
            joint.append(('gx',str(numargs[0][1])))
        joint.append(('ix',str(len(implist))))
        joint.append(('dx',str(len(tagmal))))
        joint.append(('tx',str(len(tagapi))))
        joint.append(('ox',str(len(outcall))))
        joint.append(('kx',str(len(calllist))))
        return joint

class AsmFunctionList(Choose2):
    def __init__(self, title, flags=0):
        Choose2.__init__(self,
                         title,
                         [ ["Address", 10 | Choose2.CHCOL_HEX],
                           ["Function Name", 15 | Choose2.CHCOL_PLAIN] ],
                         embedded=True, width=25, height=5, flags=flags)
        self.n = 0
        self.selcount = 0
        self.icon = 41
        self.PopulateItems()
    def PopulateItems(self):
        global funcListG
        funcListG = [ [hex(x), GetFunctionName(x), x] for x in idautils.Functions() ] #x=ea
        self.items = funcListG
    def OnClose(self):
        pass
    def OnRefresh(self, n):
        self.Refresh()
        return self.n
    def OnRefreshed(self):
        return self.n
    def OnSelectLine(self, n):
        self.selcount += 1
    def OnGetLine(self, n):
        return self.items[n]
    def OnGetSize(self):
        n = len(self.items)
        return n
		
class ReSrcForm(Form):
    def __init__(self):
        global funcListG
        self.EChooser = AsmFunctionList("asm", flags=Choose2.CH_MULTI)
        Form.__init__(self,
r"""BUTTON YES* NONE
BUTTON CANCEL Exit
BinSourcerer
{FormChangeCb}
Function Selection:
<Please use [ctrl+click] or [shift+click] for multiple selections.:{cEChooser}>
Search Configuration:
<#Selecting this option will disable the list.#Select all functions:{rSearchAll}>
<Use offline analyzer for malware:{rOfflineMal}>{cSearchCfgGroup}>
Feature Extraction:
<Path:{iPath}>
<File:{iFSave}>
<Feature Extraction:{iBatchM}><API Tagging:{iAPI}>

""", {
            'cSearchCfgGroup': Form.ChkGroupControl(("rSearchAll", "rOfflineMal")),
            'iBatchM': Form.ButtonInput(self.OnFormChange),
            'iPath': Form.DirInput(swidth=30),
	    'iFSave': Form.FileInput(save=True, swidth=30),
            'iAPI': Form.ButtonInput(self.OnFormChange),
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
            'cEChooser' : Form.EmbeddedChooserControl(self.EChooser)
        })
        
    def validPath(self, inputFPath):
        if os.access(os.path.dirname(inputFPath), os.W_OK):
            return True
        else:
            return False
        
    def OnFormChange(self, fid):
        if fid == self.iBatchM.id:
            p = self.GetControlValue(self.iPath)
	    f = self.GetControlValue(self.iFSave)
	    if p and f and self.validPath(p+"\\"+f):
                print "Status: BinSourcerer is ready for magic..."
                IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH = p + "\\"
                IDA_PRO_TEXT_EXTRACTOR_EXTRACTED_FILE = f
                print "Status: Extracting features. This could take a few minutes..."

                stt = self.GetControlValue(self.rSearchAll)
                if stt==1:
                    x = IdaProTextExtractorScript()
                    results = x.pluginExtract(0, [])
                else:
                    global funcListG
                    selFuncEAS = []
                    selectedFuncsList = self.GetControlValue(self.cEChooser)
                    if len(selectedFuncsList) > 0:
                            for funcIndex in selectedFuncsList:
                                    selFuncEAS.append( funcListG[funcIndex][2] )
                            x = IdaProTextExtractorScript(selFuncEAS, 1)
                            results = x.pluginExtract(1, selFuncEAS)
                    else:
                            idc.Warning("Please select multiple functions from the list.")
                print "Status: Features extraction completed. Saving results..."
                try:
                    f = open(IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH + IDA_PRO_TEXT_EXTRACTOR_EXTRACTED_FILE, "w")
                    if results:
                        for func in results:
                            f.write(str(func) + "\n")
			outputPath = IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH + IDA_PRO_TEXT_EXTRACTOR_EXTRACTED_FILE
                    
                        print "Status: Features file saved in " + os.path.normpath(outputPath)
                except:
                    print "I/O Error! Could not write to the output file!"
                finally:
                    f.close()
                    
            else:
                idc.Warning("Please select a valid path and file name!\n")
       
        elif fid == self.rSearchAll.id:
            stt = self.GetControlValue(self.rSearchAll)
            if stt==1:
                self.EnableField(self.cEChooser, False)
            else:
                self.EnableField(self.cEChooser, True)

	elif fid == self.iAPI.id:
            potentialTags=APITags(0)
            choice=AskYN(1,"Approximately "+ str(potentialTags)+" tags are to be assigned to functions.\nProceed with name changes?")
            if choice==1:
                APITags(1)
                self.EChooser.PopulateItems()
                self.RefreshField(self.cEChooser)
        return 1

def APITags(flag):
    apilist=[]
    if flag==1:
        save=True
    else:
        save=False
    count=0
    for idx, api_row in enumerate(api_matrix):
        l = api_row[0]
        apis = api_row[1:]
        for api in apis:
            ref_addrs = CodeRefsTo(LocByName(api),0)
            for ref in ref_addrs:
                func_addr = LocByName(GetFunctionName(ref))
                fnsAddr = GetFunctionAttr( func_addr, FUNCATTR_START )
                func_name = GetFunctionName(ref)
                if l not in func_name:
                    count +=1 #counter
                    apilist.append( ( 't', l ) ) #t: api tags
                    ss = l + '_' + func_name
                    if save:
                        MakeNameEx(func_addr , ss  , SN_NOWARN)		
    return count		
#-------------------------------------------------------
# The following vars have to be configured to matching
# values from the BinSourcerer configuration file
#-------------------------------------------------------
IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH = 'C:\\Users\\UserName\\Desktop\\Features\\'
IDA_PRO_TEXT_EXTRACTOR_EXTRACTED_FILE = 'Features.xtrak'
funcListG = []

class TextExtract_t(idaapi.plugin_t):
    flags = 0
    comment = "Extract all the available features"
    help = "Features will be saved in the output file: [extractedFeatures.extract]"
    wanted_name = "BinSourcerer Features Extractor"
    wanted_hotkey = "Ctrl-Shift-J"
   
    def init(self):
        idaapi.msg("Status: BinSourcerer Features Extractor plug-in loaded.\n")
        return idaapi.PLUGIN_KEEP
    def run(self, arg):
	global fm
	fm = ReSrcForm()
	fm.Compile()
        fm.rSearchAll.checked = False
	fm.rOfflineMal.checked = True
	cwdPath = os.getcwd() + "\\"
	IDA_PRO_TEXT_EXTRACTOR_EXTRACTION_PATH = cwdPath	
	fm.iPath.value = cwdPath
	fm.iFSave.value = "Features.xtrak"
	ok = fm.Execute()
	fm.Free()
    def term(self):
        pass

def PLUGIN_ENTRY():
    return TextExtract_t()



    


repository: zmrbak/IDAPython
file: 配套代码/ida10/python/examples/hexrays/vds17.py
description: 在线课程《IDAPython脚本化软件逆向分析入门》课程源码
wanted_name: Structure offsets (IDAPython)
comment: Sample plugin17 for Hex-Rays decompiler
created_at: 2022-05-02 02:13:34+00:00
pushed_at: 2022-05-14 10:41:43+00:00
url:
https://github.com/zmrbak/IDAPython/blob/478278339ba0c150b3ace657296233d8a89b26ea/%E9%85%8D%E5%A5%97%E4%BB%A3%E7%A0%81/ida10/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: a1ext/ida-embed-arch-disasm
file: ida-embed-arch-disasm.py
description: Allows IDA PRO to disassemble x86-64 code (WOW64) in 32-bit database
wanted_name: IDA Embed ARCH Disasm as x86_64
created_at: 2018-06-29 11:20:36+00:00
pushed_at: 2021-12-27 16:01:58+00:00
url:
https://github.com/a1ext/ida-embed-arch-disasm/blob/aa86249caff5e488e2373b7b357f5eaa6e531052/ida-embed-arch-disasm.py

content:
# -*- coding: utf-8 -*
__author__ = 'Trafimchuk Aliaksandr'

import idaapi
import idc


class MyHandler(idaapi.action_handler_t):
    @classmethod
    def get_name(cls):
        return cls.__name__
        
    @classmethod
    def get_label(cls):
        return cls.label
        
    @classmethod
    def register(cls, plugin, label):
        cls.plugin = plugin
        cls.label = label
        instance = cls()
        return idaapi.register_action(idaapi.action_desc_t(
            cls.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(cls):
        """Unregister the action.
        After unregistering the class cannot be used.
        """
        idaapi.unregister_action(cls.get_name())

    def activate(self, ctx):
        start, end = idc.read_selection_start(), idc.read_selection_end()
        if start == idaapi.BADADDR:
            print 'Please select something'
            return

        import capstone
        md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
        md.details = True
        data = idaapi.get_bytes(start, end - start)
        for insn in md.disasm(data, start):
            # print "0x%x:\t%s\t%s" % (insn.address, insn.mnemonic, insn.op_str)
            idaapi.set_cmt(insn.address, str('%s %s' % (insn.mnemonic, insn.op_str)), False)

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


class Hooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        tft = idaapi.get_widget_type(form)

        if tft == idaapi.BWN_DISASM:
            # Define a silly handler.

            # Note the 'None' as action name (1st parameter).
            # That's because the action will be deleted immediately
            # after the context menu is hidden anyway, so there's
            # really no need giving it a valid ID.
            idaapi.attach_action_to_popup(form, popup, MyHandler.get_name(), "-")


class ida_embed_arch_disasm_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = ""

    help = ""
    wanted_name = "IDA Embed ARCH Disasm as x86_64"
    wanted_hotkey = ""

    def __init__(self):
        super(ida_embed_arch_disasm_t, self).__init__()
        self._data = None
        self.hooks = None

    def init(self):
        MyHandler.register(self, self.wanted_name)
        self.hooks = Hooks()
        self.hooks.hook()
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        # refer to https://github.com/aquynh/capstone/blob/master/bindings/python/test_x86.py
        # [about actions in the menus] http://www.hexblog.com/?p=886
        pass

    def term(self):
        self._data = None
        self.hooks.unhook()
        self.hooks = None


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


repository: Russinovich/IDA7.5_SDK
file: python/examples/core/install_user_defined_prefix.py
description: SDK and IDAPython for IDA Pro 7.5
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2021-09-01 02:18:48+00:00
pushed_at: 2021-07-23 08:19:09+00:00
url:
https://github.com/Russinovich/IDA7.5_SDK/blob/5244a2b5059113e421b4da4ecf771013a6fde2f8/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: Cisco-Talos/GhIDA
file: ghida.py
url: https://github.com/Cisco-Talos/GhIDA/blob/a396916ae53e46adf3dca918b810f5db046af015/ghida.py
wanted_name: GhIDA Decompiler
comment: GhIDA Decompiler for IDA Pro
created_at: 2019-09-03 15:54:11+00:00
pushed_at: 2021-02-19 03:30:46+00:00
content:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

##############################################################################
#                                                                            #
#  GhIDA: Ghidra decompiler for IDA Pro                                      #
#                                                                            #
#  Copyright 2019 Andrea Marcelli, Cisco Talos                               #
#                                                                            #
#  Licensed under the Apache License, Version 2.0 (the "License");           #
#  you may not use this file except in compliance with the License.          #
#  You may obtain a copy of the License at                                   #
#                                                                            #
#      http://www.apache.org/licenses/LICENSE-2.0                            #
#                                                                            #
#  Unless required by applicable law or agreed to in writing, software       #
#  distributed under the License is distributed on an "AS IS" BASIS,         #
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  #
#  See the License for the specific language governing permissions and       #
#  limitations under the License.                                            #
#                                                                            #
##############################################################################

try:
    from pygments.lexers import CLexer
    from pygments.token import Token
except Exception:
    # Missing library is managed at the plugin entry
    pass


import ida_kernwin
import idaapi
import idautils
import idc

import ghida_plugin as gl

DECOMP_VIEW = None
GHIDA_CONF = None
DECOMPILED_CACHE = None
COMMENTS_CACHE = None


# ------------------------------------------------------------
#   IDA MENU - HANDLERS
# ------------------------------------------------------------


class ShowSettingsHandler(idaapi.action_handler_t):

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

    # Say hello when invoked.
    def activate(self, ctx):
        display_configuration_form()
        return 1

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


class ShowDecompWindowHandler(idaapi.action_handler_t):

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

    # Say hello when invoked.
    def activate(self, ctx):
        if DECOMP_VIEW:
            DECOMP_VIEW.Show()
        else:
            # print("GhIDA:: [DEBUG] DECOMP_VIEW non existing")
            pass
        return 1

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


# ------------------------------------------------------------
#   DECOMPILED VIEW -- POP-UP HANDLERS
# ------------------------------------------------------------

class GoToCustViewerHandler(idaapi.action_handler_t):

    def __init__(self, view):
        idaapi.action_handler_t.__init__(self)
        self.view = view

    # Say hello when invoked.
    def activate(self, ctx):
        # print("GhIDA:: [DEBUG] GoToCustViewerHandler HELLO")
        goto()
        return 1

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


class AddCommentCustViewerHandler(idaapi.action_handler_t):

    def __init__(self, view):
        idaapi.action_handler_t.__init__(self)
        self.view = view

    # Say hello when invoked.
    def activate(self, ctx):
        # print("GhIDA:: [DEBUG] AddCommentCustViewerHandler HELLO")
        DECOMP_VIEW.add_comment()
        return 1

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


class RenameCustViewerHandler(idaapi.action_handler_t):

    def __init__(self, view):
        idaapi.action_handler_t.__init__(self)
        self.view = view

    # Say hello when invoked.
    def activate(self, ctx):
        # print("GhIDA:: [DEBUG] RenameCustViewerHandler HELLO")
        DECOMP_VIEW.rename_symbol()
        return 1

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


# ------------------------------------------------------------
#   ScreenEAHook
# ------------------------------------------------------------

class ScreenEAHook(ida_kernwin.View_Hooks):

    def __init__(self):
        ida_kernwin.View_Hooks.__init__(self)
        print("GhIDA [DEBUG] ScreenEAHook initialized")

    def view_loc_changed(self, widget, curloc, prevloc):
        """
        view_loc_changed is called each time the user clicks
        somwhere. This is used to synchronize the IDA DISASM
        view with the IDA DECOM view. The synchronization is
        active only when the decompile view has been created
        and the synch option has been selected in the pop-up
        menu.
        """
        # Check if the selected address has changed
        # if curloc.plce.toea() != prevloc.plce.toea():
        #     return

        # Hooking the IDA DISASM view only
        if idaapi.get_widget_type(widget) != idaapi.BWN_DISASM:
            return

        # If the DECOMP view has already been created.
        if DECOMP_VIEW:
            # Get the new address
            ca = curloc.plce.toea()
            ea = gl.convert_address(ca)

            # This is a valid function address
            if ea:
                # The synch is active
                if GHIDA_CONF.disasm_tracker:
                    # The address in DECOMP view is different
                    if ea != DECOMP_VIEW.ea:
                        # Update DECOMP view
                        DECOMP_VIEW.switch_to_address(ea)

                # Update the selection
                return gl.highlight_symbol_in_DECOMP()

            # This is not a valid function address
            if not ea:
                # If the synch is active
                if GHIDA_CONF.disasm_tracker:
                    DECOMP_VIEW.clear(msg="[!] Function not found.",
                                      do_show=False)
        return


# ------------------------------------------------------------
#   GOTO utils
# ------------------------------------------------------------

def goto(shift=False):
    # print("GhIDA:: [DEBUG] goto called")

    symbol = None
    ret = ida_kernwin.get_highlight(ida_kernwin.get_current_viewer())
    if ret and ret[1]:
        symbol = ret[0]

    if not symbol:
        return False

    address = gl.get_address_for_symbol(symbol)
    if not address:
        return False

    print("OnDblClick, shift=%d, selection:%s, address:%s" %
          (shift, symbol, address))

    # Update IDA DISASM view
    idaapi.jumpto(address)

    # Update IDA DECOMP view
    ea = gl.convert_address(address)
    # print("GhIDA:: [DEBUG] update view to %s" % ea)
    DECOMP_VIEW.switch_to_address(ea)

    return True


# ------------------------------------------------------------
#   SIMPLECUSTVIEWER FOR THE DECOMPILED RESULT
# ------------------------------------------------------------

# Check this example: https://github.com/nologic/idaref/blob/master/idaref.py
class DecompiledViewer_t(idaapi.simplecustviewer_t):

    def color_line(self, line):
        """
        """
        lexer = CLexer()
        tokens = list(lexer.get_tokens(line))
        new_line = ""
        for t in tokens:
            ttype = t[0]
            ttext = str(t[1])
            if ttype == Token.Text:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_INSN)

            elif ttype == Token.Text.Whitespace:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_INSN)

            elif ttype == Token.Error:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_ERROR)

            elif ttype == Token.Other:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_DSTR)

            elif ttype == Token.Keyword:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_KEYWORD)

            elif ttype == Token.Name:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_LIBNAME)

            elif ttype == Token.Literal:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_LOCNAME)

            elif ttype == Token.Literal.String:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_STRING)

            elif ttype == Token.Literal.Number:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_DNUM)

            elif ttype == Token.Operator:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_ALTOP)

            elif ttype == Token.Punctuation:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_SYMBOL)

            elif ttype == Token.Comment:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_REGCMT)

            elif ttype == Token.Comment.Single:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_REGCMT)

            elif ttype == Token.Generic:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_CREFTAIL)

            else:
                new_line += idaapi.COLSTR(ttext, idaapi.SCOLOR_CREFTAIL)
        return new_line

    def add_colored_text(self, text):
        """
        Parse the code with the C lexer and display the colored text
        in the decompiled view.
        """
        for line in text.splitlines():
            new_line = self.color_line(line)
            self.AddLine(new_line)
        return

    def Create(self, decompiled, ea):
        """
        Create a new view with the decompiled code
        """
        title = "Decompiled Function"
        self.__ea = ea
        self.__decompiled = decompiled

        # Create the customviewer
        if not idaapi.simplecustviewer_t.Create(self, title):
            return False

        # for line in decompiled.splitlines():
        #     self.AddLine(str(line))
        self.add_colored_text(decompiled)

        return True

    def clear(self, msg=None, do_show=True):
        """
        Clear the view content
        """
        self.ClearLines()
        self.__ea = None
        self.__decompiled = None
        if msg:
            for line in msg.splitlines():
                self.AddLine(line)
        self.Refresh()
        if do_show:
            self.Show()
        return

    def update(self, ea, decompiled, do_show=True):
        """
        Update the content of the view with the new decompiled code
        """
        self.__ea = ea
        self.__decompiled = decompiled
        self.ClearLines()
        self.add_colored_text(decompiled)
        self.Refresh()
        if do_show:
            self.Show()

        # Update the cache
        DECOMPILED_CACHE.update_decompiled_cache(ea, decompiled)
        # print("GhIDA:: [DEBUG] GhIDA DECOM view updated to %s" % ea)
        return

    def switch_to_address(self, ea):
        """
        The IDA DIASM view switched to a new address, change
        the decompiled view accordingly.
        """
        self.__ea = ea
        decompile_function_wrapper(cache_only=True, do_show=False)
        return

    def add_comment(self):
        """
        Add a commment to the selected line
        """
        # print("GhIDA:: [DEBUG] add_comment called")
        colored_line = self.GetCurrentLine(notags=1)
        if not colored_line:
            idaapi.warning("Select a line")
            return False

        # Use pygments to parse the line to check if there are comments
        line = idaapi.tag_remove(colored_line)
        lexer = CLexer()
        tokens = list(lexer.get_tokens(line))
        text = ""
        text_comment = ""
        for t in tokens:
            ttype = t[0]
            ttext = str(t[1])
            if ttype == Token.Comment.Single:
                text_comment = ttext.replace('//', '').strip()
            else:
                text += ttext

        # Get the new comment
        comment = gl.display_comment_form(text_comment)
        if not comment or len(comment) == 0:
            return False
        comment = comment.replace("//", "").replace("\n", " ")
        comment = comment.strip()

        # Create the new text
        full_comment = "\t// %s" % comment
        text = text.rstrip()
        new_text = text + full_comment
        text_colored = self.color_line(new_text)

        num_line = self.GetLineNo()
        self.EditLine(num_line, text_colored)
        self.RefreshCurrent()

        # Add comment to cache
        COMMENTS_CACHE.add_comment_to_cache(self.__ea, num_line, full_comment)

        # print("GhIDA:: [DEBUG] Added comment to #line: %d (%s)" %
        #       (num_line, new_text))
        return

    def add_comments(self, comment_list):
        """
        Updated the view with the available comments
        """
        for item in comment_list:
            lineno = item[0]
            comment = item[1]
            if len(comment) == 0:
                continue
            line = self.GetLine(lineno)
            if not line:
                print("GhIDA:: [!] line not found")
                continue
            line_text = line[0]
            if not line_text:
                print("GhIDA:: [!] line-text not found")
                continue
            line_text = idaapi.tag_remove(line_text) + comment
            new_line = self.color_line(line_text)
            self.EditLine(lineno, new_line)

        self.Refresh()
        # print("GhIDA:: [DEBUG] updated comments terminated")
        return

    def rename_symbol(self):
        """
        Rename the symbol "symbol" with the new name
        provided by the user in the Pop-Up
        """
        # Get the symbol
        symbol = None
        ret = ida_kernwin.get_highlight(ida_kernwin.get_current_viewer())
        if ret and ret[1]:
            symbol = ret[0]

        if not symbol:
            idaapi.warning("Select a symbol")
            return False

        # Get the address
        address = gl.get_address_for_symbol(symbol)
        if not address:
            print("GhIDA:: [!] Symbol %s not found" % symbol)
            return False

        # Display a Pop-up to get the new name
        new_name = gl.display_rename_form(address, symbol)
        if not new_name or len(new_name) == 0:
            return

        # Check for white_spaces in the new symbol name
        for letter in new_name:
            if not (letter.isdigit() or letter.isalpha() or letter == '_'):
                print("GhIDA:: [!] symbol name contains invalid char")
                return

        # Check if new_name is already used
        if gl.check_if_symbol_is_used(new_name):
            print("GhIDA:: [!] symble name already used")
            return

        # Update symbol name in SYMBLE DICT:
        gl.updated_symbol_name_for_address(symbol, address, new_name)

        # Update symbol name in IDA DISASM view.
        # print("GhIDA:: [DEBUG] New symbol name: %s" % new_name)

        # Update symbol name in the decompiled view
        new_code = gl.rename_variable_in_text(
            self.__decompiled,
            symbol,
            new_name)
        self.update(self.__ea, new_code)

        # Add comments
        comment_list = COMMENTS_CACHE.get_comments_cache(self.__ea)
        if comment_list:
            self.add_comments(comment_list)

        print("GhIDA:: [INFO] Symbol name updated in IDA DECOMP view.")

        if idc.set_name(address, new_name):
            # Refresh the view
            idaapi.request_refresh(idaapi.IWID_DISASMS)
            # Highlight the new identifier
            gl.highlight_symbol_in_DISASM()
            print("GhIDA:: [INFO] Symbol name updated in IDA DISASM view.")
            return

        print("GhIDA:: [!] IDA DISASM rename error")
        return

    def OnClick(self, shift):
        """
        User clicked in the view
        @param shift: Shift flag
        @return: Boolean. True if you handled the event
        """
        return gl.highlight_symbol_in_DISASM()

    def OnDblClick(self, shift):
        """
        User dbl-clicked in the view
        @param shift: Shift flag
        @return: Boolean. True if you handled the event
        """
        return goto(shift)

    def OnKeydown(self, vkey, shift):
        """
        User pressed a key
        @param vkey: Virtual key code
        @param shift: Shift flag
        @return: Boolean. True if you handled the event
        """
        # print("OnKeydown, vk=%d shift=%d" % (vkey, shift))

        # # Esc --> close the window
        if vkey == 27:
            # TODO.. go back in the history, as disasm does
            return True

        # N or n --> rename a symbol
        if vkey == ord('N'):
            self.rename_symbol()

        # : --> add a comment
        elif vkey == 186 and shift == 1:
            self.add_comment()

        else:
            return False

        return True

    @property
    def ea(self):
        return self.__ea


# ------------------------------------------------------------
#   HANDLERS FOR THE POP-UP MENU IN DISASMS VIEW
# ------------------------------------------------------------

class InvalidateCache(idaapi.action_handler_t):

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

    # Say hello when invoked.
    def activate(self, ctx):
        # print("GhIDA:: [DEBUG] InvalidateCache HELLO")
        address = gl.get_current_address()
        if not address:
            # print("GhIDA:: [DEBUG] address not found")
            return

        DECOMPILED_CACHE.invalidate_cache(address)
        gl.force_export_XML_file()

        return 1

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


class DisasmTracker(idaapi.action_handler_t):

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

    # Say hello when invoked.
    def activate(self, ctx):
        # print("GhIDA:: [DEBUG] DisasmTracker HELLO")

        if GHIDA_CONF.disasm_tracker:
            GHIDA_CONF.disasm_tracker = False
            print("GhIDA:: [INFO] synchronization disabled")

            # Update the description in the pop-up menu.
            ida_kernwin.update_action_label(
                "my:disasmtracker", "Enable decompile view synchronization")

        else:
            GHIDA_CONF.disasm_tracker = True
            print("GhIDA:: [INFO] synchronization activated")

            # Update the description in the pop-up menu.
            ida_kernwin.update_action_label(
                "my:disasmtracker", "Disable decompile view synchronization")

            # Update the DECOMP view to the current address
            ea = gl.get_current_address()
            if ea:
                DECOMP_VIEW.switch_to_address(ea)
        return 1

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


class DisasmsHandler(idaapi.action_handler_t):

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

    # Say hello when invoked.
    def activate(self, ctx):
        # print("GhIDA:: [DEBUG] DisasmsHandler HELLO")
        decompile_function_wrapper()
        return 1

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


class DisasmsHooks(idaapi.UI_Hooks):

    def finish_populating_tform_popup(self, form, popup):
        # TODO - Attach to the functions view.
        # if idaapi.get_tform_type(form) == idaapi.BWN_FUNCS:
        #     idaapi.attach_action_to_popup(
        #         form, popup, "my:disasmsaction", None)

        # Attach to the disassembler view only
        if idaapi.get_tform_type(form) == idaapi.BWN_DISASMS:
            idaapi.attach_action_to_popup(
                form, popup, "my:disasmsaction", None)
            idaapi.attach_action_to_popup(
                form, popup, "my:disasmtracker", None)
            idaapi.attach_action_to_popup(
                form, popup, "my:invalidatecache", None)


def register_handlers():
    """
    Register the handlers for the pop-up menu to interact with the UI
    """
    # print("GhIDA:: [DEBUG] Registering handlers")

    # Load a custom icon
    icon_path = gl.plugin_resource("ghida.png")
    icon_data = str(open(icon_path, "rb").read())
    icon_ghida = idaapi.load_custom_icon(data=icon_data)

    idaapi.register_action(idaapi.action_desc_t(
        "my:disasmsaction",
        "Decompile function with GhIDA",
        DisasmsHandler(),
        None,
        'IDA plugin for Ghidra decompiler',
        icon_ghida))

    disasmtracker_action = idaapi.action_desc_t(
        "my:disasmtracker",
        "Disable decompile view synchronization",
        DisasmTracker(),
        None,
        None,
        icon_ghida)
    idaapi.register_action(disasmtracker_action)

    idaapi.register_action(idaapi.action_desc_t(
        "my:invalidatecache",
        "Clear cache for current function",
        InvalidateCache(),
        None,
        None,
        icon_ghida))

    # Add the settings item in the menu
    show_settings_action = idaapi.action_desc_t(
        'my:showsettingsaction',
        'GhIDA Settings',
        ShowSettingsHandler(),
        None,
        'GhIDA Settings',
        icon_ghida)
    idaapi.register_action(show_settings_action)

    idaapi.attach_action_to_menu(
        'Edit/Settings/GhIDA Settings',
        'my:showsettingsaction',
        idaapi.SETMENU_APP)

    # Add the view decompile window in the menu
    show_decomp_window_action = idaapi.action_desc_t(
        'my:showdecompilewindowaction',
        'GhIDA decomp view',
        ShowDecompWindowHandler(),
        None,
        'GhIDA decomp view',
        icon_ghida)
    idaapi.register_action(show_decomp_window_action)

    idaapi.attach_action_to_menu(
        'View/Open subviews/GhIDA',
        'my:showdecompilewindowaction',
        idaapi.SETMENU_APP)

    return


# ------------------------------------------------------------
#   DEFINING GLOBAL VARS AND GHIDA INITIAL CONFIGURATION
# ------------------------------------------------------------

def load_configuration():
    """
    """
    global GHIDA_CONF
    global DECOMPILED_CACHE
    global COMMENTS_CACHE

    # Loading the plugin configuration
    # print("GhIDA:: [DEBUG] Reading GhIDA configuration")
    GHIDA_CONF = gl.GhidaConfiguration()

    print("GHIDA_CONF.load_save_cached_code",
          GHIDA_CONF.load_save_cached_code)
    print("GHIDA_CONF.load_save_cached_comments",
          GHIDA_CONF.load_save_cached_comments)

    md5 = idautils.GetInputFileMD5()

    # Initalize the cache (and load cached objects)
    DECOMPILED_CACHE = gl.DecompiledCache(
        file_id=md5,
        use_cache=GHIDA_CONF.load_save_cached_code)
    COMMENTS_CACHE = gl.CommentsCache(
        file_id=md5,
        use_cache=GHIDA_CONF.load_save_cached_comments)

    return

# ------------------------------------------------------------
#   HANDLERS FOR THE POP-UP MENU IN DECOMP VIEW
# ------------------------------------------------------------


def register_actions_and_handlers_decompile_view():
    """
    Attach the following actions in the pop-up menu of the
    decompiled view.
    """
    # Load a custom icon
    icon_path = gl.plugin_resource("ghida.png")
    icon_data = str(open(icon_path, "rb").read())
    icon_ghida = idaapi.load_custom_icon(data=icon_data)

    decompiler_widget = idaapi.find_widget('Decompiled Function')
    # TODO alternative
    # decompiler_widget = idaapi.get_current_tform()

    # Add Rename to the pop-up
    action_renamecustviewer = idaapi.action_desc_t(
        'my:renamecustviewerhandler',
        'Rename',
        RenameCustViewerHandler(
            DECOMP_VIEW),
        None,
        None,
        icon_ghida)
    decompiler_widget = idaapi.find_widget('Decompiled Function')
    idaapi.register_action(action_renamecustviewer)
    idaapi.attach_action_to_popup(decompiler_widget,
                                  None,
                                  "my:renamecustviewerhandler",
                                  None)

    # Add add-comment to the pop-up
    action_addcommentcustviewer = idaapi.action_desc_t(
        'my:addcommentcustviewer',
        'Add comment',
        AddCommentCustViewerHandler(
            DECOMP_VIEW),
        None,
        None,
        icon_ghida)
    idaapi.register_action(action_addcommentcustviewer)
    idaapi.attach_action_to_popup(decompiler_widget,
                                  None,
                                  "my:addcommentcustviewer",
                                  None)

    # Add goto to the pop-up
    action_gotocustviewerhandler = idaapi.action_desc_t(
        'my:gotocustviewerhandler',
        'Goto',
        GoToCustViewerHandler(
            DECOMP_VIEW),
        None,
        None,
        icon_ghida)
    idaapi.register_action(action_gotocustviewerhandler)
    idaapi.attach_action_to_popup(decompiler_widget,
                                  None,
                                  "my:gotocustviewerhandler",
                                  None)
    return


# ------------------------------------------------------------
#   GHIDA Configuration FORM
# ------------------------------------------------------------

def display_configuration_form():
    """
    Display a configuration dialog for the user.
    """
    f = gl.GhIDASettingsForm()
    f.Compile()
    f.GhidraInstallationPath.value = GHIDA_CONF.ghidra_install_path
    f.GhidraaasURL.value = GHIDA_CONF.ghidra_server_url
    r = f.Execute()

    if r == 1:  # OK
        # Do not display the settings menu anymore
        GHIDA_CONF.global_settings = True

        if f.cGroup.value == 0:
            # Use local Ghidra
            GHIDA_CONF.use_ghidra_server = False
            GHIDA_CONF.ghidra_install_path = f.GhidraInstallationPath.value
        else:
            # Use Ghidra server (Ghidraaas)
            GHIDA_CONF.use_ghidra_server = True
            GHIDA_CONF.ghidra_server_url = f.GhidraaasURL.value

        if f.cGroup1.value == 0:
            # Do not display the popup at startup
            GHIDA_CONF.show_settings = True
        else:
            # Dispaly the menu at startup
            GHIDA_CONF.show_settings = False

        if f.cGroup2.value == 0:
            # Do not save cache to file
            GHIDA_CONF.load_save_cached_code = False
            GHIDA_CONF.load_save_cached_comments = False
        else:
            # Save cache to file
            GHIDA_CONF.load_save_cached_code = True
            GHIDA_CONF.load_save_cached_comments = True

        print("GHIDA_CONF.global_settings", GHIDA_CONF.global_settings)
        print("GHIDA_CONF.use_ghidra_server", GHIDA_CONF.use_ghidra_server)
        print("GHIDA_CONF.ghidra_install_path", GHIDA_CONF.ghidra_install_path)
        print("GHIDA_CONF.ghidra_server_url", GHIDA_CONF.ghidra_server_url)
        print("GHIDA_CONF.show_settings", GHIDA_CONF.show_settings)
        print("GHIDA_CONF.load_save_cached_code",
              GHIDA_CONF.load_save_cached_code)
        print("GHIDA_CONF.load_save_cached_comments",
              GHIDA_CONF.load_save_cached_comments)

        # Save configuration to file
        GHIDA_CONF.dump_to_json()
        return True

    # Canceled
    return False


# ------------------------------------------------------------
#   DECOMPILE FUNCTION - CORE
# ------------------------------------------------------------

def decompile_function_wrapper(cache_only=False, do_show=True):
    """
    Perform all the operations to decompile the code of the selected
    function and display in the decompile view.
    """
    try:
        global DECOMP_VIEW
        ea = gl.get_current_address()
        if not ea:
            # This is not a function
            # GhIDA can decompile only IDA recognized functions.
            return

        # Set the base_image
        image_base = idaapi.get_imagebase()
        if GHIDA_CONF.image_base is None:
            GHIDA_CONF.image_base = image_base

        # Check if the program has been rebased
        if GHIDA_CONF.image_base != image_base:
            print(
                "GhIDA:: [DEBUG] program has been rebased. Invalidating caches.")
            DECOMPILED_CACHE.invalidate_cache()
            COMMENTS_CACHE.invalidate_cache()
            gl.force_export_XML_file()

        # Display the Configuration form
        if GHIDA_CONF.show_settings and \
                not GHIDA_CONF.global_settings:
            canceled = not(display_configuration_form())
            if canceled:
                return

        # If exists, clear the decompile view.
        if DECOMP_VIEW:
            DECOMP_VIEW.clear(msg="Decompiling function...",
                              do_show=do_show)

        # Call export XML file. It also populates the SYMBOL DICT TABLE (S_D_T)
        # for the highlighting, renaming, etc.
        # Do it here because I need S_D_T it also if the result is in the
        # cache.
        gl.export_ida_project_to_xml()

        # Check the cache
        decompiled = DECOMPILED_CACHE.get_decompiled_cache(ea)

        # Cache miss - opt1: do not use Ghidra, just display what is in the
        # cache
        if not decompiled and cache_only:
            # This is a redundant check...
            if DECOMP_VIEW:
                msg = "Function 0x%s\n\n" % ea
                msg += "Decompiled code not available in cache.\n"
                msg += "Press Ctrl-Alt-D or Right click GhIDA decompiler "
                msg += "to decompile the function."
                DECOMP_VIEW.clear(msg=msg, do_show=do_show)
            # print("GhIDA:: [DEBUG] Function code not available in cache.")
            return

        # Cache miss - opt2: Use Ghidra to decompile the function
        if not decompiled:
            decompiled = gl.decompile_function(
                address=ea,
                use_ghidra_server=GHIDA_CONF.use_ghidra_server,
                ghidra_headless_path=GHIDA_CONF.ghidra_headless_path,
                ghidra_plugins_path=GHIDA_CONF.ghidra_plugins_path,
                ghidra_server_url=GHIDA_CONF.ghidra_server_url)

            if decompiled:
                # Add decompiled_code to cache
                DECOMPILED_CACHE.add_decompiled_to_cache(ea, decompiled)
            else:
                # Something went wrong or was interrupted
                if DECOMP_VIEW:
                    DECOMP_VIEW.clear(msg="Decompiling interrupted.",
                                      do_show=do_show)
                print("GhIDA:: [!] Decompilation interrupted.")
                return

        # Decompiled code is available.
        if DECOMP_VIEW:
            # The view exists, update it
            DECOMP_VIEW.update(ea, decompiled, do_show=do_show)
        else:
            # Create the view
            DECOMP_VIEW = DecompiledViewer_t()
            if not DECOMP_VIEW.Create(decompiled, ea):
                print("GhIDA:: [!] Error creating the view")
                return
            DECOMP_VIEW.Show()

        register_actions_and_handlers_decompile_view()

        # Add comments
        comment_list = COMMENTS_CACHE.get_comments_cache(ea)
        if comment_list:
            DECOMP_VIEW.add_comments(comment_list)

        return
    except Exception as e:
        print("GhIDA:: [!] Decompilation wrapper error: {}".format(e))
        idaapi.warning("GhIDA decompilation wrapper error {}".format(e))


# ------------------------------------------------------------
#   GHIDRA DECOMPILER PLUGIN
# ------------------------------------------------------------

class GhIDADecompiler_t(idaapi.plugin_t):
    comment = "GhIDA Decompiler for IDA Pro"
    help = "GhIDA Decompiler shortcut key is Ctrl-Alt-D"
    wanted_name = "GhIDA Decompiler"
    wanted_hotkey = "Ctrl-Alt-D"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        # Print header
        print("=" * 60)
        print("GhIDA Decompiler v{0}".format(gl.ghida_vv))
        print("Andrea Marcelli <anmarcel@cisco.com>")
        print("Cisco Talos, June 2019")
        print("GhIDA Decompiler shortcut key is Ctrl-Alt-D")
        print("=" * 60)

        self.__uihooks = None
        self.__seh = None

        try:
            import pygments
        except Exception:
            print("GhIDA:: [!] pygments library is missing")
            print("pip2 install pygments")
            return idaapi.PLUGIN_SKIP

        try:
            import requests
        except Exception:
            print("GhIDA:: [!] requests library is missing")
            print("pip2 install requests")
            return idaapi.PLUGIN_SKIP

        load_configuration()
        register_handlers()

        # Avoid displaying Running python script dialog
        # Otherwise, it breaks the UI and Cancel button
        idaapi.disable_script_timeout()

        # Hooking
        self.__uihooks = DisasmsHooks()
        self.__uihooks.hook()

        self.__seh = ScreenEAHook()
        self.__seh.hook()
        return idaapi.PLUGIN_KEEP

    def term(self):
        # Unhook
        if self.__uihooks:
            self.__uihooks.unhook()
        if self.__seh:
            self.__seh.unhook()

        # Remove temporary files
        gl.ghida_finalize(GHIDA_CONF.use_ghidra_server,
                          GHIDA_CONF.ghidra_server_url)

        # Dump the cache to file
        if GHIDA_CONF.load_save_cached_code:
            DECOMPILED_CACHE.dump_cache_to_json()
        if GHIDA_CONF.load_save_cached_comments:
            COMMENTS_CACHE.dump_cache_to_json()

    def run(self, arg):
        # Decompile the function
        decompile_function_wrapper()


def PLUGIN_ENTRY():
    # Check IDA version
    if not gl.is_ida_version_supported:
        return

    return GhIDADecompiler_t()


repository: 19h/ida-plugins
file: rcomdump/xcomdump.py
wanted_name: XCOM Call Graph Generator
comment: Generates a Graphviz DOT file showing the call graph from the current function.
created_at: 2025-04-08 22:54:27+00:00
pushed_at: 2025-04-08 23:07:17+00:00
url:
https://github.com/19h/ida-plugins/blob/b045b6f6bc9ae2a53104986fb74f3fe1198e5d2e/rcomdump/xcomdump.py

content:
import idaapi
import idautils
import idc
import ida_kernwin

def generate_call_graph_graphviz():
    """
    Performs a breadth-first search starting from the current function to identify all
    reachable functions in the call graph. Then generates a Graphviz DOT file
    representing the call graph structure. Each node will have a maximum of 2 incoming edges
    from the nodes closest to the root.
    """
    output_file = ida_kernwin.ask_file(False, "*.dot", "Select output file for call graph")
    if not output_file:
        print("No output file selected.")
        return

    # Get the current effective address and the function containing it
    ea = idaapi.get_screen_ea()
    func = idaapi.get_func(ea)

    if not func:
        print("No function at the current address.")
        return

    # Initialize data structures
    start_ea = func.start_ea
    visited = set()       # Set of reachable function start addresses
    callers = {}          # Dictionary mapping function EA to set of caller EAs
    distances = {start_ea: 0}  # Dictionary mapping function EA to distance from root

    # Perform BFS to build the call graph with shortest paths
    bfs(start_ea, visited, callers, distances)

    try:
        with open(output_file, "w") as f:
            # Write Graphviz header with strict digraph to avoid duplicate edges
            f.write("strict digraph CallGraph {\n")
            f.write("    node [shape=box, style=filled, fillcolor=lightblue];\n")
            f.write("    rankdir=LR;\n\n")
            
            # Write nodes with function names as labels
            f.write("    // Nodes\n")
            for ea in visited:
                func_name = idc.get_func_name(ea) or f"func_0x{ea:X}"
                safe_name = f"func_{ea}"  # Use address as node ID to avoid invalid chars
                f.write(f'    {safe_name} [label="{func_name}"];\n')
            
            f.write("\n    // Edges\n")
            # Write edges - max 2 edges from the closest callers for each function
            for callee_ea, caller_set in callers.items():
                if callee_ea == start_ea:  # Skip the root node
                    continue
                    
                callee_name = f"func_{callee_ea}"
                
                # Find up to 2 callers with the shortest distance to root
                valid_callers = [caller for caller in caller_set if caller in distances]
                closest_callers = sorted(valid_callers, key=lambda x: distances.get(x, float('inf')))[:4]
                
                # Add edges from the closest callers (up to 2)
                for caller_ea in closest_callers:
                    caller_name = f"func_{caller_ea}"
                    f.write(f"    {caller_name} -> {callee_name};\n")
            
            # Close the graph
            f.write("}\n")

        print(f"[*] Call graph written to {output_file}")
    except Exception as e:
        print(f"Error writing output file: {e}")


def bfs(start_ea, visited, callers, distances):
    """
    Perform breadth-first search to find shortest paths from root to all reachable functions.
    """
    queue = [start_ea]
    visited.add(start_ea)
    
    while queue:
        func_ea = queue.pop(0)
        current_distance = distances[func_ea]
        func = idaapi.get_func(func_ea)
        if not func:
            continue

        # Iterate over the function's basic blocks and instructions
        for block in idaapi.FlowChart(func):
            for head in idautils.Heads(block.start_ea, block.end_ea):
                if idaapi.is_call_insn(head):
                    for ref in idautils.CodeRefsFrom(head, 0):
                        called_func = idaapi.get_func(ref)
                        # Ensure the reference is the start of a function
                        if called_func and called_func.start_ea == ref:
                            callee_ea = ref
                            # Initialize callers set for the callee if not present
                            if callee_ea not in callers:
                                callers[callee_ea] = set()
                            # Add the current function as a caller
                            callers[callee_ea].add(func_ea)
                            
                            # If we haven't visited this function yet, add it to the queue
                            if callee_ea not in visited:
                                visited.add(callee_ea)
                                distances[callee_ea] = current_distance + 1
                                queue.append(callee_ea)


class XCOMCallGraphPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "Generates a Graphviz DOT file showing the call graph from the current function."
    help = ("This plugin performs a breadth-first search from the current function, "
            "identifies all reachable functions, and generates a call graph in Graphviz DOT format.")
    wanted_name = "XCOM Call Graph Generator"
    wanted_hotkey = ""

    def init(self):
        print("[*] XCOM Call Graph Generator plugin initialized.")
        return idaapi.PLUGIN_OK

    def run(self, arg):
        generate_call_graph_graphviz()

    def term(self):
        print("[*] XCOM Call Graph Generator plugin terminated.")


def PLUGIN_ENTRY():
    return XCOMCallGraphPlugin()


repository: The-Itach1/Audit
file: MipsAudit/MipsAudit.py
description: 使用idapython写的一个辅助二进制固件漏洞挖掘的代码审计脚本
wanted_name: MipsAudit
comment: MIPS Audit plugin for IDA Pro
created_at: 2022-08-03 14:54:03+00:00
pushed_at: 2022-08-03 15:22:14+00:00
url:
https://github.com/The-Itach1/Audit/blob/10691236ddff1ad923e0b9f365c66f8bdef76743/MipsAudit/MipsAudit.py

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

# reference
# 《ida pro 权威指南》
# 《python 灰帽子》
# 《家用路由器0day漏洞挖掘》
# https://github.com/wangzery/SearchOverflow/blob/master/SearchOverflow.py

from idaapi import *
import idaapi
import idc
from prettytable import PrettyTable

if idaapi.IDA_SDK_VERSION > 700:
    import ida_search
    from idc import (
        print_operand
    )
    from ida_bytes import (
        get_strlit_contents
    )
else:
    from idc import (
        GetOpnd as print_operand,
        GetString
    )
    def get_strlit_contents(*args): return GetString(args[0])

DEBUG = True

# fgetc,fgets,fread,fprintf,
# vspritnf

# set function_name
dangerous_functions = [
    "strcpy",
    "strcat",
    "sprintf",
    "read",
    "getenv"
]

attention_function = [
    "memcpy",
    "malloc",
    "strncpy",
    "sscanf",
    "strncat",
    "snprintf",
    "vprintf",
    "printf"
]

command_execution_function = [
    "system",
    "execve",
    "popen",
    "unlink"
]

# describe arg num of function

one_arg_function = [
    "getenv",
    "system",
    "unlink"
]

two_arg_function = [
    "strcpy", 
    "strcat",
    "popen"
]

three_arg_function = [
    "strncpy",
    "strncat", 
    "memcpy",
    "execve",
    "read"
]

format_function_offset_dict = {
    "sprintf":1,
    "sscanf":1,
    "snprintf":2,
    "vprintf":0,
    "printf":0
}


def printFunc(func_name):
    string1 = "========================================"
    string2 = "========== Auditing " + func_name + " "
    strlen = len(string1) - len(string2)
    return string1 + "\n" + string2 + '=' * strlen + "\n" + string1

def getFuncAddr(func_name):
    func_addr = idc.get_name_ea_simple(func_name)
    if func_addr != BADADDR:
        print(printFunc(func_name))
        return func_addr
    return False

def getFormatString(addr):
    op_num = 1
    # idc.get_operand_type Return value
    #define o_void        0  // No Operand                           ----------
    #define o_reg         1  // General Register (al, ax, es, ds...) reg
    #define o_mem         2  // Direct Memory Reference  (DATA)      addr
    #define o_phrase      3  // Memory Ref [Base Reg + Index Reg]    phrase
    #define o_displ       4  // Memory Reg [Base Reg + Index Reg + Displacement] phrase+addr
    #define o_imm         5  // Immediate Value                      value
    #define o_far         6  // Immediate Far Address  (CODE)        addr
    #define o_near        7  // Immediate Near Address (CODE)        addr
    #define o_idpspec0    8  // IDP specific type
    #define o_idpspec1    9  // IDP specific type
    #define o_idpspec2   10  // IDP specific type
    #define o_idpspec3   11  // IDP specific type
    #define o_idpspec4   12  // IDP specific type
    #define o_idpspec5   13  // IDP specific type
    # 如果第二个不是立即数则下一个

    #print(hex(addr))
    if(idc.get_operand_type(addr ,op_num) != 5):
        op_num = op_num + 1
    if idc.get_operand_type(addr ,op_num) != 5:
        return "get fail"

    op_string = print_operand(addr, op_num).split(" ")[0].split("+")[0].split("-")[0].replace("(", "")
    string_addr = idc.get_name_ea_simple(op_string)
    if string_addr == BADADDR:
        return "get fail"
    string = str(get_strlit_contents(string_addr, -1, STRTYPE_TERMCHR))
    return [string_addr, string ,op_string]


def get_Tpye_ArgAddr(start_addr, regNum, argType):
    mipscondition = ["bn", "be" , "bg", "bl"]
    scan_deep = 80
    count = 0
    reg = argType + str(regNum)

    # try to get before
    before_addr = get_first_cref_to(start_addr)
    while before_addr != BADADDR:
        if reg == print_operand(before_addr, 0):
            Mnemonics = print_insn_mnem(before_addr)
            if Mnemonics[0:2] in mipscondition:
                pass
            elif Mnemonics[0:1] == "j":
                pass
            else:
                return before_addr
        count = count + 1
        if count > scan_deep:
            break
        before_addr = get_first_cref_to(before_addr)
    return BADADDR


def getArg(start_addr, regNum):
    mipsmov = ["move", "lw", "li", "lb", "lui", "lhu", "lbu", "la"]
    arg_addr = get_Tpye_ArgAddr(start_addr, regNum, "$a")
    if arg_addr != BADADDR:
        Mnemonics = print_insn_mnem(arg_addr) 
        if Mnemonics[0:3] == "add":
            if print_operand(arg_addr, 2) == "":
                arg = print_operand(arg_addr, 0) + "+" + print_operand(arg_addr, 1)
            else:
                arg = print_operand(arg_addr, 1) + "+" +  print_operand(arg_addr, 2)
        elif Mnemonics[0:3] == "sub":
            if print_operand(arg_addr, 2) == "":
                arg = print_operand(arg_addr, 0) + "-" + print_operand(arg_addr, 1)
            else:
                arg = print_operand(arg_addr, 1) + "-" +  print_operand(arg_addr, 2)
        elif Mnemonics in mipsmov:
            arg = print_operand(arg_addr, 1) 
        else:
            arg = idc.generate_disasm_line(arg_addr,1).split("#")[0]
        set_cmt(arg_addr, "addr: 0x%x " % start_addr  + "-------> arg" + str((int(regNum)+1)) + " : " + arg, 0)
        return arg
    else:
        return "get fail"

def audit(func_name):
    func_addr = getFuncAddr(func_name)  
    if func_addr == False:
        return False

    # get arg num and set table
    if func_name in one_arg_function:
        arg_num = 1
    elif func_name in two_arg_function:
        arg_num = 2
    elif func_name in three_arg_function:
        arg_num = 3
    elif func_name in format_function_offset_dict:
        arg_num = format_function_offset_dict[func_name] + 1
    else:
        print("The %s function didn't write in the describe arg num of function array,please add it to,such as add to `two_arg_function` arary" % func_name)
        return
    table_head = ["func_name", "addr"]
    for num in range(0,arg_num):
        table_head.append("arg"+str(num+1))
    if func_name in format_function_offset_dict:
        table_head.append("format&value[string_addr, num of '%', fmt_arg...]")
    table_head.append("local_buf_size")
    table = PrettyTable(table_head)

    # get first call
    call_addr = get_first_cref_to(func_addr)
    while call_addr != BADADDR:
        idc.set_color(call_addr, idc.CIC_ITEM, 0xffff00)
        Mnemonics = print_insn_mnem(call_addr)
        if Mnemonics[0:1] == "j" or Mnemonics[0:1] == "b":
            if func_name in format_function_offset_dict:
                info = auditFormat(call_addr, func_name, arg_num)
            else:
                info = auditAddr(call_addr, func_name, arg_num)
            table.add_row(info)
        call_addr = get_next_cref_to(func_addr, call_addr)
    print(table)

def auditAddr(call_addr, func_name, arg_num):
    addr = "0x%x" % call_addr
    ret_list = [func_name, addr]
    # local buf size
    local_buf_size = idc.get_func_attr(call_addr , idc.FUNCATTR_FRSIZE)
    if local_buf_size == BADADDR :
        print("debug 236")
        local_buf_size = "get fail"
    else:
        local_buf_size = "0x%x" % local_buf_size
    # get arg
    for num in range(0,arg_num):
        ret_list.append(getArg(call_addr, num)) 
    ret_list.append(local_buf_size)
    return ret_list

def auditFormat(call_addr, func_name, arg_num):
    addr = "0x%x" % call_addr
    ret_list = [func_name, addr]
    # local buf size
    local_buf_size = idc.get_func_attr(call_addr , idc.FUNCATTR_FRSIZE)
    if local_buf_size == BADADDR :
        print("debug 252")
        local_buf_size = "get fail"
    else:
        local_buf_size = "0x%x" % local_buf_size
    # get arg
    for num in range(0,arg_num):
        ret_list.append(getArg(call_addr, num)) 
    mid_arg_addr = get_Tpye_ArgAddr(call_addr, format_function_offset_dict[func_name], "$a")
    xx_arg=getArg(call_addr,format_function_offset_dict[func_name])

    if xx_arg in ['$v0','$v1','$v2']:
        true_arg_addr = get_Tpye_ArgAddr(mid_arg_addr, int(xx_arg[-1]),"$v")
    elif xx_arg in ['$s1','$s0']:
        true_arg_addr = get_Tpye_ArgAddr(mid_arg_addr, int(xx_arg[-1]), "$s")
    elif xx_arg in ["$a0","$a1","$a2"] and int(xx_arg[-1])!= format_function_offset_dict[func_name]:
        true_arg_addr = get_Tpye_ArgAddr(mid_arg_addr, int(xx_arg[-1]), "$a")
    else:
        true_arg_addr=mid_arg_addr
    string_and_addr =  getFormatString(true_arg_addr)

    format_arg=string_and_addr[2]
    ret_list.pop(-1)
    ret_list.append(format_arg)


    format_and_value = []
    if string_and_addr == "get fail":
        ret_list.append("get fail")
    else:
        string_addr = "0x%x" % string_and_addr[0]
        format_and_value.append(string_addr)
        string = string_and_addr[1]
        fmt_num = string.count("%")
        format_and_value.append(fmt_num)
        # mips arg reg is from a0 to a3
        if fmt_num + arg_num >= 4:
            fmt_num = 4
        else:
            fmt_num = fmt_num + arg_num
        for num in range(arg_num, fmt_num):
            format_and_value.append(getArg(call_addr, num))
        ret_list.append(format_and_value)
    ret_list.append(local_buf_size)
    return ret_list

def MipsAudit():
    # the word create with figlet
    print("Auditing dangerous functions ......")
    for func_name in dangerous_functions:
        audit(func_name)
        
    print("Auditing attention function ......")
    for func_name in attention_function:
        audit(func_name)

    print("Auditing command execution function ......")
    for func_name in command_execution_function:
        audit(func_name)
        
    print("Finished! Enjoy the result ~")

m_initialized = False

class MipsAudit_Plugin_t(idaapi.plugin_t):
    comment = "MIPS Audit plugin for IDA Pro"
    help = "todo"
    wanted_name = "MipsAudit"
    wanted_hotkey = "Ctrl-Alt-M"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        global m_initialized

        if m_initialized is False:
            m_initialized = True
            print("=" * 80)
            start = '''
             __  __ _              _             _ _ _   
            |  \/  (_)_ __  ___   / \  _   _  __| (_) |_ 
            | |\/| | | '_ \/ __| / _ \| | | |/ _` | | __|
            | |  | | | |_) \__ \/ ___ \ |_| | (_| | | |_ 
            |_|  |_|_| .__/|___/_/   \_\__,_|\__,_|_|\__|
                     |_|                                 
                     
                            re-edit by The_Itach1  2022.8.3
            '''
            print(start)
            print("=" * 80)

            return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        info = idaapi.get_inf_structure()
        if 'mips' in info.procName:
            MipsAudit()
        else:
            print('MipsAudit is not supported on the current arch')

def PLUGIN_ENTRY():
    return MipsAudit_Plugin_t()

if __name__ == '__main__':
    MipsAudit()


repository: clovme/WTools
file: software/Reverse Cracking/IDA Professional 9.0/python/examples/merge/py_mex1.py
wanted_name: IDAPython: Merge example 1
comment: IDAPython: An example 1 how to implement IDA merge functionality
created_at: 2025-03-27 00:14:54+00:00
pushed_at: 2025-05-11 10:18:21+00:00
url:
https://github.com/clovme/WTools/blob/f12f01c833a6a6241799a9d01e9117cc28a65dbc/software/Reverse%20Cracking/IDA%20Professional%209.0/python/examples/merge/py_mex1.py

content:
# pylint: disable=line-too-long,invalid-name,import-error

"""
summary: add merge functionality to a simple plugin, example 1

description:
  This is a primitive plugin which asks user for some info and saves it for
  some addresses.

  We will add a merge functionality to plugin.

  An IDA plugin may have two kinds of data with permanent storage:
    1. Data common for entire database (e.g. the options).
       To describe them we will use the idbattr_info_t type.
    2. Data specific to a particular address.
       To describe them we will use the merge_node_info_t type.

  Also, see SDK/plugins/mex1 example
"""


import ida_idaapi
import ida_ida
import ida_kernwin
import ida_netnode
import ida_funcs
import ida_merge
import ida_mergemod
import ida_idp

# --------------------------------------------------------------------------
# netnode to store plugin data
MEX_NODE_NAME = "$ idapython mex1"
# user input
MEX_OPTION_FLAGS_IDX = ida_idaapi.BADADDR & -1   # atag
MEX_OPTION_IDENT_IDX = ida_idaapi.BADADDR & -2   # stag
# EA marks
MEX_EA_TAG = 'm'

# mex_ctx_t::flags bits
MEX_FLAGS_0 = 0x01
MEX_FLAGS_1 = 0x02

# --------------------------------------------------------------------------
class idp_listener_t(ida_idp.IDP_Hooks):
    """
    we need an event listener to catch processor_t::ev_create_merge_handlers
    """

    def __init__(self, ctx):
        # by default IDP_Hooks uses ida_idp.HKCB_GLOBAL hkcb_flags,
        # in that case IDP events are sent to all DB instances.
        # Such behaviour does not matter for IDA Pro
        # but IDA Teams needs to have only one IDP event sent to plugin.
        # We set hkcb_flags to 0
        ida_idp.IDP_Hooks.__init__(self, 0, 0)
        self.ctx = ctx

    def ev_ending_undo(self):
        """
        A well behaving plugin should restore its state from the database
        upon ev_ending_undo. Otherwise its state may be conflicting with the
        database.
        """
        self.ctx.restore_from_idb()
        return 0

    def ev_create_merge_handlers(self, md):
        """
        This event occurs when IDA is performing a 3-way merge (for IDA Teams)
        Our plugins should create and register merge handler(s) for its data.
        """
        self.ctx.create_merge_handlers(md)
        return 0


    def ev_cvt64_supval(self, node, tag, idx, data):
        """
        Converter to i64 database
        """
        nn = ida_netnode.netnode(MEX_NODE_NAME)
        if nn == node:
            if tag == ida_netnode.stag:
                nn.supset(MEX_OPTION_IDENT_IDX, data)
                return 1
            if tag == ida_netnode.atag and len(data):
                val = int.from_bytes(data, 'little')
                nn.altset(MEX_OPTION_FLAGS_IDX, val)
                return 1
            if chr(tag) == MEX_EA_TAG:
                nn.supset(idx, data, MEX_EA_TAG)
                return 1
        return 0


# --------------------------------------------------------------------------
class mex_ctx_t(ida_idaapi.plugmod_t):
    """
    Regular plugin implementation below.
    For example, in our case the plugin asks for 2 bit values and a string value.
    Then the plugin stores this data in the database.
    And mark the start address of the current function.
    These data will be merged later.
    """

    def __init__(self):
        # bit flags, see above MEX_FLAGS_0/MEX_FLAGS_1
        self.flags = 0
        # unique database ident
        self.ident = ""
        # Restore the plugin data from the database into the memory.
        self.restore_from_idb()
        # Hook an event listener, to catch the merge-related event(s).
        self.idp_listener = idp_listener_t(self)
        self.idp_listener.hook()
        # MERGE: the following data must exist during plugin lifetime
        self.modmerger_helper = None
        self.idpopts_info = None
        self.merge_node_info = None

    def save_to_idb(self):
        """ Save the plugin state to the idb. """
        nn = ida_netnode.netnode(MEX_NODE_NAME, 0, True)
        nn.altset(MEX_OPTION_FLAGS_IDX, self.flags)
        nn.supset(MEX_OPTION_IDENT_IDX, self.ident)

    def restore_from_idb(self):
        """ Restore plugin variables from the idb. """
        nn = ida_netnode.netnode(MEX_NODE_NAME)
        if nn != ida_netnode.BADNODE:
            self.flags = nn.altval(MEX_OPTION_FLAGS_IDX)
            self.ident = nn.supstr(MEX_OPTION_IDENT_IDX)

    def run(self, _):
        """
        Ask user for the data and save them to database.
        Add mark for current EA.
        """
        if self._ask_form():
            self.save_to_idb()
        # Our plugin stores a string for the current function.
        # Just for illustration purposes of how plugins should merge address-specific info
        # stored in a netnode.
        ea = ida_kernwin.get_screen_ea()
        pfn = ida_funcs.get_func(ea)
        if pfn:
            one = ""
            if (self.flags & MEX_FLAGS_0) != 0:
                one = " one"
            two = ""
            if (self.flags & MEX_FLAGS_1) != 0:
                one = " two"
            mark = "IPMEX1" + one + two
            nn = ida_netnode.netnode(MEX_NODE_NAME, 0, True)
            nn.supset_ea(pfn.start_ea, mark, MEX_EA_TAG)

    def _ask_form(self):
        class MexForm(ida_kernwin.Form):
            def __init__(self):
                ida_kernwin.Form.__init__(self,
                r"""
IDAPython: merge example 1

<Flag 0:{optFlag0}>
<Flag 1:{optFlag1}>{grpFlags}>
<Ident prefix:{ident}>
                """,
                {
                    "grpFlags": ida_kernwin.Form.ChkGroupControl(("optFlag0", "optFlag1")),
                    "ident"   : ida_kernwin.Form.StringInput(swidth=10),
                })

        form = MexForm()
        form, _ = form.Compile()
        form.grpFlags.value = self.flags
        form.ident.value = self.ident
        ok = form.Execute()
        if ok == 1:
            self.flags = form.grpFlags.value
            self.ident = form.ident.value
        form.Free()
        return ok == 1

    def create_merge_handlers(self, md):
        """
        Create merge handlers for plugin
        """

        # 1. Data common for entire database (e.g. the options).
        #
        # This example shows how to merge the data from database.
        # We will describe the items to merge and pass the description
        # to create_std_modmerge_handlers(), which will do all the work for us.
        sizeof_flags = ida_netnode.SIZEOF_nodeidx_t
        self.idpopts_info = [
            # Describe both flags
            ida_ida.idbattr_info_t("MEX flag 0", MEX_OPTION_FLAGS_IDX, sizeof_flags, MEX_FLAGS_0, ida_netnode.atag, ida_ida.IDI_ALTVAL|ida_ida.IDI_SCALAR),
            ida_ida.idbattr_info_t("MEX flag 1", MEX_OPTION_FLAGS_IDX, sizeof_flags, MEX_FLAGS_1, ida_netnode.atag, ida_ida.IDI_ALTVAL|ida_ida.IDI_SCALAR),
            # Describe ident
            ida_ida.idbattr_info_t("MEX ident",  MEX_OPTION_IDENT_IDX, 0,            0,           ida_netnode.stag, ida_ida.IDI_SUPVAL|ida_ida.IDI_CSTR),
        ]

        # The descriptions are ready. Now create an instance of the standard helper
        # class to be passed to the kernel, and the kernel will take care of organizing
        # the merge process for them.

        # helper instance name
        self.modmerger_helper = ida_merge.moddata_diff_helper_t(
            "Sample merge data", # label: prefix for the attribute names, e.g. "Sample merge data.MEX flag 0"
            MEX_NODE_NAME,       # netnode name for idpopts_info and merge_node_info
            self.idpopts_info)   # field descriptions

        # 2. Data specific to a particular address.
        # We describe how the data is stored in a netnode.
        self.merge_node_info = [
            ida_merge.merge_node_info_t(
                "Function marks",       # label of the merge handler, e.g. "Plugins/Merge example 1/Function marks"
                MEX_EA_TAG,             # netnode tag
                ida_merge.NDS_MAP_IDX|ida_merge.NDS_IS_STR
                                        # netnode value descriptors and modificators, see \ref nds_flags_t
            ),
        ]

        # Now we should combine together in function create_std_modmerge_handlers.
        # This function will be called on processor_t::ev_create_merge_handlers event.
        # As a result, two merge handlers with labels
        #   "Plugins/Merge example 1/Database attributes"
        #   "Plugins/Merge example 1/Function marks"
        # will be created.
        mhp = ida_merge.merge_handler_params_t(
            md,                                  # merge handler data
            "Plugins/IDAPython merge example 1", # Label of the merge handler
            ida_merge.MERGE_KIND_NONE,           # allocate a merge kind
            ida_merge.MERGE_KIND_END,            # insert to the end of handler list
            0)                                   # merge handler flags
        ida_mergemod.create_std_modmerge_handlers(
             mhp,
             self.modmerger_helper,
             self.merge_node_info)


# --------------------------------------------------------------------------
class mex1_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_MULTI | ida_idaapi.PLUGIN_MOD
    wanted_name = "IDAPython: Merge example 1"
    comment = "IDAPython: An example 1 how to implement IDA merge functionality"
    wanted_hotkey = ""
    help = ""
    def init(self):
        return mex_ctx_t()
    def term(self):
        pass
    def run(self, arg):
        pass
def PLUGIN_ENTRY():
    return mex1_plugin_t()


repository: alphanin9/CyberpunkReversing
file: address-helper.py
comment: Get address hash of function/vtable
created_at: 2024-12-11 03:52:17+00:00
pushed_at: 2025-02-18 23:21:09+00:00
description:
Various reverse engineering stuff for Cyberpunk 2077 - because hosting it is more convenient here than on Discord

url:
https://github.com/alphanin9/CyberpunkReversing/blob/a83031fda772dfd769c2800095c003f8b8807ee8/address-helper.py

content:
VERSION = '1.0.0'
__AUTHOR__ = "not_alphanine"

PLUGIN_NAME = "Get Cyberpunk 2077 address hashes"

import os
import sys
import idc
import idaapi
import idautils
import ida_funcs
import ida_name
import ida_ua
import ida_segment

import json

import ida_ida
UA_MAXOP=ida_ida.UA_MAXOP

major, minor = map(int, idaapi.get_kernel_version().split("."))
using_ida7api = (major > 6)
using_pyqt5 = using_ida7api or (major == 6 and minor >= 9)

idaver_74newer = (major == 7 and minor >= 4)
idaver_8newer = (major >= 8)

if idaver_74newer or idaver_8newer:
    newer_version_compatible = True
else:
    newer_version_compatible = False

if newer_version_compatible:
    #IDA 7.4+
    #https://hex-rays.com/products/ida/support/ida74_idapython_no_bc695_porting_guide.shtml
    import ida_ida
    import ida_kernwin

if using_pyqt5:
    import PyQt5.QtGui as QtGui
    import PyQt5.QtCore as QtCore
    import PyQt5.QtWidgets as QtWidgets
    from PyQt5.Qt import QApplication

else:
    import PySide.QtGui as QtGui
    import PySide.QtCore as QtCore
    QtWidgets = QtGui
    QtCore.pyqtSignal = QtCore.Signal
    QtCore.pyqtSlot = QtCore.Slot
    from PySide.QtGui import QApplication

def set_clipboard_text(data):
    cb = QApplication.clipboard()
    cb.clear(mode=cb.Clipboard )
    cb.setText(data, mode=cb.Clipboard)

class CP2077AddressHelper(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Get address hash of function/vtable"
    help = "Select whatever line you want (with functions) or the vtable, right click and press \"Lookup hash\""
    wanted_name = PLUGIN_NAME
    
    ACTION_LOOKUP_HASH = "prefix:cp2077_lookup_hash"
    ACTION_JUMP_TO_HASH = "prefix:cp2077_jump_to_hash"

    def init_action_lookup_hash(self):
        if idaapi.unregister_action(self.ACTION_LOOKUP_HASH):
            idaapi.msg("Warning: action was already registered, unregistering it first\n")
        
        if idaapi.unregister_action(self.ACTION_JUMP_TO_HASH):
            idaapi.msg("Warning: action was already registered, unregistering it first\n")

        if (sys.version_info > (3, 0)):
            # Describe the action using python3 copy
            action_desc = idaapi.action_desc_t(
                self.ACTION_LOOKUP_HASH,                                                    # The action name.
                "Lookup CP2077 address hash",                                               # The action text.
                IDACtxEntry(lookup_current_address_hash),                                   # The action handler.
            )

            action_jump_desc = idaapi.action_desc_t(
                self.ACTION_JUMP_TO_HASH,
                "Jump to CP2077 addr hash",
                IDACtxEntry(goto_hash)
            )
        else:
            # Describe the action using python2 copy
            action_desc = idaapi.action_desc_t(
                self.ACTION_LOOKUP_HASH,                                        # The action name.
                "Lookup CP2077 address hash",                                   # The action text.
                IDACtxEntry(lookup_current_address_hash),                      # The action handler.
            )

            action_jump_desc = idaapi.action_desc_t(
                self.ACTION_JUMP_TO_HASH,
                "Jump to CP2077 addr hash",
                IDACtxEntry(goto_hash)
            )

        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"
        assert idaapi.register_action(action_jump_desc), "Action 2 registration failed"

    def del_action_lookup_hash(self):
        idaapi.unregister_action(self.ACTION_LOOKUP_HASH)
        idaapi.unregister_action(self.ACTION_JUMP_TO_HASH)

    def init(self):
        if idc.get_root_filename() != "Cyberpunk2077.exe":
            # This is only for CP2077
            return idaapi.PLUGIN_SKIP

        self.init_action_lookup_hash()
        self._init_hooks()

        idaapi.msg("%s started" % self.wanted_name)

        return idaapi.PLUGIN_KEEP
    
    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.msg("%s cannot be run as a script.\n" % self.wanted_name)

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

        if hasattr(self, '_hooks'):
            self._hooks.unhook()
        self.del_action_lookup_hash()

        # done
        idaapi.msg("%s terminated...\n" % self.wanted_name)

    def _init_hooks(self):
        self._hooks = PluginHooks()
        self._hooks.ready_to_run = self._init_hexrays_hooks
        self._hooks.hook()

    def _init_hexrays_hooks(self):
        """
        Install Hex-Rrays hooks (when available).
        NOTE: This is called when the ui_ready_to_run event fires.
        """
        if idaapi.init_hexrays_plugin():
            idaapi.install_hexrays_callback(self._hooks.hxe_callback)

    

class PluginHooks(idaapi.UI_Hooks):
    def __init__(self):
        # Call the __init__ method of the superclass
        super(PluginHooks, self).__init__()

        # Get the IDA version
        major, minor = map(int, idaapi.get_kernel_version().split("."))
        self.newer_version_compatible = (major == 7 and minor >= 4) or (major >= 8)
        
        # If the IDA version is less than 7.4, define finish_populating_tform_popup
        if not self.newer_version_compatible:
            self.finish_populating_tform_popup = self._finish_populating_tform_popup

    def finish_populating_widget_popup(self, widget, popup_handle, ctx=None):
        """
        A right click menu is about to be shown. (IDA 7.x)
        """
        inject_lookup_address_hash(widget, popup_handle, idaapi.get_widget_type(widget))
        return 0


    def _finish_populating_tform_popup(self, form, popup):
        """
        A right click menu is about to be shown. (IDA 6.x)
        """
        inject_lookup_address_hash(form, popup, idaapi.get_tform_type(form))
        return 0
    
    def hxe_callback(self, event, *args):
        """
        HexRays event callback.
        """

        #
        # if the event callback indicates that this is a popup menu event
        # (in the hexrays window), we may want to install our prefix menu
        # actions depending on what the cursor right clicked.
        #

        if event == idaapi.hxe_populating_popup:
            form, popup, vu = args

            idaapi.attach_action_to_popup(
                form,
                popup,
                CP2077AddressHelper.ACTION_LOOKUP_HASH,
                "Lookup CP2077 address hash",
                idaapi.SETMENU_APP
            )

            idaapi.attach_action_to_popup(
                form,
                popup,
                CP2077AddressHelper.ACTION_JUMP_TO_HASH,
                "Jump to CP2077 address hash",
                idaapi.SETMENU_APP
            )

        # done
        return 0
    
def inject_lookup_address_hash(widget, popup_handle, widget_type):
    disasm_type = None
    if major > 8:
        disasm_type = idaapi.BWN_DISASM
    else:
        disasm_type = idaapi.BWN_DISASMS

    if widget_type == disasm_type or widget_type == idaapi.BWN_PSEUDOCODE:
        idaapi.attach_action_to_popup(
            widget,
            popup_handle,
            CP2077AddressHelper.ACTION_LOOKUP_HASH,
            "Lookup CP2077 address hash",
            idaapi.SETMENU_APP
        )

        idaapi.attach_action_to_popup(
            widget,
            popup_handle,
            CP2077AddressHelper.ACTION_JUMP_TO_HASH,
            "Jump to CP2077 address hash",
            idaapi.SETMENU_APP
        )
    return 0

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

cp2077_addresses_cache = None

def build_address_cache():
    global cp2077_addresses_cache

    # Build address cache first...
    if cp2077_addresses_cache is None:
        print("Building addr cache...")
        current_path = os.path.split(idc.get_idb_path())
        address_path = os.path.join(current_path[0], "cyberpunk2077_addresses.json")

        print("File is at ", current_path)

        cp2077_addresses_cache = {}

        addr_to_hash = {}
        hash_to_addr = {}

        if not os.path.exists(address_path):
            print("Failed to find CP2077 addresses!")
            return
        
        with open(address_path) as file:
            base_of_code = ida_segment.get_segm_by_name(".text").start_ea
            base_of_data = ida_segment.get_segm_by_name(".data").start_ea
            # Fix: IDA-generated imports segment messing things up
            # Could also be fixed by parsing PE ourselves, but... why?
            seg_rdata = ida_segment.get_segm_by_name(".idata")
            if not seg_rdata:
                seg_rdata = ida_segment.get_segm_by_name(".rdata")
            base_of_rdata = seg_rdata.start_ea

            # 1 = code
            # 2 = rdata
            # 3 = data

            json_addresses = json.load(file)["Addresses"]
            addr_count = 0
            for addr in json_addresses:
                offset = addr["offset"].split(":")

                # 1st: seg
                # 2nd: addr

                addrValue = int(offset[1], 16)
                
                final_addr = addrValue

                if offset[0] == "0001":
                    final_addr = addrValue + base_of_code
                elif offset[0] == "0002":
                    final_addr = addrValue + base_of_rdata
                elif offset[0] == "0003":
                    final_addr = addrValue + base_of_data

                # Maybe make a pass that renames autonamed functions with defined syms to their names in the .json...

                if "symbol" in addr:
                    func_start = idc.get_func_attr(final_addr, idc.FUNCATTR_START)
                    func_name = idc.get_func_name(func_start)

                    if func_name.startswith("sub_"):
                        # Automatically named, rename it?
                        idaapi.set_name(func_start, addr["symbol"], idaapi.SN_FORCE)
                    else:
                        print(f'{addr["symbol"]}@{hex(final_addr)} = {func_name}')

                if not (final_addr in addr_to_hash):
                    addr_to_hash[final_addr] = []

                addr_to_hash[final_addr].append(addr["hash"])
                
                hash_to_addr[int(addr["hash"])] = final_addr

                addr_count += 1
            print(f"Loaded {addr_count} addresses into map")

            cp2077_addresses_cache["hash_to_addr"] = hash_to_addr
            cp2077_addresses_cache["addr_to_hash"] = addr_to_hash

def get_hash_for_address(addr):
    global cp2077_addresses_cache
    build_address_cache()

    if addr in cp2077_addresses_cache["addr_to_hash"]:
        return cp2077_addresses_cache["addr_to_hash"][addr]
    
    return None

def get_addr_for_hash(hash):
    global cp2077_addresses_cache
    build_address_cache()

    if hash in cp2077_addresses_cache["hash_to_addr"]:
        return cp2077_addresses_cache["hash_to_addr"][hash]

    return None


def lookup_current_address_hash():
    """
    If we're reversing CP2077, tries to get the address hash of the current fn and put it in clipboard/output
    """

    screen_address = idc.get_screen_ea()

    if screen_address != idaapi.BADADDR:
        addr_to_look_for = screen_address
        func_start = idc.get_func_attr(screen_address, idc.FUNCATTR_START)
        if func_start != idaapi.BADADDR:
            addr_to_look_for = func_start

        addr_hash = get_hash_for_address(addr_to_look_for)

        if addr_hash:
            if len(addr_hash) > 1:
                print("Multiple hashes!")
                for i in range(len(addr_hash)):
                    print(i, ": ", addr_hash[i])
                set_clipboard_text(addr_hash[0])
            else:
                print("Hash:", addr_hash[0])
                set_clipboard_text(addr_hash[0])
        else:
            print("Failed to find address in cache!")

def goto_hash():
    hash_str = None
    if newer_version_compatible:
        hash_str = ida_kernwin.ask_str("0", 0, "Hash to jump towards")    
    else:
        hash_str = idc.AskStr("0", "Hash to jump towards")

    if not hash_str:
        return
    
    hash_int = int(hash_str, 0)

    addr = get_addr_for_hash(hash_int)

    if not addr:
        print("Failed to find addy for hash!")
        return
    
    if newer_version_compatible:
        ida_kernwin.jumpto(addr)
    else:
        idc.Jump(addr)
        
class IDACtxEntry(idaapi.action_handler_t):
    def __init__(self, action_function):
        idaapi.action_handler_t.__init__(self)
        self.action_function = action_function

    def activate(self, ctx):
        self.action_function()
        return 1

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


repository: SapphireServer/ffxiv-ida-scripts
file: xiv/ffxiv-exdgetters.py
description: Useful scripts for the lazy developer
wanted_name: FFXIV - Annotate EXD Getters
comment: Automagically names EXD getter funcs
created_at: 2019-02-26 09:49:29+00:00
pushed_at: 2021-03-27 23:41:13+00:00
url:
https://github.com/SapphireServer/ffxiv-ida-scripts/blob/de97b09ae78d98998adbe98ccddac6092d0724d2/xiv/ffxiv-exdgetters.py

content:
# ffxiv-exdgetters.py
#
# Automagically labels most exd getter functions along with a hint indicating which sheet/sheet id its fetching from
#

import idaapi

# nb: "pattern": "func suffix" OR None
exd_func_patterns = {
    "48 83 EC 28 48 8B 05 ? ? ? ? 44 8B C1 BA ? ? ? ? 48 8B 88 ? ? ? ? E8 ? ? ? ? 48 85 C0 75 05 48 83 C4 28 C3 48 8B 00 48 83 C4 28 C3": None,
    "48 83 EC 28 48 8B 05 ? ? ? ? BA ? ? ? ? 44 0F B6 C1 48 8B 88 ? ? ? ? E8 ? ? ? ? 48 85 C0 75 05 48 83 C4 28 C3 48 8B 00 48 83 C4 28 C3": None,
    "48 83 EC 28 48 8B 05 ? ? ? ? 44 8D 81 ? ? ? ? BA ? ? ? ? 48 8B 88 ? ? ? ? E8 ? ? ? ? 48 85 C0 75 05 48 83 C4 28 C3 48 8B 00 48 83 C4 28 C3": None,
    "48 83 EC 38 48 8B 05 ? ? ? ? 44 8B CA 44 8B C1 48 C7 44 24 ? ? ? ? ? BA ? ? ? ? 48 C7 44 24 ? ? ? ? ? 48 8B 88 ? ? ? ? E8 ? ? ? ? 48 85 C0 75 05 48 83 C4 38 C3 48 8B 00 48 83 C4 38 C3": None,
    "48 83 EC 28 48 8B 05 ? ? ? ? BA ? ? ? ? 48 8B 88 ? ? ? ? E8 ? ? ? ? 48 85 C0 74 14 48 8B 10 48 8B C8 FF 52 08 84 C0 75 07 B0 01 48 83 C4 28 C3 32 C0 48 83 C4 28 C3": None,
    "48 83 EC 28 85 C9 74 20 48 8B 05 ? ? ? ? 44 8B C1 BA ? ? ? ? 48 8B 88 ? ? ? ? E8 ? ? ? ? 48 85 C0 75 07 33 C0 48 83 C4 28 C3 48 8B 00 48 83 C4 28 C3": None,
    "48 83 EC 28 48 8B 05 ? ? ? ? 44 8B C1 BA ? ? ? ? 48 8B 88 ? ? ? ? E8 ? ? ? ? 48 85 C0 74 17 48 8B 08 48 85 C9 74 0F 8B 01 25 ? ? ? ? 48 03 C1 48 83 C4 28 C3 33 C0 48 83 C4 28 C3": None,

    # unsure if this is totally accurate but it looks to be the case
    "48 8B 05 ? ? ? ? BA ? ? ? ? 48 8B 88 ? ? ? ? E9 ? ? ? ?": "::rowCount"
}

# todo: figure out how/where these exd getters are used
# .text:0000000140622200                         sub_140622200   proc near               ; CODE XREF: sub_14067D8E0+D3
# .text:0000000140622200 48 8B 05 F1 7F 46 01                    mov     rax, cs:qword_141A8A1F8
# .text:0000000140622207 BA 59 01 00 00                          mov     edx, 159h
# .text:000000014062220C 48 8B 88 E8 2B 00 00                    mov     rcx, [rax+2BE8h]
# .text:0000000140622213 E9 28 E2 E2 FF                          jmp     sub_140450440
# .text:0000000140622213                         sub_140622200   endp

exd_map = {
    209: "Achievement",
    210: "AchievementCategory",
    211: "AchievementKind",
    4: "Action",
    656: "ActionCastTimeline",
    624: "ActionCastVFX",
    3: "ActionCategory",
    642: "ActionComboRoute",
    651: "ActionComboRouteTransient",
    431: "ActionIndirection",
    751: "ActionInit",
    613: "ActionParam",
    430: "ActionProcStatus",
    100: "ActionTimeline",
    709: "ActionTimelineMove",
    608: "ActionTimelineReplace",
    568: "ActionTransient",
    646: "ActivityFeedButtons",
    619: "ActivityFeedCaptions",
    620: "ActivityFeedGroupCaptions",
    618: "ActivityFeedImages",
    0: "Addon",
    251: "AddonHud",
    640: "AddonHudLayout",
    667: "AddonHudSize",
    731: "AddonLayout",
    250: "AddonParam",
    596: "AddonTalkParam",
    402: "AddonTransient",
    316: "Adventure",
    648: "AdventureExPhase",
    452: "AetherCurrent",
    440: "AetherCurrentCompFlgSet",
    357: "AetherialWheel",
    252: "Aetheryte",
    253: "AetheryteSystemDefine",
    409: "AirshipExplorationLevel",
    412: "AirshipExplorationLog",
    437: "AirshipExplorationParamType",
    408: "AirshipExplorationPart",
    410: "AirshipExplorationPoint",
    487: "AirshipSkyIsland",
    242: "AnimationLOD",
    532: "AnimaWeapon5",
    530: "AnimaWeapon5Param",
    525: "AnimaWeapon5PatternGroup",
    545: "AnimaWeapon5SpiritTalk",
    548: "AnimaWeapon5SpiritTalkParam",
    544: "AnimaWeapon5SpiritTalkType",
    524: "AnimaWeapon5TradeItem",
    57: "AnimaWeaponFUITalk",
    200: "AnimaWeaponFUITalkParam",
    531: "AnimaWeaponIcon",
    533: "AnimaWeaponItem",
    734: "AozAction",
    747: "AozActionTransient",
    766: "AOZArrangement",
    767: "AOZBoss",
    763: "AOZContent",
    765: "AOZContentBriefingBNpc",
    761: "AOZScore",
    777: "AOZWeeklyReward",
    561: "AquariumFish",
    560: "AquariumWater",
    61: "ArrayEventHandler",
    101: "AttackType",
    64: "Attributive",
    605: "BacklightColor",
    270: "Ballista",
    6: "Balloon",
    234: "BaseParam",
    63: "Battalion",
    77: "BattleLeve",
    79: "BattleLeveRule",
    511: "BeastRankBonus",
    284: "BeastReputationRank",
    285: "BeastTribe",
    7: "Behavior",
    8: "BehaviorPath",
    662: "BgcArmyAction",
    664: "BgcArmyActionTransient",
    134: "BGM",
    137: "BGMFade",
    138: "BGMFadeType",
    135: "BGMScene",
    136: "BGMSituation",
    672: "BGMSwitch",
    139: "BGMSystemDefine",
    629: "BNpcAnnounceIcon",
    46: "BNpcBase",
    48: "BNpcCustomize",
    47: "BNpcName",
    49: "BNpcParts",
    50: "BNpcState",
    230: "Buddy",
    228: "BuddyAction",
    229: "BuddyEquip",
    233: "BuddyItem",
    231: "BuddyRank",
    232: "BuddySkill",
    243: "Cabinet",
    631: "CabinetCategory",
    161: "Calendar",
    282: "Carry",
    266: "Channeling",
    175: "CharaMakeClassEquip",
    177: "CharaMakeCustomize",
    176: "CharaMakeName",
    174: "CharaMakeType",
    362: "ChocoboRace",
    354: "ChocoboRaceAbility",
    353: "ChocoboRaceAbilityType",
    365: "ChocoboRaceCalculateParam",
    497: "ChocoboRaceChallenge",
    355: "ChocoboRaceItem",
    369: "ChocoboRaceRank",
    397: "ChocoboRaceRanking",
    370: "ChocoboRaceStatus",
    371: "ChocoboRaceTerritory",
    356: "ChocoboRaceTutorial",
    396: "ChocoboRaceWeather",
    202: "ChocoboTaxi",
    203: "ChocoboTaxiStand",
    59: "ClassJob",
    60: "ClassJobCategory",
    454: "ClassJobResident",
    553: "ColorFilter",
    570: "Colosseum",
    495: "ColosseumMatchRank",
    172: "Companion",
    173: "CompanionMove",
    455: "CompanionTransient",
    264: "CompanyAction",
    429: "CompanyCraftDraft",
    425: "CompanyCraftDraftCategory",
    448: "CompanyCraftManufactoryState",
    423: "CompanyCraftPart",
    422: "CompanyCraftProcess",
    400: "CompanyCraftSequence",
    449: "CompanyCraftSupplyItem",
    428: "CompanyCraftType",
    91: "CompanyLeve",
    92: "CompanyLeveRule",
    333: "CompleteJournal",
    334: "CompleteJournalCategory",
    66: "Completion",
    700: "Condition",
    103: "ConfigKey",
    638: "ContentAttributeRect",
    718: "ContentCloseCycle",
    652: "ContentDirectorManagedSG",
    695: "ContentEffectiveTime",
    704: "ContentEntry",
    601: "ContentExAction",
    476: "ContentFinderCondition",
    555: "ContentFinderConditionTransient",
    615: "ContentGauge",
    616: "ContentGaugeColor",
    475: "ContentMemberType",
    151: "ContentNpcTalk",
    671: "ContentRewardCondition",
    153: "ContentRoulette",
    707: "ContentRouletteOpenRule",
    670: "ContentRouletteRoleBonus",
    289: "ContentsNote",
    290: "ContentsNoteCategory",
    291: "ContentsNoteLevel",
    725: "ContentsNoteRewardEurekaEXP",
    149: "ContentTalk",
    150: "ContentTalkParam",
    292: "ContentType",
    509: "ContentUICategory",
    54: "CraftAction",
    78: "CraftLeve",
    80: "CraftLeveTalk",
    52: "CraftType",
    445: "Credit",
    769: "CreditBackImage",
    446: "CreditCast",
    776: "CreditDataSet",
    770: "CreditFont",
    773: "CreditList",
    775: "CreditListText",
    787: "CreditVersion",
    416: "Currency",
    417: "CurrencyLimit",
    114: "CustomTalk",
    758: "CustomTalkDynamicIcon",
    781: "CustomTalkNestHandlers",
    128: "Cutscene",
    414: "CutsceneActorSize",
    439: "CutsceneEventMotion",
    184: "CutsceneMotion",
    733: "CutsceneName",
    129: "CutsceneWorkIndex",
    288: "CutScreenImage",
    286: "CycleTime",
    317: "DailySupplyItem",
    541: "DeepDungeon",
    546: "DeepDungeonBan",
    547: "DeepDungeonDanger",
    540: "DeepDungeonEquipment",
    549: "DeepDungeonFloorEffectUI",
    550: "DeepDungeonGrowData",
    539: "DeepDungeonItem",
    534: "DeepDungeonLayer",
    714: "DeepDungeonMagicStone",
    518: "DeepDungeonMap5X",
    515: "DeepDungeonRoom",
    543: "DeepDungeonStatus",
    113: "DefaultTalk",
    180: "DefaultTalkLipSyncType",
    35: "DeliveryQuest",
    720: "Description",
    713: "DescriptionPage",
    712: "DescriptionSection",
    729: "DescriptionString",
    323: "DirectorSystemDefine",
    324: "DirectorType",
    577: "DisposalShop",
    576: "DisposalShopFilterType",
    578: "DisposalShopItem",
    724: "DomaStoryProgress",
    492: "DpsChallenge",
    491: "DpsChallengeOfficer",
    581: "DpsChallengeTransient",
    752: "EmjAddon",
    762: "EmjCharaViewCamera",
    778: "EmjDani",
    181: "Emote",
    183: "EmoteCategory",
    182: "EmoteMode",
    97: "ENpcBase",
    784: "ENpcDressUp",
    783: "ENpcDressUpDress",
    96: "ENpcResident",
    106: "EObj",
    563: "EObjName",
    645: "EquipRaceCategory",
    14: "EquipSlotCategory",
    68: "Error",
    723: "Eureka",
    722: "EurekaAetherItem",
    708: "EurekaAethernet",
    782: "EurekaDungeonPortal",
    705: "EurekaGrowData",
    748: "EurekaLogosMixerProbability",
    742: "EurekaMagiaAction",
    744: "EurekaMagiciteItem",
    743: "EurekaMagiciteItemType",
    706: "EurekaSphereElementAdjust",
    710: "EurekaStoryProgress",
    123: "EventAction",
    126: "EventIconPriority",
    127: "EventIconType",
    124: "EventItem",
    464: "EventItemCastTimeline",
    512: "EventItemCategory",
    564: "EventItemHelp",
    465: "EventItemTimeline",
    125: "EventSystemDefine",
    573: "EventVfx",
    692: "ExHotbarCrossbarIndexType",
    155: "ExportedSG",
    574: "ExtraCommand",
    419: "ExVersion",
    699: "FashionCheckThemeCategory",
    701: "FashionCheckWeeklyTheme",
    164: "Fate",
    165: "FateEvent",
    214: "FCActivity",
    224: "FCActivityCategory",
    215: "FCAuthority",
    217: "FCAuthorityCategory",
    221: "FCChestName",
    223: "FCCrestSymbol",
    421: "FccShop",
    222: "FCDefine",
    216: "FCHierarchy",
    225: "FCProfile",
    218: "FCRank",
    219: "FCReputation",
    220: "FCRights",
    262: "Festival",
    296: "FieldMarker",
    249: "FishingRecordType",
    636: "FishingRecordTypeTransient",
    246: "FishingSpot",
    247: "FishParameter",
    406: "Frontline",
    208: "Frontline01",
    373: "Frontline02",
    424: "Frontline03",
    527: "Frontline04",
    293: "GardeningSeed",
    28: "GatheringCondition",
    25: "GatheringExp",
    21: "GatheringItem",
    413: "GatheringItemLevelConvertTable",
    335: "GatheringItemPoint",
    88: "GatheringLeve",
    90: "GatheringLeveBNpcEntry",
    89: "GatheringLeveRoute",
    87: "GatheringLeveRule",
    29: "GatheringNotebook",
    336: "GatheringNotebookList",
    24: "GatheringPoint",
    22: "GatheringPointBase",
    447: "GatheringPointBonus",
    23: "GatheringPointBonusType",
    27: "GatheringPointName",
    442: "GatheringSubCategory",
    26: "GatheringType",
    554: "GcArmyCandidateCategory",
    660: "GcArmyCapture",
    668: "GcArmyCaptureTactics",
    669: "GcArmyEquipPreset",
    519: "GcArmyExpedition",
    520: "GcArmyExpeditionMemberBonus",
    521: "GcArmyExpeditionTrait",
    522: "GcArmyExpeditionTraitCond",
    529: "GcArmyExpeditionType",
    280: "GcArmyMember",
    542: "GcArmyMemberGrow",
    537: "GcArmyMemberGrowExp",
    516: "GcArmyProgress",
    517: "GcArmyTraining",
    190: "GCRankGridaniaFemaleText",
    189: "GCRankGridaniaMaleText",
    188: "GCRankLimsaFemaleText",
    187: "GCRankLimsaMaleText",
    192: "GCRankUldahFemaleText",
    191: "GCRankUldahMaleText",
    485: "GCScripShopCategory",
    486: "GCScripShopItem",
    196: "GCShop",
    194: "GCShopItemCategory",
    199: "GCSupplyDefine",
    197: "GCSupplyDuty",
    198: "GCSupplyDutyReward",
    104: "GeneralAction",
    388: "GFATE",
    361: "GFateClimbing",
    738: "GFateClimbing2",
    746: "GFateClimbing2Content",
    741: "GFateClimbing2TotemType",
    377: "GFateDance",
    364: "GFateHiddenObject",
    760: "GFateRideShooting",
    378: "GFateRoulette",
    372: "GFateStelth",
    603: "GilShop",
    736: "GilShopInfo",
    602: "GilShopItem",
    120: "GimmickAccessor",
    119: "GimmickBill",
    314: "GimmickJump",
    121: "GimmickRect",
    122: "GimmickYesNo",
    346: "GoldSaucerArcadeMachine",
    556: "GoldSaucerContent",
    327: "GoldSaucerTalk",
    366: "GoldSaucerTextData",
    186: "GrandCompany",
    193: "GrandCompanyRank",
    575: "GroupPoseCharaStatus",
    185: "GuardianDeity",
    76: "GuildleveAssignment",
    75: "GuildleveAssignmentCategory",
    85: "GuildleveAssignmentTalk",
    72: "GuildleveEvaluation",
    146: "GuildOrder",
    147: "GuildOrderGuide",
    148: "GuildOrderOfficer",
    287: "HairMakeType",
    461: "HouseRetainerPose",
    635: "HousingAethernet",
    740: "HousingAppeal",
    678: "HousingEmploymentNpcList",
    676: "HousingEmploymentNpcRace",
    272: "HousingExterior",
    277: "HousingFurniture",
    273: "HousingInterior",
    279: "HousingLandSet",
    569: "HousingMapMarkerInfo",
    469: "HousingMateAuthority",
    737: "HousingMerchantPose",
    719: "HousingPileLimit",
    471: "HousingPlacement",
    276: "HousingPreset",
    426: "HousingUnitedExterior",
    513: "HousingUnplacement",
    275: "HousingYardObject",
    140: "HowTo",
    142: "HowToCategory",
    141: "HowToPage",
    143: "InstanceContent",
    154: "InstanceContentBuff",
    785: "InstanceContentCSBonus",
    152: "InstanceContentGuide",
    528: "InstanceContentRewardItem",
    145: "InstanceContentTextData",
    144: "InstanceContentType",
    10: "Item",
    16: "ItemAction",
    11: "ItemFood",
    17: "ItemLevel",
    9: "ItemSearchCategory",
    12: "ItemSeries",
    13: "ItemSpecialBonus",
    15: "ItemUICategory",
    265: "Jingle",
    162: "JournalCategory",
    163: "JournalGenre",
    503: "JournalSection",
    5: "Knockback",
    33: "LegacyQuest",
    73: "Leve",
    71: "LeveAssignmentType",
    74: "LeveClient",
    168: "Level",
    81: "LeveRewardItem",
    82: "LeveRewardItemGroup",
    86: "LeveString",
    83: "LeveSystemDefine",
    84: "LeveVfx",
    1: "LinkRace",
    663: "LoadingImage",
    348: "LoadingTips",
    349: "LoadingTipsSub",
    62: "Lobby",
    294: "Lockon",
    160: "LogFilter",
    157: "LogKind",
    159: "LogKindCategoryText",
    158: "LogMessage",
    715: "LootModeType",
    579: "LotteryExchangeShop",
    263: "MacroIcon",
    340: "MacroIconRedirectOld",
    239: "MainCommand",
    240: "MainCommandCategory",
    674: "Maneuvers",
    673: "ManeuversArmor",
    107: "Map",
    108: "MapMarker",
    630: "MapMarkerRegion",
    109: "MapSymbol",
    95: "Marker",
    443: "MasterpieceSupplyDuty",
    617: "MasterpieceSupplyMultiplier",
    470: "MateAuthorityCategory",
    19: "Materia",
    18: "MateriaJoinRate",
    562: "MateriaParam",
    20: "MateriaTomestoneRate",
    749: "MiniGameRA",
    750: "MiniGameRANotes",
    479: "MinionRace",
    473: "MinionRules",
    484: "MinionSkillType",
    436: "MinionStage",
    434: "MobHuntOrder",
    611: "MobHuntOrderType",
    435: "MobHuntReward",
    415: "MobHuntRewardCap",
    320: "MobHuntTarget",
    687: "ModelAttribute",
    38: "ModelChara",
    689: "ModelScale",
    39: "ModelSkeleton",
    693: "ModelState",
    212: "MonsterNote",
    213: "MonsterNoteTarget",
    110: "MotionTimeline",
    112: "MotionTimelineAdvanceBlend",
    111: "MotionTimelineBlendTable",
    133: "Mount",
    418: "MountAction",
    698: "MountCustomize",
    404: "MountFlyingCondition",
    641: "MountSpeed",
    463: "MountTransient",
    331: "MoveControl",
    523: "MoveTimeline",
    587: "MoveVfx",
    258: "MovieSubtitle",
    259: "MovieSubtitleVoyage",
    69: "NotebookDivision",
    786: "NotebookDivisionCategory",
    102: "NotebookList",
    607: "NotoriousMonster",
    606: "NotoriousMonsterTerritory",
    204: "NpcEquip",
    179: "NpcYell",
    238: "Omen",
    156: "OnlineStatus",
    226: "Opening",
    227: "OpeningSystemDefine",
    499: "Orchestrion",
    505: "OrchestrionCategory",
    504: "OrchestrionPath",
    506: "OrchestrionUiparam",
    169: "ParamGrow",
    481: "PartyContent",
    480: "PartyContentCutscene",
    483: "PartyContentTextData",
    490: "PartyContentTransient",
    321: "PatchMark",
    679: "Perform",
    680: "PerformTransient",
    330: "Permission",
    260: "Pet",
    261: "PetAction",
    30: "PetMirage",
    255: "PhysicsGroup",
    256: "PhysicsOffGroup",
    257: "PhysicsWind",
    582: "Picture",
    585: "PictureCategory",
    99: "PlaceName",
    508: "PlantPotFlowerSeed",
    717: "PreHandler",
    117: "PresetCamera",
    241: "PresetCameraAdjust",
    166: "PublicContent",
    339: "PublicContentCutscene",
    167: "PublicContentTextData",
    450: "Purify",
    205: "PvPAction",
    643: "PvPActionSort",
    647: "PvPInitialSelectActionTrait",
    206: "PvPRank",
    468: "PvPRankTransient",
    639: "PvPSelectTrait",
    655: "PvPSelectTraitTransient",
    207: "PvPTrait",
    622: "QTE",
    32: "Quest",
    170: "QuestBattle",
    171: "QuestBattleSystemDefine",
    474: "QuestClassJobReward",
    711: "QuestClassJobSupply",
    494: "QuestEquipModel",
    34: "QuestRepeatFlag",
    432: "QuestRewardOther",
    31: "QuestSystemDefine",
    593: "QuickChat",
    594: "QuickChatTransient",
    178: "Race",
    385: "RacingChocoboGrade",
    363: "RacingChocoboItem",
    375: "RacingChocoboName",
    382: "RacingChocoboNameCategory",
    384: "RacingChocoboNameInfo",
    379: "RacingChocoboParam",
    105: "RecastNavimesh",
    51: "Recipe",
    55: "RecipeElement",
    56: "RecipeLevelTable",
    332: "RecipeNotebookList",
    269: "RecommendContents",
    306: "Relic",
    310: "Relic3",
    311: "Relic3Materia",
    312: "Relic3Rate",
    313: "Relic3RatePattern",
    394: "Relic6Magicite",
    344: "RelicItem",
    305: "RelicMateria",
    307: "RelicNote",
    309: "RelicNoteCategory",
    325: "Resident",
    586: "ResidentMotionType",
    297: "RetainerTask",
    298: "RetainerTaskLvRange",
    299: "RetainerTaskNormal",
    302: "RetainerTaskParameter",
    552: "RetainerTaskParameterLvDiff",
    301: "RetainerTaskRandom",
    768: "RideShooting",
    753: "RideShootingScheduler",
    756: "RideShootingTarget",
    755: "RideShootingTargetScheduler",
    757: "RideShootingTargetType",
    759: "RideShootingTextData",
    130: "Role",
    716: "RPParameter",
    661: "Salvage",
    591: "SatisfactionNpc",
    590: "SatisfactionSupply",
    592: "SatisfactionSupplyReward",
    666: "SatisfactionSupplyRewardExp",
    612: "ScenarioTree",
    614: "ScenarioTreeTips",
    690: "ScenarioTreeTipsClassQuest",
    688: "ScenarioTreeTipsQuest",
    659: "ScenarioType",
    322: "ScreenImage",
    131: "SE",
    653: "SEBattle",
    466: "SecretRecipeBook",
    565: "SkyIsland",
    588: "SkyIsland2",
    584: "SkyIsland2Mission",
    583: "SkyIsland2MissionDetail",
    598: "SkyIsland2MissionType",
    595: "SkyIsland2RangeType",
    489: "SkyIslandMapMarker",
    488: "SkyIslandSubject",
    621: "Snipe",
    623: "SnipeCollision",
    633: "SnipeElementId",
    697: "SnipeHitEvent",
    625: "SnipePerformanceCamera",
    628: "SpearfishingEcology",
    610: "SpearfishingItem",
    634: "SpearfishingNotebook",
    637: "SpearfishingRecordPage",
    267: "SpecialShop",
    268: "SpecialShopItemCategory",
    526: "Spectator",
    271: "Stain",
    472: "StainTransient",
    70: "Status",
    451: "StatusHitEffect",
    599: "StatusLoopVFX",
    244: "Story",
    245: "StorySystemDefine",
    683: "SubmarineExploration",
    694: "SubmarineExplorationLog",
    684: "SubmarinePart",
    685: "SubmarineRank",
    682: "SubmarineSpecCategory",
    304: "SwitchTalk",
    53: "TerritoryChatRule",
    93: "TerritoryIntendedUse",
    94: "TerritoryType",
    498: "TerritoryTypeTransient",
    65: "TextCommand",
    67: "TextCommandParam",
    116: "Title",
    644: "TomestoneConvert",
    303: "Tomestones",
    571: "TomestonesItem",
    507: "TopicSelect",
    42: "Town",
    98: "Trait",
    345: "TraitRecast",
    467: "TraitTransient",
    315: "Transformation",
    36: "Treasure",
    403: "TreasureHuntRank",
    405: "TreasureHuntTexture",
    37: "TreasureModel",
    283: "TreasureSpot",
    58: "Tribe",
    352: "TripleTriad",
    347: "TripleTriadCard",
    392: "TripleTriadCardRarity",
    395: "TripleTriadCardResident",
    374: "TripleTriadCardType",
    390: "TripleTriadCompetition",
    391: "TripleTriadDefine",
    387: "TripleTriadResident",
    360: "TripleTriadRule",
    493: "Tutorial",
    502: "TutorialDPS",
    501: "TutorialHealer",
    500: "TutorialTank",
    735: "UIColor",
    535: "Vase",
    510: "VaseFlower",
    118: "VFX",
    235: "Warp",
    236: "WarpCondition",
    237: "WarpLogic",
    115: "WeaponTimeline",
    43: "Weather",
    45: "WeatherGroup",
    44: "WeatherRate",
    626: "WeatherReportReplace",
    343: "WeddingBGM",
    342: "WeddingFlowerColor",
    341: "WeddingPlan",
    558: "WeeklyBingoOrderData",
    559: "WeeklyBingoRewardData",
    557: "WeeklyBingoText",
    368: "WeeklyLotBonus",
    367: "WeeklyLotBonusThreshold",
    328: "World",
    329: "WorldDCGroupType",
    696: "XPVPGroupActivity",
    514: "YKW",
    441: "ZoneSharedGroup",
    433: "ZoneTimeline"
}

def do_pattern(pattern, suffix = ""):
    ea = 0

    while True:
        ea = FindBinary(ea + 1, SEARCH_DOWN, pattern)

        if ea == 0xFFFFFFFFFFFFFFFF:
            break

        # this is mega retarded but it works rofl
        ins = FindBinary(ea, SEARCH_DOWN, "BA ? ? ? ?")
        sheetIdx = idc.Dword(ins + 1)

        origName = GetFunctionName(ea)

        # don't rename any funcs that are already named
        if origName[0:4] == "sub_":
            sheetName = exd_map[sheetIdx]

            if suffix == None:
                suffix = ""

            fnName = "Client::ExdData::get%s%s" % (exd_map[sheetIdx], suffix)

            print("found exd func @ %x -> mapped to %s (%i)" % (ea, sheetName, sheetIdx))

            MakeName(ea, fnName)
            MakeComm(ins, "Sheet: %s (%i)" % (sheetName, sheetIdx))


def run():
    # todo: this doesnt find all getters, there's a few slightly different ones
    # along with others that call different virts in slightly different ways/different args
    for pattern, suffix in exd_func_patterns.items():
        do_pattern(pattern, suffix)


class ffxiv_exdgetters_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL

    wanted_name = "FFXIV - Annotate EXD Getters"
    wanted_hotkey = ""

    comment = 'Automagically names EXD getter funcs'
    help = 'no'
 
    def init(self):
        return idaapi.PLUGIN_OK
 
    def run(self, arg):
        run()
 
    def term(self):
        pass
 
def PLUGIN_ENTRY():
    return ffxiv_exdgetters_t()


repository: computerline1z/IDAPatcher6.1
file: idapatcher-r1.1.2.py
description: IDA Patcher for v6.1
wanted_name: IDA Patcher
comment: Enhances manipulation and application of patched bytes.
created_at: 2014-03-31 04:44:29+00:00
pushed_at: 2014-03-31 12:04:13+00:00
url:
https://github.com/computerline1z/IDAPatcher6.1/blob/b7cc364b58b8deadf7b9c60d05218ab7f75227c8/idapatcher-r1.1.2.py

content:
#!/usr/bin/env python
#
# IDA Patcher is a plugin for Hex-Ray's IDA Pro disassembler designed to 
# enhance IDA's ability to patch binary files and memory. The plugin is 
# useful for tasks related to malware analysis, exploit development as well
# as bug patching. IDA Patcher blends into the standard IDA user interface
# through the addition of a subview and several menu items. 
#
# VERSION 0.1.1
#
# Copyright (C) 2013 Peter Kacherginsky
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met: 
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer. 
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution. 
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import idaapi
import idautils
import idc

from idaapi import Form, Choose2, plugin_t

import os
import shutil
import struct
import binascii

IDAPATCHER_VERSION = "0.1.1.1"

#--------------------------------------------------------------------------
# Forms
#--------------------------------------------------------------------------
class PatchRestoreForm(Form):
    """
    Form to aid in restoring patched bytes to their original values.
    """
    def __init__(self, addr_str, fpos_str, patch_str, org_str):
        Form.__init__(self, 
r"""BUTTON YES* Restore
BUTTON CANCEL Cancel
Restore patch bytes (Please don't edit)

Address        {strAddr}
File offset    {strFpos}
<:{strOrg}>
""", {
        'strAddr': Form.StringLabel(addr_str),
        'strFpos': Form.StringLabel(fpos_str),
        #'strOrg': Form.MultiLineTextControl(text=org_str, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_READONLY),
        'strOrg': Form.StringInput(tp=None, width=1024, swidth=40, hlp=None, value=org_str, size=None),
        })
        self.Compile()

#--------------------------------------------------------------------------
class PatchEditForm(Form):
    """
    Form to edit patched bytes.
    """
    def __init__(self, addr_str, fpos_str, patch_str, org_str):
        Form.__init__(self,
r"""Edit patch bytes

Address        {strAddr}
File offset    {strFpos}
<:{strPatch}>
""", {
        'strAddr':  Form.StringLabel(addr_str),
        'strFpos':  Form.StringLabel(fpos_str),
        #'strPatch': Form.MultiLineTextControl(text=patch_str, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT),
        'strPatch': Form.StringInput(tp=None, width=1024, swidth=40, hlp=None, value=patch_str, size=None),
        })

        self.Compile()

#--------------------------------------------------------------------------
class PatchApplyForm(Form):
    """
    Form to prompt for target file, backup file, and the address
    range to save patched bytes.
    """
    def __init__(self, start_ea, end_ea, org_file, bkp_file):
        Form.__init__(self,
r"""Apply patches to input file

{FormChangeCb}
<##Start EA   :{intStartEA}>
<##End EA     :{intEndEA}>
<##Input file :{orgFile}>
<##Backup file:{bkpFile}>

<##Create backup:{rBackup}>
<##Restore original bytes:{rRestore}>{cGroup1}>
""", {
        'intStartEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=start_ea),
        'intEndEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=end_ea),
        'orgFile': Form.FileInput(swidth=50, open=True, value=org_file),
        'bkpFile': Form.FileInput(swidth=50, open=True, value=bkp_file),
        'cGroup1': Form.ChkGroupControl(("rBackup", "rRestore")),
        'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
        })

        self.Compile()

    def OnFormChange(self, fid):
        # Set initial state
        if fid == -1:
            self.EnableField(self.bkpFile, False)

        # Toggle backup checkbox
        elif fid == self.rBackup.id:
            self.rBackup.checked = not self.rBackup.checked
            self.EnableField(self.bkpFile, self.rBackup.checked)

        # Toggle restore checkbox
        elif fid == self.rRestore.id:
            self.rRestore.checked = not self.rRestore.checked

        return 1

#--------------------------------------------------------------------------
class PatchFillForm(Form):
    """
    Form to fill a range of addresses with a specified byte value.
    """
    def __init__(self, start_ea, end_ea, fill_value):

        Form.__init__(self,
r"""BUTTON YES* Fill
Fill bytes

<##Start EA   :{intStartEA}>
<##End EA     :{intEndEA}>
<##Value      :{intPatch}>
""", {
        'intStartEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=start_ea),
        'intEndEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=end_ea),
        'intPatch': Form.NumericInput(swidth=40,tp=Form.FT_HEX,value=fill_value),
        })

        self.Compile()

class DataImportForm(Form):
    """
    Form to import data of various types into selected area.
    """
    def __init__(self, start_ea, end_ea):
        Form.__init__(self,
r"""BUTTON YES* Import
Import data

{FormChangeCb}
<##Start EA   :{intStartEA}>
<##End EA     :{intEndEA}>

Import type:                    Patching options:
<hex string:{rHex}><##Trim to selection:{cSize}>{cGroup}>
<string literal:{rString}>
<binary file:{rFile}>{rGroup}>

<:{strPatch}>
<##Import BIN file:{impFile}>
""", {        
        'intStartEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=start_ea),
        'intEndEA': Form.NumericInput(swidth=40,tp=Form.FT_ADDR,value=end_ea),

        'cGroup': Form.ChkGroupControl(("cSize",)),
        'rGroup': Form.RadGroupControl(("rHex", "rString", "rFile")),

        #'strPatch': Form.MultiLineTextControl(swidth=80, flags=Form.MultiLineTextControl.TXTF_FIXEDFONT),
        'strPatch': Form.StringInput(tp=None, width=1024, swidth=40, hlp=None, value=None, size=None),
        'impFile': Form.FileInput(swidth=50, open=True),

        'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
        })

        self.Compile()

    def OnFormChange(self, fid):
        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.strPatch)
            self.EnableField(self.strPatch, True)
            self.EnableField(self.impFile, False)

        # Form OK pressed
        elif fid == -2:
            pass

        # Form from text box
        elif fid == self.rHex.id or fid == self.rString.id:
            self.SetFocusedField(self.strPatch)
            self.EnableField(self.strPatch, True)
            self.EnableField(self.impFile, False)

        # Form import from file
        elif fid == self.rFile.id:
            self.SetFocusedField(self.rFile)
            self.EnableField(self.impFile, True)
            self.EnableField(self.strPatch, False)

        return 1

#--------------------------------------------------------------------------
# Chooser
#--------------------------------------------------------------------------
class PatchView(Choose2):
    """
    Chooser class to display and manage patched bytes in the database.
    """
    def __init__(self):
        Choose2.__init__(self,
                         "Patches",
                         [ ["Address",  10 | Choose2.CHCOL_HEX], 
                           ["Name",     18 | Choose2.CHCOL_PLAIN], 
                           ["Size",      4 | Choose2.CHCOL_DEC], 
                           ["Modified", 10 | Choose2.CHCOL_HEX],
                           ["Original", 10 | Choose2.CHCOL_HEX], 
                           ["Comment",  30 | Choose2.CHCOL_PLAIN]
                         ],
                         flags = Choose2.CH_MULTI_EDIT)

        self.popup_names = ["Insert", "Delete", "Edit", "Refresh"]
        
        self.icon = 47

        # Items for display and corresponding data
        # NOTE: Could become desynchronized, so to avoid this
        #       refresh the view after each change.
        self.items = []
        self.items_data  = []

        # Initialize/Refresh the view
        self.refreshitems()

        # Data members
        self.patch_file = None
        self.restore = False

        # Command callbacks
        self.cmd_apply_patches = None
        self.cmd_restore_bytes = None

    def show(self):
        # Attempt to open the view
        if self.Show() < 0: return False

        # Add extra context menu commands
        # NOTE: Make sure you check for duplicates.
        if self.cmd_apply_patches == None:
            self.cmd_apply_patches = self.AddCommand("Apply patches to input file...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_NO_SELECTION, icon=27)
        if self.cmd_restore_bytes == None:
            self.cmd_restore_bytes = self.AddCommand("Restore original byte(s)...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=139)

        return True

    # Patch byte visitor callback to apply the patches
    # NOTE: Only bytes with fpos > -1 can be applied.
    def apply_patch_byte(self, ea, fpos, org_val, patch_val):
        if fpos != -1:
            self.patch_file.seek(fpos)

            if self.restore:
                self.patch_file.write(struct.pack('B', org_val))
            else:
                self.patch_file.write(struct.pack('B', patch_val))

        return 0

    # Patch byte visitor callback to collect and aggregate bytes
    def get_patch_byte(self, ea, fpos, org_val, patch_val):

        # Aggregate contiguous bytes (base ea + length)
        # NOTE: Looking at the last item [-1] is sufficient
        #       since we are dealing with sorted data.
        if len(self.items_data) and (ea - self.items_data[-1][0] == self.items_data[-1][2]):

            # Increment length
            self.items_data[-1][2] += 1
            self.items[-1][2] = str(self.items_data[-1][2])

            # Append patched bytes
            self.items_data[-1][3].append(patch_val)
            self.items[-1][3] = " ".join(["%02X" % x for x in self.items_data[-1][3]])

            # Append original bytes
            self.items_data[-1][4].append(org_val)
            self.items[-1][4] =  " ".join(["%02X" % x for x in self.items_data[-1][4]])


        # Add new patch byte to the list
        else:

            name = SegName(ea)

            if GetFunctionName(ea) or Name(ea):
                name += ": %s" % GetFunctionName(ea) or Name(ea)


            comment = Comment(ea) or RptCmt(ea) or ""
            # DATA STORAGE FORMAT:      address, function / fpos, len,    patched byte(s), original byte(s), comments
            self.items.append(     ["%08X" % ea,            name, "1", "%02X" % patch_val, "%02X" % org_val, comment])
            self.items_data.append([         ea,            fpos,   1,        [patch_val],        [org_val], None]   )

        return 0

    def refreshitems(self):
        print "[*] IdaPatcher::refeshitems()"
        self.items_data = []
        self.items = []
        visit_patched_bytes(0, idaapi.BADADDR, self.get_patch_byte)

    def OnCommand(self, n, cmd_id):

        # Apply patches to a file
        if cmd_id == self.cmd_apply_patches:

            # Set initial start/end EA values
            start_ea = 0x0
            end_ea = idaapi.cvar.inf.maxEA

            # Set initial output file values
            org_file = GetInputFilePath()
            bkp_file = "%s.bak" % org_file

            # Create the form
            f = PatchApplyForm(start_ea, end_ea, org_file, bkp_file)

            # Execute the form
            ok = f.Execute()
            if ok == 1:
                # Get restore checkbox
                self.restore = f.rRestore.checked

                # Get updated ea max/min
                start_ea = f.intStartEA.value
                end_ea = f.intEndEA.value

                # Get updated file path
                new_org_file = f.orgFile.value

                # Backup the file before replacing
                if f.rBackup.checked:
                    bkp_file = f.bkpFile.value
                    shutil.copyfile(org_file, bkp_file)

                # Apply patches
                try:
                    self.patch_file = open(new_org_file,'rb+')
                except Exception, e:
                    idaapi.warning("Cannot update file '%s'" % new_org_file)
                else:
                    #r = idaapi.visit_patched_bytes(start_ea, end_ea, self.apply_patch_byte)
                    r = visit_patched_bytes(start_ea, end_ea, self.apply_patch_byte)
                    self.patch_file.close()
                    self.restore = False

                    # Update db input file, so we are working
                    # with a patched version.
                    #if not org_file == new_org_file:
                    #    idaapi.set_root_filename(new_org_file)
                    #    org_file = new_org_file

            # Dispose the form
            f.Free()

        # Restore selected byte(s)
        elif cmd_id == self.cmd_restore_bytes:

            # Empty list

            if n == -2 or n ==-3:
                return 1

            elif not len(self.items) > 0:
                idaapi.warning("There are no patches to restore.")
                return 1

            # Nothing selected
            elif n == -1:
                idaapi.warning("Please select bytes to restore.")
                return 1

            ea = self.items_data[n][0]
            fpos =  self.items_data[n][1]
            buf = self.items_data[n][4]

            addr_str = "%#x" % ea
            fpos_str = "%#x" % fpos if fpos != -1 else "N/A"  
            patch_str = self.items[n][3]
            org_str = self.items[n][4]

            # Create the form
            f = PatchRestoreForm(addr_str, fpos_str, patch_str, org_str)

            # Execute the form
            ok = f.Execute()
            if ok == 1:

                # Restore original bytes
                idaapi.put_many_bytes(ea, struct.pack("B"*len(buf), *buf))

                # Refresh all IDA views
                self.refreshitems()
                idaapi.refresh_idaview_anyway()

            # Dispose the form
            f.Free()

        return 1

    def OnClose(self):
        self.cmd_apply_patches = None
        self.cmd_restore_bytes = None

    def OnEditLine(self, n):

        # Empty list
        if n == -1:
            return

        # Multiselect START_SEL/END_SEL protocol
        if n == -2 or n ==-3:
            return

        ea = self.items_data[n][0]
        fpos =  self.items_data[n][1]
        patch_buf = self.items_data[n][3]
        orig_buf = self.items_data[n][4]

        addr_str = "%#x" % ea
        fpos_str = "%#x" % fpos if fpos != -1 else "N/A"     
        patch_str = self.items[n][3]
        org_str = self.items[n][4]    

        # Create the form
        f = PatchEditForm(addr_str, fpos_str, patch_str, org_str)

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            # Convert hex bytes to binary
            buf = f.strPatch.value
            buf = buf.replace(' ','')       # remove spaces
            buf = buf.replace('\\x','')     # remove '\x' prefixes
            buf = buf.replace('0x','')      # remove '0x' prefixes
            try:
                buf = binascii.unhexlify(buf)   # convert to bytes
            except Exception, e:
                idaapi.warning("Invalid input: %s" % e)
                f.Free()
                return

            # Restore original bytes first
            idaapi.put_many_bytes(ea, struct.pack("B"*len(orig_buf), *orig_buf))

            # Now apply newly patched bytes
            idaapi.patch_many_bytes(ea, buf)

            # Refresh all IDA views
            self.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()

    def OnSelectLine(self, n):
        idaapi.jumpto(self.items_data[n][0])

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

    def OnGetIcon(self, n):

        # Empty list
        if not len(self.items) > 0:
            return -1

        if self.items_data[n][1] == -1:
            return 138
        else:
            return 137

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

    def OnRefresh(self, n):
        self.refreshitems()
        return n

    def OnActivate(self):
        self.refreshitems()

#--------------------------------------------------------------------------
# Manager
#--------------------------------------------------------------------------
class PatchManager():
    """ Class that manages GUI forms and patching methods of the plugin. """
    
    def __init__(self): 
        self.addmenu_item_ctxs = list()
        self.patch_view = PatchView()

    #--------------------------------------------------------------------------
    # Menu Items
    #--------------------------------------------------------------------------
    def add_menu_item_helper(self, menupath, name, hotkey, flags, pyfunc, args):

        # add menu item and report on errors
        addmenu_item_ctx = idaapi.add_menu_item(menupath, name, hotkey, flags, pyfunc, args)
        if addmenu_item_ctx is None:
            return 1
        else:
            self.addmenu_item_ctxs.append(addmenu_item_ctx)
            print "[*] Add Menu Item -> %s" % (menupath + name)
            return 0

    def add_menu_items(self):
        #print "[*] IdaPatcher::add_menu_items()"
        if self.add_menu_item_helper("View/Open subviews/Problems", "Patches", "", 1, self.show_patches_view, None): return 1
        if self.add_menu_item_helper("Edit/Patch program/", "Edit selection...", "", 0, self.show_edit_form, None):  return 1
        if self.add_menu_item_helper("Edit/Patch program/", "Fill selection...", "", 0, self.show_fill_form, None):  return 1
        if self.add_menu_item_helper("Edit/", "Import data...", "Shift-I", 0, self.show_import_form, None):   return 1
        # The last must be fail :(  EPIC FAIL HERE !!!!
        if self.add_menu_item_helper("Edit/Export data...", "Import data...", "Shift-I", 1, self.show_import_form, None):   return 1
        return 0

    def del_menu_items(self):
        for addmenu_item_ctx in self.addmenu_item_ctxs:
            idaapi.del_menu_item(addmenu_item_ctx)

    #--------------------------------------------------------------------------
    # View Callbacks
    #--------------------------------------------------------------------------   
    
    # Patches View
    def show_patches_view(self):
        print "[*] IdaPatcher::show_fatches_view()"
        self.patch_view.show()

    # Patches Edit Dialog
    def show_edit_form(self):
        selection, start_ea, end_ea = idaapi.read_selection()
        
        if not selection:
            start_ea = idaapi.get_screen_ea()
            end_ea = start_ea + 1

        buf_len = end_ea - start_ea
        buf = get_many_bytes(start_ea, buf_len) or "\xFF"*buf_len
        buf_str = " ".join(["%02X" % ord(x) for x in buf])

        fpos = idaapi.get_fileregion_offset(start_ea)

        addr_str = "%#X" % start_ea
        fpos_str = "%#x" % fpos if fpos != -1 else "N/A" 

        f = PatchEditForm(addr_str, fpos_str, buf_str, buf_str)

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            # Convert hex bytes to binary
            buf = f.strPatch.value
            buf = buf.replace(' ','')       # remove spaces
            buf = buf.replace('\\x','')     # remove '\x' prefixes
            buf = buf.replace('0x','')      # remove '0x' prefixes
            try:
                buf = binascii.unhexlify(buf)   # convert to bytes
            except Exception, e:
                idaapi.warning("Invalid input: %s" % e)
                f.Free()
                return

            # Now apply newly patched bytes
            idaapi.patch_many_bytes(start_ea, buf)

            # Refresh all IDA views
            self.patch_view.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()

    # Fill range with a value form
    def show_fill_form(self):
        selection, start_ea, end_ea = idaapi.read_selection()
        
        if not selection:
            start_ea = idaapi.get_screen_ea()
            end_ea = start_ea + 1
        
        # Default fill value
        fill_value = 0x00

        # Create the form
        f = PatchFillForm(start_ea, end_ea, fill_value)

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            # Get updated values
            start_ea = f.intStartEA.value
            end_ea = f.intEndEA.value
            fill_value = f.intPatch.value

            # Now apply newly patched bytes
            # NOTE: fill_value is expected to be one byte
            #       so if a user provides a larger patch_byte()
            #       will trim the value as expected.


            for ea in range(start_ea, end_ea):
                idaapi.patch_byte(ea, fill_value)

            # Refresh all IDA views
            self.patch_view.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()

    # Import data form
    def show_import_form(self):
        selection, start_ea, end_ea = idaapi.read_selection()

        if not selection:
            start_ea = idaapi.get_screen_ea()
            end_ea = start_ea + 1

        # Create the form
        f = DataImportForm(start_ea, end_ea);

        # Execute the form
        ok = f.Execute()
        if ok == 1:

            start_ea = f.intStartEA.value
            end_ea = f.intEndEA.value

            if f.rFile.selected:
                imp_file = f.impFile.value

                try:
                    f_imp_file = open(imp_file,'rb+')
                except Exception, e:
                    idaapi.warning("File I/O error({0}): {1}".format(e.errno, e.strerror))
                    return
                else:
                    buf = f_imp_file.read()
                    f_imp_file.close()

            else:

                buf = f.strPatch.value

                # Hex values, unlike string literal, needs additional processing
                if f.rHex.selected:
                    buf = buf.replace(' ','')       # remove spaces
                    buf = buf.replace('\\x','')     # remove '\x' prefixes
                    buf = buf.replace('0x','')      # remove '0x' prefixes
                    try:
                        buf = binascii.unhexlify(buf)   # convert to bytes
                    except Exception, e:
                        idaapi.warning("Invalid input: %s" % e)
                        f.Free()
                        return

            if not len(buf):
                idaapi.warning("There was nothing to import.")
                return

            # Trim to selection if needed:
            if f.cSize.checked:
                buf_size = end_ea - start_ea
                buf = buf[0:buf_size]

            # Now apply newly patched bytes
            idaapi.patch_many_bytes(start_ea, buf)

            # Refresh all IDA views
            self.patch_view.refreshitems()
            idaapi.refresh_idaview_anyway()

        # Dispose the form
        f.Free()

#--------------------------------------------------------------------------
# Patch Notify Hook
#--------------------------------------------------------------------------

patched_bytes_dict = {}

def visit_patched_bytes(startEA, endEA, callback):
    for key, val in patched_bytes_dict.items():
        callback(key, val[0], val[1], val[2])
    return 0


class PatchNotifyHook(idaapi.IDB_Hooks):
    # Overide
    def byte_patched(self, ea):
        if ea != idaapi.BADADDR:
            patched_byte = get_byte(ea)
            orginal_byte = get_original_byte(ea)
            file_offset = idaapi.get_fileregion_offset(ea)
            if file_offset != -1:
                patched_bytes_dict[ea] = [file_offset, orginal_byte, patched_byte]
            else:
                patched_bytes_dict[ea] = [-1, orginal_byte, patched_byte]
            print "[*] Byte patched at (ea %x file_offset %x) : %x -> %x" % (ea, file_offset, orginal_byte, patched_byte)
            
        return 0


#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class idapatcher_t(plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = "Enhances manipulation and application of patched bytes."
    help = "Enhances manipulation and application of patched bytes."
    wanted_name = "IDA Patcher"
    wanted_hotkey = ""

    def init(self):  
        global idapatcher_manager
        global patch_notify_hook

        print("[*] IDA Patcher v%s (c) Peter Kacherginsky <iphelix@thesprawl.org>" % IDAPATCHER_VERSION)
        print("[*] Initialize ...")
        # Check if already initialized
        if not 'idapatcher_manager' in globals():
            if not 'patch_notify_hook' in globals():
                idapatcher_manager = PatchManager()
                patch_notify_hook = PatchNotifyHook()
                patch_notify_hook.hook() # Install the hook
                if not idapatcher_manager.add_menu_items():
                    print "[!] Failed to initialize IDA Patcher."
                    idapatcher_manager.del_menu_items()
                    del idapatcher_manager
                    return idaapi.PLUGIN_SKIP
                else:
                    print("[*] IDA Patcher v%s (c) Peter Kacherginsky <iphelix@thesprawl.org>" % IDAPATCHER_VERSION)
        print("[*] Initialize Finished !")         
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        print("[*] IDA Patcher : Run ...")
        global idapatcher_manager

        idapatcher_manager.show_patches_view()

    # NOTE: term() gets called immediately after run()
    #       as the plugin is expected to complete execution
    #       so just pass as this plugin is expected to run
    #       for the duration of IDA session.
    def term(self):
        print("[*] IDA Patcher : Term ...")
        pass
        

def PLUGIN_ENTRY():
    return idapatcher_t()

#--------------------------------------------------------------------------
# Script / Testing
#--------------------------------------------------------------------------
def idapatcher_main():
    global idapatcher_manager

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

    idapatcher_manager = PatchManager()
    idapatcher_manager.add_menu_items()
    idapatcher_manager.show_patches_view()

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


repository: vAlerainTech/vAlerain-Ark
file: cmake-build-debug/IDA_Pro_7.7/plugins/keypatch.py
wanted_name: Keypatch Patcher
comment: Keypatch plugin for IDA Pro (using Keystone framework)
created_at: 2024-03-25 14:26:36+00:00
pushed_at: 2024-07-27 05:07:46+00:00
description:
🔭This will be my first long-term update project, and vAlerainArk is committed to helping security personnel with testing.

url:
https://github.com/vAlerainTech/vAlerain-Ark/blob/df8bcafeda091d30ceade7bb2c7c73b5ae351a38/cmake-build-debug/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: tmr232/BRUTAL-IDA
file: brutal_ida.py
description: Block Redo & Undo To Achieve Legacy IDA
url: https://github.com/tmr232/BRUTAL-IDA/blob/0e2c7fa9ae25d365922d255c7956ab89575fcfbb/brutal_ida.py
wanted_name: BRUTAL-IDA
comment: BRUTAL IDA
created_at: 2019-07-31 18:35:07+00:00
pushed_at: 2019-08-01 08:35:59+00:00
content:
'''
BRUTAL IDA

Block Redo & Undo To Achieve Legacy IDA
'''
import os

import idaapi


def get_brutal_icon_path(brutal_icon_name):
    return idaapi.load_custom_icon(os.path.join(os.path.dirname(__file__), 'BRUTAL-ICONS', brutal_icon_name))


BRUTAL5_ICON = get_brutal_icon_path('BRUTAL_IDA5.png')
BRUTAL6_ICON = get_brutal_icon_path('BRUTAL_IDA6.png')
BRUTAL_LETTERS = {
    brutal_letter: get_brutal_icon_path('BRUTAL_{}.png'.format(brutal_letter)) for brutal_letter in 'BRUTAL'
}


class BrutalLetterHandler(idaapi.action_handler_t):
    def __init__(brutal_self):
        idaapi.action_handler_t.__init__(brutal_self)

    def activate(brutal_self, brutal_context):
        pass

    def update(brutal_self, brutal_context):
        return idaapi.AST_DISABLE_ALWAYS


class BrutalActionHandler(idaapi.action_handler_t):
    def __init__(brutal_self):
        idaapi.action_handler_t.__init__(brutal_self)
        brutal_self.brutal_mode = '6.x'

    def activate(brutal_self, brutal_context):
        brutal_self.brutal_mode = {
            '5.x': '6.x',
            '6.x': '5.x',
        }[brutal_self.brutal_mode]

    @property
    def icon(brutal_self):
        return {
            '5.x': BRUTAL5_ICON,
            '6.x': BRUTAL6_ICON,
        }[brutal_self.brutal_mode]

    def update(brutal_self, brutal_context):
        idaapi.update_action_icon(brutal_context.action, brutal_self.icon)
        return idaapi.AST_ENABLE


class BrutalIDA(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = 'BRUTAL IDA'
    help = 'Block Redo & Undo To Achieve Legacy IDA'
    wanted_name = 'BRUTAL-IDA'
    wanted_hotkey = ''

    def handle_5x(brutal_self):
        idaapi.error('bTree error: Brutal. Just like the good ol\' days!')

    def handle_6x(brutal_self):
        pass

    def dispatch_brutality(brutal_self):
        {
            '5.x': brutal_self.handle_5x,
            '6.x': brutal_self.handle_6x,
        }[brutal_self.brutal_action_handler.brutal_mode]()

    def init(brutal_self):
        idaapi.unregister_action('Undo')
        idaapi.unregister_action('Redo')

        brutal_self.brutal_action_handler = BrutalActionHandler()
        brutal_action_desc = idaapi.action_desc_t('BRUTAL', 'BRUTAL IDA', brutal_self.brutal_action_handler, '',
                                                  'IDA', BRUTAL6_ICON)
        idaapi.register_action(brutal_action_desc)
        idaapi.create_toolbar('BRUTAL IDA', 'BRUTAL IDA')

        brutal_self.brutal_letter_handlers = []

        for brutal_letter in 'BRUTAL':
            brutal_letter_handler = BrutalLetterHandler()
            brutal_self.brutal_letter_handlers.append(brutal_letter_handler)

            brutal_label = 'BRUTAL {}'.format(brutal_letter)
            brutal_letter_desc = idaapi.action_desc_t(brutal_label,
                                                      brutal_label,
                                                      brutal_letter_handler,
                                                      '',
                                                      brutal_letter,
                                                      BRUTAL_LETTERS[brutal_letter])
            idaapi.register_action(brutal_letter_desc)
            idaapi.attach_action_to_toolbar('BRUTAL IDA', brutal_label)

        idaapi.attach_action_to_toolbar('BRUTAL IDA', 'BRUTAL')

        brutal_self.brutal_hotkey = idaapi.add_hotkey('Ctrl+Z', brutal_self.dispatch_brutality)

        return idaapi.PLUGIN_KEEP

    def term(brutal_self):
        pass

    def run(brutal_self, arg):
        pass


def PLUGIN_ENTRY():
    return BrutalIDA()


repository: JPCERTCC/GobRAT-Analysis
file: YaraGenerateTool/YaraGene_GobRAT_MultiArch.py
wanted_name: YaraGeneGobRAT
comment: YaraGenelate GobRAT for Multi Arch.
created_at: 2023-11-27 00:33:16+00:00
pushed_at: 2023-12-21 01:00:24+00:00
url:
https://github.com/JPCERTCC/GobRAT-Analysis/blob/064630f6c90894f1e2bd5b296f5116774e2e2fe3/YaraGenerateTool/YaraGene_GobRAT_MultiArch.py

content:
#---------------------------------------------------------------------------------

#---------------------------------------------------------------------------------
__date__ = "20231031"
__version__ = "1.0"
__author__ = "Yuma Masubuchi JPCERT/CC"
#---------------------------------------------------------------------------------
# Import module
#---------------------------------------------------------------------------------
from idautils import *
from idaapi import *
from idc import *
import idaapi
import ida_kernwin
import ida_nalt
import os
import re
#---------------------------------------------------------------------------------

#---------------------------------------------------------------------------------
# Config
#---------------------------------------------------------------------------------
class ConfigVar:
    """
    A number that determines whethre the hunt is acceptable or not based on how many of
    the byte sequences created match.
    """
    matchNumber = 6
    
    """
    To hit the various libraries, this tool extracts a sequence of bytes from among all
    functions divided into quarters.A number determines how many byte sequences to use
    for each quarter.
    ie 
    """
    NumFunc_eachPart = 2
    
    """
    In yara conditions, the file size to be hit is detemined in Mega Bytes.
    """
    LessthanSize_MB = 10
    GreatorthanSize_MB = 2

    """
    Output filename
    """
    YaraFileName = "malware_GobRAT.yara"

    """
    Other
    """
    Other_RuleName = "GobRAT_arch"
    Other_Description = "Hunting GobRAT malware"
    Other_Auther = "JPCERT/CC Incident Response Group"
    Other_Usage = "Hunting"

    # On command prompt by use of idat.exe or idat64.exe
    BATCH_MODE = False

    # Target only Sub_* Functions
    TARGET_SUBFUNC = True
#---------------------------------------------------------------------------------
  
#---------------------------------------------------------------------------------


class YaraCompo:
    GENERATED_YARA = ""
    TEMPLATE = '''

import "[MODULE]"

rule [RULENAME] {
    meta:
        description = "[DESCRIPTION]"
        author = "[AUTHOR]"
        rule_usage = "[USAGE]"
        hash = "[HASH]"
    
    strings:
[STRINGS]
    condition:
        [CONDITION]
}
    '''
    MODULE = ""
    #RULENAME = ConfigVar.Other_RuleName
    DESCRIPTION = ConfigVar.Other_Description
    AUTHOR = ConfigVar.Other_Auther
    USAGE = ConfigVar.Other_Usage
    HASH = ""
    STRINGS = ""
    CONDITION = ""

    def set_RULENAME(self, in_Rulename):
        self.RULENAME = in_Rulename

    def set_MODULE(self, in_module):
        self.MODULE = in_module

    def set_HASH(self, in_hash):
        self.HASH = in_hash

    def set_STRINGS(self, in_strings):
        self.STRINGS = in_strings

    def set_CONDITION(self, in_condition):
        self.CONDITION = in_condition

    def print_all_compo(self):
        print("[*] INFO MODULE: {}".format(self.MODULE))
        print("[*] INFO RULENAME: {}".format(self.RULENAME))
        print("[*] INFO DESCRIPTION: {}".format(self.DESCRIPTION))
        print("[*] INFO AUTHOR: {}".format(self.AUTHOR))
        print("[*] INFO USAGE: {}".format(self.USAGE))
        print("[*] INFO HASH: {}".format(self.HASH))
        print("[*] INFO STRINGS:\n{}".format(self.STRINGS))
        print("[*] INFO CONDITION:\n{}".format(self.CONDITION))

    def print_generated_yara_text(self):
        print("[*] GENERATED_YARA:\n{}".format(self.GENERATED_YARA))

    def get_generated_yara_text(self):
        return self.GENERATED_YARA

    def generate_yara_text(self):
        self.GENERATED_YARA = self.TEMPLATE 
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[MODULE]", self.MODULE)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[RULENAME]", self.RULENAME)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[DESCRIPTION]", self.DESCRIPTION)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[AUTHOR]", self.AUTHOR)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[USAGE]", self.USAGE)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[HASH]", self.HASH)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[STRINGS]", self.STRINGS)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[MODULE]", self.MODULE)
        self.GENERATED_YARA = self.GENERATED_YARA.replace("[CONDITION]", self.CONDITION)
         
   
class YaraModule():
    import_module = ""
    
    def set_import_module(self):
        fileformat = idaapi.get_file_type_name()
        if "ELF" in fileformat:
            self.import_module = "elf"
        elif "PE" in fileformat:
            self.import_module = "pe"
        else:
            self.import_module == "elf" # default
            print("[*] Warning: The file is not PE or ELF format. please check the file format.")
        
    def get_import_module(self):
        return self.import_module
        
        
class YaraHash():
    sha256hash = ""
    
    def Convert_str_from_Hex_bytes(self, bytes):
        for byte in bytes:
            self.sha256hash += '{:02X}'.format(byte)
    
    def set_hash_sha256(self):
        sha256_bytes = ida_nalt.retrieve_input_file_sha256()
        if sha256_bytes != None:
            self.Convert_str_from_Hex_bytes(sha256_bytes)
        else:
            self.sha256hash = "NoFile"

    def get_hash_sha256(self):
        return self.sha256hash


class IDAPythoOperandCtrl:
    opTypeDict = { 0:"NoOp",  # No Operand
                   1:"GER",   # General Register 
                   2:"DATA",  # Direct Memory Reference (DATA)
                   3:"MEM",   # Memory Ref [Base Reg + Index Reg]
                   4:"VAR",   # Memory Reg [Base Reg + Index Reg + Displacement]
                   5:"IMM",   # Immediate Value
                   6:"FADDR", # Immediate Far Address (CODE)
                   7:"NADDR", # Immediate Near Address (CODE)
                   8:"PSY1",  # processor specific type
                   9:"PSY2",  # processor specific type
                   10:"PSY3", # processor specific type
                   11:"PSY4", # processor specific type
                   12:"PSY5", # processor specific type
                   13:"PSY6", # processor specific type
                 }

    ReplaceNumOperand = [2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13]


class YaraStrings:
    yara_strings = ''

    instructions_offset = 2 
    convert_flag_offset = 0
    inst_length_offset = 1
    bytes_offset = 2
    instruction_offset = 3
    function_all_info = []
    extracted_most_long_inst = []
    yara_bytes = []

    sp8 = ' ' * 8
    func_comment01 = sp8 + '/* Function Address: '
    func_comment02 = sp8 + '*/\n'
    func_bytes01 = sp8 + '$func'
    func_bytes02 = ' = { '
    func_bytes03 = '}\n\n'
 
    def Convert_Hex_bytes_to_Yaraformat(self, bytes):
        ret_str = ''
        for byte in bytes:
            ret_str += '{:02X} '.format(byte)
        return ret_str
    
    def inst_extract(self, funcea):    
        ret_extracted = []
        f = ida_funcs.get_func(funcea)

        for ea in Heads(f.start_ea, f.end_ea):
            insn = idaapi.insn_t()
            # get length 
            length = decode_insn(insn, ea)
            # get bytes(strings)
            inst_bytes = get_bytes(ea, length)   
            if inst_bytes == None:
                break
            inst_bytes_str = self.Convert_Hex_bytes_to_Yaraformat(inst_bytes)            
            # get instruction
            re_inst = generate_disasm_line(ea, 0)

            # get convert instruction
            convert_flag = 0 
            re_inst_conv = re_inst
            for opNo in range(0, 3):
                if ( get_operand_type(ea, opNo) in IDAPythoOperandCtrl.ReplaceNumOperand ):
                    re_inst_conv = re_inst.replace( print_operand(ea, opNo),
                                            IDAPythoOperandCtrl.opTypeDict[get_operand_type(ea, opNo)],
                                            1)
                    convert_flag = 1
            ret_extracted.append([convert_flag, length, inst_bytes_str, re_inst, re_inst_conv])
        return ret_extracted
    
    def get_all_function_info(self):
        for segea in Segments():
            for funcea in Functions(segea, get_segm_end(segea)):
                functionName = get_func_name(funcea)

                # Target functions are sub_* 
                #if ((functionName.startswith("main")) == True) or ((functionName.startswith("aaa_com_bbb")) == True):
                if ((functionName.startswith("aaa.com_bbb")) == True):

                    # get a flag status on functions to exclude library functions for making yara rule
                    flags = idc.get_func_flags(funcea)
                    if (flags & FUNC_LIB) == 0:
                        function_info = []
                        function_info = self.inst_extract(funcea)

                        self.function_all_info.append(["0x{:02x}".format(funcea), functionName, function_info])
 
    def extract_long_inst4yara(self):        
        numberOfFunc = len(self.function_all_info)
        for fnum in range(numberOfFunc):
            noConvINst_count = max_noConv_count = max_inum = 0
            
            numberOfInst = len(self.function_all_info[fnum][self.instructions_offset])
            
            # extract most long instrunction length and the inum into the function
            current_inum_head = current_len = max_len = max_inum_head = 0
            for inum in range(numberOfInst):
                convert_flag = self.function_all_info[fnum][self.instructions_offset][inum][self.convert_flag_offset]
                inst_length = self.function_all_info[fnum][self.instructions_offset][inum][self.inst_length_offset]
                if convert_flag == 0:
                    current_len += inst_length
                    if max_len < current_len:
                        max_len = current_len
                        max_inum_head = current_inum_head
                elif convert_flag == 1:
                    current_inum_head = inum + 1
                    current_len = 0
                else:
                    print("[*] Error: extract_long_inst4yara()")
                    return             
            self.extracted_most_long_inst.append([fnum, max_inum_head, max_len])    
    
    def pickup_yara_bytes(self): 
        split_len = len(self.extracted_most_long_inst) // 4
        
        part1 = self.extracted_most_long_inst[:split_len]
        part2 = self.extracted_most_long_inst[split_len:2*split_len]
        part3 = self.extracted_most_long_inst[2*split_len:3*split_len]
        part4 = self.extracted_most_long_inst[3*split_len:]
        
        part1.sort(key=lambda x:x[2], reverse=True)   
        part2.sort(key=lambda x:x[2], reverse=True)   
        part3.sort(key=lambda x:x[2], reverse=True)
        part4.sort(key=lambda x:x[2], reverse=True)   
            
        concat_allpart = [part1[:ConfigVar.NumFunc_eachPart], part2[:ConfigVar.NumFunc_eachPart], part3[:ConfigVar.NumFunc_eachPart], part4[:ConfigVar.NumFunc_eachPart]]
        concat_allpart = sum(concat_allpart, [])
        
        for long_funcinfo in concat_allpart:
            fnum = long_funcinfo[0]
            inum = long_funcinfo[1]
            
            bytes_strings = ""
            instructions_strings = ""
        
            funcAddr = self.function_all_info[fnum][0]
            funcName = self.function_all_info[fnum][1]

            func_Addr_and_Name = funcAddr +" : "+ funcName
            for n in range(300):
                try:
                    convert_flag = self.function_all_info[fnum][self.instructions_offset][inum + n][self.convert_flag_offset]
                except:
                    break
                if convert_flag == 1:
                    break
    
                current_bytes = self.function_all_info[fnum][self.instructions_offset][inum + n][self.bytes_offset]
                bytes_strings += current_bytes
                instructions_strings += self.sp8 + "{:36}".format(current_bytes) + "{:30}".format(self.function_all_info[fnum][self.instructions_offset][inum + n][self.instruction_offset]) + "\n"

            # exclude same bytes_strings in the list(yara_bytes)
            sameFlag = True
            for ele in self.yara_bytes:
                if bytes_strings == ele[1]:
                    sameFlag = False
            if sameFlag:
                self.yara_bytes.append([func_Addr_and_Name, bytes_strings, instructions_strings])
    
    
    def print_main_list(self):
        print("[*] INFO function_all_info[] : {}".format(self.function_all_info))
        print("[*] INFO extracted_most_long_inst[] : {}".format(self.extracted_most_long_inst))
        print("[*] INFO yara_bytes[] : {}".format(self.yara_bytes))
        ### -- Following is each data structure -- 
        ### function_all_info = [
        ###     ['0x804dce9', '.term_proc',   ### Function-address, Function-name ### 
        ###         [                           ### Convert-flag, Bytes-length, Instruction ### 
        ###             [0, 3, '83 EC 0C ', 'sub     esp, 0Ch', 'sub     esp, 0Ch'],
        ###             [1, 5, 'E8 BF A3 FF FF ', 'call    sub_80480B0', 'call    NADDR'],
        ###             [0, 3, '83 C4 0C ', 'add     esp, 0Ch', 'add     esp, 0Ch'],
        ###             [0, 1, 'C3 ', 'retn', 'retn']
        ###         ]
        ###     ], ...
        ### ]        
        ### extracted_most_long_inst = [ [ Function-Number, max-Instruction-startPoint, max-Instruction-length ], ... ]
        ### yara_bytes = [ [ Function-Address, FunctionBytes, FunctionInstructions ], ... ]
        

    def generate_yara_strings(self):
        # set each function bytes
        numberOfFunction = len(self.yara_bytes)
        for num in range(numberOfFunction):
            self.yara_strings += self.func_comment01
            # function addr
            self.yara_strings += self.yara_bytes[num][0] + "\n"
            # function instructions
            self.yara_strings += self.yara_bytes[num][2]
            self.yara_strings += self.func_comment02
            self.yara_strings += self.func_bytes01
            self.yara_strings += str(num)
            self.yara_strings += self.func_bytes02
            # funciton bytes
            self.yara_strings += self.yara_bytes[num][1]
            self.yara_strings += self.func_bytes03

    def get_yara_strings(self):
        return self.yara_strings


class YaraCondition():
    yara_condition = ""
    file_format = ""
    sp8 = ' ' * 8

    filefomrmat_elf = '(uint32(0) == 0x464C457F)\n'
    filefomrmat_dos = '(uint16(0) == 0x5A4D)\n'

    elf_x64 = sp8+ 'and (elf.machine == elf.EM_X86_64)\n'
    elf_x86 = sp8+ 'and (elf.machine == elf.EM_386)\n'
    elf_ARM = sp8+ 'and (elf.machine == elf.EM_ARM)\n'
    elf_MIPS = sp8+ 'and (elf.machine == elf.EM_MIPS)\n'
    elf_arch = ""

    con_filesize = sp8+ 'and (filesize < ' + str(ConfigVar.LessthanSize_MB) + 'MB)\n'
    con_filesize2 = sp8+ 'and (filesize > ' + str(ConfigVar.GreatorthanSize_MB) + 'MB)\n'
    func_number01 = sp8 + 'and ( '
    match_numberOfFunc = ""
    func_number02 = ' of ($func*) )\n'
    
    def set_file_format(self):
        file_t = idaapi.get_file_type_name()
        if "ELF" in file_t:
            self.file_format = "ELF"
        elif "PE" in file_t:
            self.file_format = "PE"
        else:
            self.file_format == "None"

    def set_condition_arch(self):
        info = get_inf_structure()
        self.proc_name = info.procname
        if self.proc_name == "metapc":
            if info.is_64bit() == True:
                self.elf_arch = self.elf_x64
            elif info.is_32bit() == True:
                self.elf_arch = self.elf_x86
            else:
                print("[*] Error: cpu architecture is not found()")
        elif self.proc_name == "ARM":
            self.elf_arch = self.elf_ARM
        elif self.proc_name == "mipsb":
            self.elf_arch = self.elf_MIPS
        else:
            print("[*] Error: cpu architecture is not found. detect cpu: {}".format(info.procname))

    def set_yara_condition_matchNumber(self):
        self.match_numberOfFunc = str(ConfigVar.matchNumber)

    def generate_yara_condition(self): 
        if self.file_format == "ELF":
            self.yara_condition += self.filefomrmat_elf
            self.yara_condition += self.elf_arch
            self.yara_condition += self.con_filesize
            self.yara_condition += self.con_filesize2
        elif self.file_format == "PE":
            self.yara_condition += self.filefomrmat_dos
            self.yara_condition += self.con_filesize
        else:
            self.yara_condition += self.con_filesize
            
        self.yara_condition += self.func_number01
        self.yara_condition += self.match_numberOfFunc
        self.yara_condition += self.func_number02

    def get_yara_condition(self):
        return self.yara_condition
    

class OutputYaraFile:
    yarafilename = ConfigVar.YaraFileName
    currentdir = ""
    filepath = ""
    yara_strings = ""

    def __init__(self):
        self.currentdir = os.getcwd()
        self.filepath = self.currentdir + "\\" + self.yarafilename

    def reset_filepath(self, input_text):
        self.filepath = self.currentdir + "\\" + input_text

    def set_yara_text(self, input_text):
        self.yara_strings = input_text
        
    def get_filepath(self):
        return self.filepath

    def output_yarafile(self):
        textfile = open(self.filepath, "w")
        textfile.write(self.yara_strings)
        textfile.close()
        print("[*] Generated Yara File! : {}".format(self.filepath))


class SimpleSubViewer(simplecustviewer_t):
    filepath = ""

    def set_filepath(self, in_filepath):
        self.filepath = in_filepath

    def Create(self):
        title = "GobRAT Yara Rule"
        simplecustviewer_t.Create(self, title)
        yarafile = open(self.filepath, "r")
        yaraLines = yarafile.readlines()
        yarafile.close()
        
        for line in yaraLines:
            self.AddLine(line)
        self.Show()

    def OnDblClick(self, value):
        curWord = self.GetCurrentWord()
        lineN = self.GetLineNo()
        if not curWord:
            curWord = "<None>"
        elif curWord[:2] == '0x':
            jumpto(int(curWord, 16))
            
        return True
    
    def OnKeydown(self, vkey, shift):
        if vkey == ord('E'):
            lineN = self.GetLineNo()
            if lineN is not None:
                lineS = self.GetCurrentLine(notags=1)
                lineSnew = ida_kernwin.ask_str("{}".format(lineS), 0, "Edit")
                self.edit_line(lineN, lineS, lineSnew)                
                
        elif vkey == ord('I'):
            lineN = self.GetLineNo()
            lineS = self.GetCurrentLine(notags=1)
            if lineN is not None:
                self.insert_line(lineN, lineS)
                            
        elif vkey == 46: # Delete
            lineN = self.GetLineNo()
            lineS = self.GetCurrentLine(notags=1)
            if lineN is not None:
                yn = True #ida_kernwin.ask_yn(0, "Delete line?")
                if yn == True:
                    self.delete_line(lineN, lineS)
                    print("[*] INFO delete >> {}".format(lineS))

        elif vkey == 27: # Escape
            print("[*] INFO window Quit")
            self.Close()

        return True

    def save_file(self, lines):
        yara_text = ""
        for line in lines:
            yara_text += line
        textfile = open(self.filepath, "w")
        textfile.write(yara_text)
        textfile.close()
        print("[*] Save Yara File! : {}".format(self.filepath))
        
    def refresh_file_contents(self):
        yaraLines = self.get_yarafile_lines()
        self.ClearLines()
        
        for line in yaraLines:
            self.AddLine(line)
        self.Refresh()
       
    def get_yarafile_lines(self):
        yarafile = open(self.filepath, "r")
        yaraLines = yarafile.readlines()
        yarafile.close()
        return yaraLines
        
    def edit_line(self, lineNum, oldcontents, newcontents):
        yaraLines = self.get_yarafile_lines()
        # check file contetnts
        if yaraLines[lineNum] == oldcontents:
            newcontents = newcontents.replace("\n", "")
            yaraLines[lineNum] = newcontents + "\n"
            self.save_file(yaraLines)    
            self.refresh_file_contents()

    def insert_line(self, lineNum, oldcontents):
        yaraLines = self.get_yarafile_lines()
        # check file contetnts
        if yaraLines[lineNum] == oldcontents:
            yaraLines[lineNum] = yaraLines[lineNum] + "\n"
            self.save_file(yaraLines)    
            self.refresh_file_contents()
            
    def delete_line(self, lineNum, oldcontents):
        yaraLines = self.get_yarafile_lines()       
        # check file contetnts
        if yaraLines[lineNum] == oldcontents:
            yaraLines[lineNum] = ""
            self.save_file(yaraLines)    
            self.refresh_file_contents()    
    
    def OnHint(self, lineno):
        return (6, "  E  : edit line\n  I  : insert line\n[Del]: delete line\n[ESC]: quit")


def main():
    print("[*] Start!!")
    # batch mode
    if ConfigVar.BATCH_MODE:
        ida_auto.auto_wait()

    # Generate yaraCompo Class
    YaraCompoClass = YaraCompo()
    
    # Set MODULE
    YaraModuleClass = YaraModule()
    YaraModuleClass.set_import_module()
    module_t = YaraModuleClass.get_import_module()
    YaraCompoClass.set_MODULE(module_t)

    # Set HASH
    YaraHashClass = YaraHash()
    YaraHashClass.set_hash_sha256()
    hash_t = YaraHashClass.get_hash_sha256()
    YaraCompoClass.set_HASH(hash_t)

    # add!   batch mode
    YaraCompoClass.set_RULENAME("RULE_" + hash_t)
    
    # Set Strings
    YaraStringsClass = YaraStrings()
    YaraStringsClass.get_all_function_info()
    YaraStringsClass.extract_long_inst4yara()
    YaraStringsClass.pickup_yara_bytes()
    YaraStringsClass.generate_yara_strings()
    strings_t = YaraStringsClass.get_yara_strings()
    YaraCompoClass.set_STRINGS(strings_t)
    
    # Set Condition
    YaraConditionClass = YaraCondition()
    YaraConditionClass.set_file_format()
    YaraConditionClass.set_condition_arch()
    YaraConditionClass.set_yara_condition_matchNumber()
    YaraConditionClass.generate_yara_condition()
    condition_t = YaraConditionClass.get_yara_condition()
    YaraCompoClass.set_CONDITION(condition_t)

    # Generate yara text data
    YaraCompoClass.generate_yara_text()
    text_t = YaraCompoClass.get_generated_yara_text()

    
    # Output .yara file
    OutputYaraFileClass = OutputYaraFile()
     # add!   batch mode
    OutputYaraFileClass.reset_filepath("RULE_" + hash_t + ".yara")
    OutputYaraFileClass.set_yara_text(text_t)
    OutputYaraFileClass.output_yarafile()
    
    # Open Subviewer on IDA
    SimpleSubViewerClass = SimpleSubViewer()
    filepath_t = OutputYaraFileClass.get_filepath()
    SimpleSubViewerClass.set_filepath(filepath_t)
    SimpleSubViewerClass.Create()
    
    print("[*] Finish!!")

    # batch mode
    if ConfigVar.BATCH_MODE:
        idc.exit()


if __name__ == '__main__': 
    main()


class YaraGeneGobRAT_Plugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_DRAW
    comment = "YaraGenelate GobRAT for Multi Arch."
    help = "YaraGeneGobRAT"
    wanted_name = "YaraGeneGobRAT"
    wanted_hotkey = "" # = "Alt-F10"

    def init(self):
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        main()
        return

    def term(self):
        pass


def PLUGIN_ENTRY():
    return YaraGeneGobRAT_Plugin()


repository: 19h/ida-plugins
file: exportdecomp/exportdecomp.py
wanted_name: Export Decompiler
comment: Decompiles all exported functions to C
created_at: 2025-04-08 22:54:27+00:00
pushed_at: 2025-04-08 23:07:17+00:00
url:
https://github.com/19h/ida-plugins/blob/b045b6f6bc9ae2a53104986fb74f3fe1198e5d2e/exportdecomp/exportdecomp.py

content:
import idaapi
import idautils
import ida_hexrays
import ida_funcs
import ida_name
import ida_kernwin
import ida_typeinf
import ida_bytes
import ida_segment
import ida_auto
import ida_pro
import os
import ida_xref
import ida_entry # Use ida_entry for modern export handling
import idc       # For idc.GetDisasm and other legacy functions if needed
import traceback # For detailed error reporting

# Helper function to get a descriptive name for xref types
# (ida_xref.get_xref_type_name might not exist in all versions, provide fallback)
def get_xref_type_name_fallback(xtype):
    """Provides a string representation for an xref type."""
    xref_types = {
        ida_xref.fl_CF: "Call Far", ida_xref.fl_CN: "Call Near",
        ida_xref.fl_JF: "Jump Far", ida_xref.fl_JN: "Jump Near",
        ida_xref.fl_F:  "Ordinary Flow", ida_xref.dr_R:  "Data Read",
        ida_xref.dr_W:  "Data Write", ida_xref.dr_O:  "Data Offset",
        ida_xref.dr_I:  "Data Info",
    }
    base_type = xtype & ida_xref.XREF_MASK
    name = xref_types.get(base_type, f"Unknown ({base_type:#x})")
    if xtype & ida_xref.XREF_USER: name += " (User)"
    return name

# Use the official function if available, otherwise the fallback
if hasattr(ida_xref, 'get_xref_type_name'):
    get_xref_type_name_ida = ida_xref.get_xref_type_name
else:
    get_xref_type_name_ida = get_xref_type_name_fallback


class ExportDecompiler(idaapi.plugin_t):
    """
    IDA Pro plugin for decompiling all exported functions to C.
    Utilizes Hex-Rays decompiler to extract accurate source representations.
    Correctly enumerates exports using ida_entry.
    """
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_MOD
    comment = "Decompiles all exported functions to C"
    help = "Generates C file with decompiled exported functions"
    wanted_name = "Export Decompiler"
    wanted_hotkey = "Alt-Shift-E"

    def init(self):
        """
        Plugin initialization.
        Checks for Hex-Rays should be done in run(), not here.
        """
        print("[Export Decompiler] Plugin initialized (check for Hex-Rays deferred to run()).")
        # Simply indicate the plugin is ready to be loaded.
        # The actual check for Hex-Rays usability happens in run().
        return idaapi.PLUGIN_OK

    def run(self, arg):
        """Main execution function - orchestrates decompilation workflow."""
        print("[Export Decompiler] Plugin execution started.")

        # --- Check for Hex-Rays availability *now* ---
        hexrays_available = False
        try:
            # Check if Hex-Rays API is responsive
            if ida_hexrays.get_hexrays_version():
                # Try initializing it specifically for this run if needed (might not be necessary if already loaded)
                if ida_hexrays.init_hexrays_plugin():
                     hexrays_available = True
                else:
                     # It might already be initialized, check if API works anyway
                     try:
                         ida_hexrays.decompile # Access a common function
                         hexrays_available = True
                     except AttributeError:
                         print("[Export Decompiler] Hex-Rays init returned false and core functions missing.")
            else:
                 print("[Export Decompiler] Hex-Rays get_hexrays_version() failed.")
        except Exception as e:
            print(f"[Export Decompiler] Error during Hex-Rays runtime check: {e}")

        if not hexrays_available:
            ida_kernwin.warning("Hex-Rays decompiler is not available or failed to initialize.")
            print("[Export Decompiler] Hex-Rays unavailable at runtime.")
            return
        # --- End Hex-Rays Check ---

        # Ensure autoanalysis completion
        if ida_auto.get_auto_state() != ida_auto.AU_FINAL:
            print("[Export Decompiler] Autoanalysis is not complete.")
            if ida_kernwin.ask_yn(1, "IDA's autoanalysis has not finished yet.\n"
                                   "The decompilation results might be incomplete or inaccurate.\n\n"
                                   "Continue anyway?") != 1:
                print("[Export Decompiler] User cancelled due to incomplete autoanalysis.")
                return
            print("[Export Decompiler] Continuing despite incomplete autoanalysis.")

        # Get output file path
        default_filename = idaapi.get_root_filename() + "_exports.c"
        output_path = ida_kernwin.ask_file(1, default_filename, "Select output C file")
        if not output_path:
            print("[Export Decompiler] User cancelled file selection.")
            return
        print(f"[Export Decompiler] Output file selected: {output_path}")

        # Configure options
        options = self._get_options()
        if not options:
            print("[Export Decompiler] User cancelled options configuration.")
            return
        print(f"[Export Decompiler] Options selected: {options}")

        # Process exports
        self._process_exports(output_path, options)
        print("[Export Decompiler] Plugin execution finished.")

    def term(self):
        """Plugin termination function."""
        # No explicit Hex-Rays termination usually needed here either.
        print("[Export Decompiler] Plugin terminated.")
        pass

    def _get_options(self):
        """Configure decompilation options via dialog."""
        # Define control variables using ChkGroupControl.
        # The tuple contains the *internal name* of the checkbox within the group.
        typeinclude_ctrl = ida_kernwin.Form.ChkGroupControl(("typeinclude_cb",)) # Use unique names for clarity
        xrefs_ctrl = ida_kernwin.Form.ChkGroupControl(("xrefs_cb",))
        asm_ctrl = ida_kernwin.Form.ChkGroupControl(("asm_cb",))
        guards_ctrl = ida_kernwin.Form.ChkGroupControl(("guards_cb",))
        paramnames_ctrl = ida_kernwin.Form.ChkGroupControl(("paramnames_cb",))
        optimize_ctrl = ida_kernwin.Form.ChkGroupControl(("optimize_cb",))

        # The names in {} MUST match the dictionary keys passed to the Form constructor.
        form_definition = r"""BUTTON YES* OK
BUTTON CANCEL Cancel
Export Decompiler Configuration

Select features to include in the output C file:

<##Include type definitions (structs, enums, etc.):{typeinclude_grp}>
<##Include cross-reference comments:{xrefs_grp}>
<##Include assembly listing comments:{asm_grp}>
<##Generate header guards (#ifndef...):{guards_grp}>
<##Attempt parameter name recovery (Hex-Rays often does this):{paramnames_grp}>
<##Apply extra optimization passes (Experimental):{optimize_grp}>
"""
        # The dictionary keys MUST match the names inside the {} in the form definition.
        # These keys link the placeholder in the form string to the control object.
        controls = {
            'typeinclude_grp': typeinclude_ctrl,
            'xrefs_grp': xrefs_ctrl,
            'asm_grp': asm_ctrl,
            'guards_grp': guards_ctrl,
            'paramnames_grp': paramnames_ctrl,
            'optimize_grp': optimize_ctrl,
        }

        form = ida_kernwin.Form(form_definition, controls)
        form.Compile()

        # --- Set default values AFTER Compile() ---
        # Access the individual checkbox controls using the *internal names*
        # provided in the ChkGroupControl tuples as attributes of the form object.
        form.typeinclude_cb.checked = True
        form.xrefs_cb.checked = True
        form.asm_cb.checked = False # Default off
        form.guards_cb.checked = True
        form.paramnames_cb.checked = True
        form.optimize_cb.checked = False # Default off

        # --- Execute the form ---
        ok = form.Execute()
        options = None
        if ok == 1:
            # --- Read values back AFTER Execute() ---
            # Read from the individual checkbox controls' .checked attribute.
            options = {
                'include_types': form.typeinclude_cb.checked,
                'include_xrefs': form.xrefs_cb.checked,
                'include_asm': form.asm_cb.checked,
                'use_guards': form.guards_cb.checked,
                'recover_params': form.paramnames_cb.checked,
                'optimize': form.optimize_cb.checked
            }

        form.Free()
        return options

    def _get_header_guard(self, filename):
        """Generate C preprocessor header guard from filename."""
        base = os.path.basename(filename)
        name = os.path.splitext(base)[0]
        guard = ''.join(c if c.isalnum() else '_' for c in name).upper()
        if not guard or not (guard[0].isalpha() or guard[0] == '_'):
            guard = '_' + guard
        guard = guard.strip('_')
        while "__" in guard: guard = guard.replace('__', '_')
        if not guard: guard = "GENERATED_EXPORTS_H"
        else: guard += "_EXPORTS_H"
        return guard

    def _get_asm_listing(self, func_start, func_end):
        """Extract assembly listing for function."""
        result = []
        addr = func_start
        while addr != idaapi.BADADDR and addr < func_end:
            disasm = idc.GetDisasm(addr)
            if disasm is None: disasm = "; Failed to get disassembly"

            instr = idaapi.insn_t()
            try:
                instr_len = idaapi.decode_insn(instr, addr)
                if instr_len <= 0:
                     bytes_str = ""
                     result.append(f"/* 0x{addr:08X}: {'':<15} {disasm} (decode error or zero length) */")
                     addr += 1;
                     if addr >= func_end: break
                     continue
            except Exception as e:
                 print(f"[Export Decompiler] Error decoding instruction at 0x{addr:X}: {e}")
                 bytes_str = ""; instr_len = 1 # Guess
                 result.append(f"/* 0x{addr:08X}: Error decoding instruction: {e} */")

            bytes_data = ida_bytes.get_bytes(addr, instr_len) if instr_len > 0 else None
            bytes_str = ' '.join([f"{b:02X}" for b in bytes_data]) if bytes_data else ""
            result.append(f"/* 0x{addr:08X}: {bytes_str:<15} {disasm} */")

            next_addr = ida_bytes.next_head(addr, func_end)
            if next_addr <= addr:
                print(f"[Export Decompiler] Warning: next_head failed to advance from 0x{addr:X}. Stopping assembly listing.")
                break
            addr = next_addr
        return '\n'.join(result)

    def _get_custom_types(self):
        """Extract all custom types (structs, enums, unions, typedefs) defined in IDB."""
        types = {}
        ti = ida_typeinf.get_idati()
        if not ti: print("[Export Decompiler] Error: Could not get local TIL."); return types

        try: ordinal_qty = ida_typeinf.get_ordinal_qty(ti)
        except AttributeError: print("[Export Decompiler] Warning: get_ordinal_qty not found. Limiting type scan."); ordinal_qty = 65536

        print(f"[Export Decompiler] Scanning up to {ordinal_qty-1} potential local types.")
        processed_ordinals = set()

        for ordinal in range(1, ordinal_qty):
            if ordinal in processed_ordinals: continue

            tinfo = ida_typeinf.tinfo_t(); success = False
            try:
                 py_obj = ida_typeinf.get_numbered_type(ti, ordinal)
                 if py_obj and isinstance(py_obj, tuple) and len(py_obj) >= 2 and isinstance(py_obj[1], ida_typeinf.tinfo_t):
                     tinfo = py_obj[1]; success = True
                 elif py_obj and isinstance(py_obj, ida_typeinf.tinfo_t):
                     tinfo = py_obj; success = True
            except Exception as e:
                 if ordinal_qty == 65536 and ordinal > 1: print(f"Stopping type scan at {ordinal}: {e}"); break
                 success = False

            if success:
                original_ordinal = ordinal; is_alias = False
                target_ordinal = ida_typeinf.get_alias_target(ti, ordinal); alias_name = None

                if target_ordinal != 0 and target_ordinal != ordinal:
                    is_alias = True; alias_name = ida_typeinf.get_numbered_type_name(ti, ordinal)
                    ordinal = target_ordinal; processed_ordinals.add(original_ordinal)
                    success_target = False
                    try:
                         py_obj_target = ida_typeinf.get_numbered_type(ti, ordinal)
                         if py_obj_target and isinstance(py_obj_target, tuple) and len(py_obj_target) >= 2 and isinstance(py_obj_target[1], ida_typeinf.tinfo_t):
                             tinfo = py_obj_target[1]; success_target = True
                         elif py_obj_target and isinstance(py_obj_target, ida_typeinf.tinfo_t):
                             tinfo = py_obj_target; success_target = True
                    except: pass
                    if not success_target: print(f"[Export Decompiler] Warning: Failed get alias target {ordinal}"); continue

                if ordinal in processed_ordinals:
                    if is_alias and alias_name:
                        target_name = ida_typeinf.get_numbered_type_name(ti, ordinal) or tinfo.dstr()
                        if target_name and alias_name not in types:
                            if target_name not in types:
                                target_decl_flags = (ida_typeinf.PRTYPE_MULTI | ida_typeinf.PRTYPE_TYPE | ida_typeinf.PRTYPE_SEMI | ida_typeinf.PRTYPE_DEF)
                                target_type_decl = ida_typeinf.print_tinfo(None, 0, 0, target_decl_flags, tinfo, target_name, None)
                                if target_type_decl: types[target_name] = target_type_decl
                            if target_name in types or target_name: types[alias_name] = f"typedef {target_name} {alias_name};"
                    continue

                type_name = ida_typeinf.get_numbered_type_name(ti, ordinal)
                if not type_name:
                    if tinfo.is_struct() or tinfo.is_union() or tinfo.is_enum():
                       type_name = tinfo.dstr();
                       if not type_name: type_name = f"__unnamed_type_{ordinal}"
                    else: continue

                if type_name in types: processed_ordinals.add(ordinal); continue

                decl_flags = (ida_typeinf.PRTYPE_MULTI | ida_typeinf.PRTYPE_TYPE | ida_typeinf.PRTYPE_SEMI | ida_typeinf.PRTYPE_DEF)
                type_decl = ida_typeinf.print_tinfo(None, 0, 0, decl_flags, tinfo, type_name, None)

                if type_decl:
                    if is_alias and alias_name:
                        types[type_name] = type_decl; types[alias_name] = f"typedef {type_name} {alias_name};"
                    else: types[type_name] = type_decl
                else: print(f"[Export Decompiler] Warning: Could not format type '{type_name}' (ordinal {ordinal})")
                processed_ordinals.add(ordinal)

        print(f"[Export Decompiler] Extracted {len(types)} unique custom type declarations.")
        return dict(sorted(types.items()))

    def _get_xrefs(self, func_addr):
        """Extract cross-references *from* function instructions."""
        xrefs_list = []; func = ida_funcs.get_func(func_addr)
        if not func: return xrefs_list
        fii = ida_funcs.func_item_iterator_t(func); ok = fii.first()
        while ok:
            addr = fii.current()
            try:
                for xref in idautils.XrefsFrom(addr, ida_xref.XREF_ALL): # type: ignore
                    target_addr = xref.to
                    if target_addr == idaapi.BADADDR: continue
                    target_name = ida_name.get_ea_name(target_addr, ida_name.GN_VISIBLE)
                    if not target_name:
                        seg = ida_segment.getseg(target_addr)
                        target_name = f"{ida_segment.get_segm_name(seg)}:0x{target_addr:X}" if seg else f"0x{target_addr:X}"
                    xref_type_str = get_xref_type_name_ida(xref.type)
                    xrefs_list.append((addr, target_addr, target_name, xref_type_str))
            except Exception as e: print(f"[Export Decompiler] Error processing xrefs at 0x{addr:X}: {e}")
            ok = fii.next_code()
        xrefs_list.sort(key=lambda x: x[0]); return xrefs_list

    def _optimize_cfunc(self, cfunc):
        """Apply optimization passes to decompiled function. (Experimental)"""
        print(f"[Export Decompiler] Applying 'optimize' option for 0x{cfunc.entry_ea:X} (marking dirty and attempting refresh).")
        if ida_hexrays.mark_cfunc_dirty(cfunc.entry_ea):
            try:
                cfunc.refresh_func_ctext()
                hxe_ignore = getattr(ida_hexrays, 'HX_IGNORE_ERRORS', 0)
                new_cfunc = ida_hexrays.decompile(cfunc.entry_ea, flags=hxe_ignore)
                if new_cfunc: print("[Export Decompiler] Refreshed cfunc after marking dirty."); return new_cfunc
                else: print("[Export Decompiler] Warning: Failed to re-decompile after marking dirty.")
            except Exception as e: print(f"[Export Decompiler] Error during cfunc refresh for optimization: {e}")
        else: print("[Export Decompiler] Failed to mark cfunc dirty.")
        return cfunc

    def _process_exports(self, output_path, options):
        """Core function for export enumeration and decompilation."""
        print("[Export Decompiler] Starting export processing...")
        binary_path = idaapi.get_input_file_path()
        binary_md5_bytes = idaapi.retrieve_input_file_md5()
        binary_md5_str = binary_md5_bytes.hex() if binary_md5_bytes else "N/A"
        binary_type = idaapi.get_file_type_name()
        ida_version = idaapi.get_kernel_version()

        custom_types = {}
        if options['include_types']:
            print("[Export Decompiler] Collecting custom types...")
            custom_types = self._get_custom_types()

        header_guard = self._get_header_guard(output_path) if options['use_guards'] else None

        exports = []; export_qty = ida_entry.get_entry_qty()
        print(f"[Export Decompiler] Found {export_qty} entries in the export table.")
        for i in range(export_qty):
            ordinal = ida_entry.get_entry_ordinal(i)
            ea = ida_entry.get_entry(ordinal)
            name = ida_entry.get_entry_name(i)
            if ea != idaapi.BADADDR and ida_segment.getseg(ea) is not None:
                if not name: name = f"ord_{ordinal}" if ordinal != 0 else f"unk_export_{i}"
                exports.append({'ea': ea, 'name': name, 'ordinal': ordinal})
            else:
                forwarder = ida_entry.get_entry_forwarder(i)
                if forwarder: print(f"[Export Decompiler] Skipping forwarded export index {i}: Ord {ordinal}, Name {name or '(none)'} -> {forwarder}")
                elif name or ea != idaapi.BADADDR: print(f"[Export Decompiler] Warning: Skipping invalid export index {i}: Ord {ordinal}, EA=0x{ea:X}, Name={name or '(none)'}")
        exports.sort(key=lambda x: x['ea'])
        print(f"[Export Decompiler] Processing {len(exports)} valid non-forwarded exports.")

        export_func_count = 0; successful = 0; failed = 0
        processed_eas = set()

        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(f"/*\n * Decompiled C code for exported functions\n"
                        f" * Source File: {os.path.basename(binary_path)}\n"
                        f" * Full Path:   {binary_path}\n"
                        f" * MD5:         {binary_md5_str}\n"
                        f" * Binary type: {binary_type}\n"
                        f" * IDA Version: {ida_version}\n"
                        f" * Generated by: IDA Pro Export Decompiler Plugin\n"
                        f" * Options: {options}\n */\n\n")

                if options['use_guards'] and header_guard:
                    f.write(f"#ifndef {header_guard}\n#define {header_guard}\n\n")
                f.write("/* Standard Includes */\n#include <stdint.h>\n#include <stddef.h>\n\n")

                if options['include_types']:
                    f.write("/* Custom Type Definitions (from IDA Local Types) */\n")
                    if not custom_types: f.write("/* (No custom types found or extracted) */\n")
                    else:
                        for name, type_decl in custom_types.items():
                            if type_decl and type_decl.strip() and type_decl.strip() != ';':
                                f.write(f"// Type: {name}\n{type_decl}\n\n")
                    f.write("\n")

                # --- Forward Declarations ---
                f.write("/* Forward Declarations of Exported Functions */\n")
                declared_eas = set(); temp_export_func_count = 0
                vnt_flags_decl = ida_name.VNT_VISIBLE | ida_name.VNT_TYPE
                for export in exports:
                    ea = export['ea']; name = export['name']; ordinal = export['ordinal']
                    if ea in declared_eas: continue
                    func = ida_funcs.get_func(ea);
                    if not func: continue
                    temp_export_func_count += 1
                    func_name_at_ea = ida_name.get_ea_name(ea, ida_name.GN_VISIBLE | ida_name.GN_DEMANGLED)
                    raw_name = func_name_at_ea or name
                    clean_name = ida_name.validate_name(raw_name, vnt_flags_decl) or f"func_{ea:X}"
                    f.write(f"/* Export Name: {name}, Ordinal: {ordinal}, Address: 0x{ea:X} */\n")
                    prototype_str = None; tinfo = ida_typeinf.tinfo_t()
                    if idaapi.get_tinfo(tinfo, ea) or ida_typeinf.guess_tinfo(tinfo, ea):
                        if tinfo.is_func(): prototype_str = ida_typeinf.print_tinfo(None, 0, 0, ida_typeinf.PRTYPE_1LINE | ida_typeinf.PRTYPE_SEMI, tinfo, clean_name, None)
                    f.write(f"{prototype_str}\n\n" if prototype_str else f"void* {clean_name}(void); // Fallback prototype\n\n")
                    declared_eas.add(ea)
                export_func_count = temp_export_func_count
                # --- End Forward Declarations ---

                f.write(f"\n/* Decompiled Implementations ({export_func_count} functions expected) */\n\n")

                # --- Decompilation ---
                hxe_ignore = getattr(ida_hexrays, 'HX_IGNORE_ERRORS', 0)
                vnt_flags_impl = vnt_flags_decl
                for export in exports:
                    ea = export['ea']; name = export['name']; ordinal = export['ordinal']
                    if ea in processed_eas: continue
                    func = ida_funcs.get_func(ea);
                    if not func: continue
                    func_name_at_ea = ida_name.get_ea_name(ea, ida_name.GN_VISIBLE | ida_name.GN_DEMANGLED)
                    raw_name = func_name_at_ea or name
                    clean_name = ida_name.validate_name(raw_name, vnt_flags_impl) or f"func_{ea:X}"
                    print(f"[Export Decompiler] Decompiling: {clean_name} (Export: {name}, Ord: {ordinal}) at 0x{ea:X}")
                    ida_auto.show_addr(ea)
                    f.write(f"//-------------------------------------------------------------------------\n"
                            f"// Exported Function: {name} (Ordinal: {ordinal})\n// Mapped Name:       {clean_name}\n"
                            f"// Address:           0x{ea:X}\n// Size:              {func.size()} bytes (0x{func.start_ea:X} - 0x{func.end_ea:X})\n")
                    segment = ida_segment.getseg(ea);
                    if segment: f.write(f"// Segment:           {ida_segment.get_segm_name(segment)}\n")
                    if options['include_xrefs']:
                        xrefs = self._get_xrefs(ea); f.write(f"// Cross-references from this function:\n")
                        if xrefs:
                            max_xrefs = 50
                            for i, (xaddr, taddr, tname, rtype) in enumerate(xrefs):
                                if i >= max_xrefs: f.write(f"//   ... (omitting {len(xrefs) - i} more xrefs)\n"); break
                                f.write(f"//   - 0x{xaddr:X} -> {tname} (0x{taddr:X}) [{rtype}]\n")
                        else: f.write(f"//   None found\n")
                    if options['include_asm']: f.write(f"// Assembly Listing:\n{self._get_asm_listing(func.start_ea, func.end_ea)}\n")
                    f.write(f"//-------------------------------------------------------------------------\n")
                    cfunc = None
                    try:
                        cfunc = ida_hexrays.decompile(ea, flags=hxe_ignore)
                        if cfunc:
                            if options['optimize']: cfunc = self._optimize_cfunc(cfunc)
                            f.write(str(cfunc) + "\n\n"); successful += 1
                        else:
                            err_msg = "Hex-Rays failed to decompile (returned None)"
                            print(f"[Export Decompiler] Failed: {clean_name} at 0x{ea:X}. Reason: {err_msg}")
                            f.write(f"/* Failed decomp: {clean_name} (0x{ea:X}), Export: {name}, Ord: {ordinal}. Err: {err_msg} */\n\n"); failed += 1
                    except ida_hexrays.DecompilationFailure as e:
                        print(f"[Export Decompiler] DecompilationFailure: {clean_name} at 0x{ea:X}: {e}")
                        f.write(f"/* Decompilation Failure: {clean_name} (0x{ea:X}), Export: {name}, Ord: {ordinal}. Exc: {e} */\n\n"); failed += 1
                    except Exception as e:
                        err_str = traceback.format_exc()
                        print(f"[Export Decompiler] Unexpected Error: {clean_name} at 0x{ea:X}: {e}\n{err_str}")
                        f.write(f"/* Unexpected Error: {clean_name} (0x{ea:X}), Export: {name}, Ord: {ordinal}. Exc: {e}\n{err_str} */\n\n"); failed += 1
                    finally: ida_kernwin.hide_wait_box()
                    processed_eas.add(ea)
                # --- End Decompilation ---

                f.write(f"/*\n * --- End of Decompilation ---\n"
                        f" * Total Valid Exports Found: {len(exports)}\n"
                        f" * Exported Functions Processed: {export_func_count}\n"
                        f" * Successfully Decompiled: {successful}\n"
                        f" * Failed to Decompile: {failed}\n */\n")
                if options['use_guards'] and header_guard: f.write(f"\n#endif // {header_guard}\n")

            ida_kernwin.msg(f"[Export Decompiler] Complete.\nOutput: {output_path}\n"
                            f"Exports: {len(exports)}, Functions: {export_func_count}, "
                            f"Success: {successful}, Failed: {failed}\n")

        except IOError as e:
            print(f"[Export Decompiler] File I/O Error: {e}")
            ida_kernwin.warning(f"Failed to write output file:\n{e}")
        except Exception as e:
            err_str = traceback.format_exc()
            print(f"[Export Decompiler] An unexpected error occurred during processing:\n{err_str}")
            ida_kernwin.warning(f"An unexpected error occurred:\n{e}\nSee console for details.")
        finally:
             ida_kernwin.hide_wait_box()

# Plugin registration function
def PLUGIN_ENTRY():
    return ExportDecompiler()


repository: 0xMirasio/IDALibcAutoFlags
file: autolibcflags.py
description: a simple plugin that replace decimals flags with enum on standard libc function
wanted_name: AutoLibcFlags
comment: Print comment on Flags enum/explanation near libc function call
created_at: 2023-02-11 13:17:33+00:00
pushed_at: 2024-03-11 19:14:51+00:00
url:
https://github.com/0xMirasio/IDALibcAutoFlags/blob/e1cfa359d5579b9c703c361544b6fb75c93ef683/autolibcflags.py

content:
import idaapi
import idc
import idautils
import os
import sys
import json
import platform
import ida_enum

VERSION="1.0"
p_initialized = False
#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class AutoLibcFlags(idaapi.plugin_t):
    comment = "Print comment on Flags enum/explanation near libc function call"
    help = "AutoLibcFlags Help"
    wanted_name = "AutoLibcFlags"
    flags = idaapi.PLUGIN_KEEP | idaapi.PLUGIN_MOD | idaapi.PLUGIN_PROC
    wanted_hotkey= "Ctrl+Shift+F1"


    def init(self):
        global p_initialized
        if p_initialized is False:
            p_initialized = True
            idaapi.register_action(idaapi.action_desc_t(
                "AutoLibcFlags",
                "AutoLibcFlags imports",
                None,
                None,
                None,
                0))
            print(f"AutoLibcFlags {VERSION} | Thibault Poncetta")
            self.enum = {}

        print(f"AutoLibcFlags {VERSION}")
        if self.sanityCheck():
            return -1

        self.registerFunctionSupported()
        print("Registered all functions")
        self.ImportEnum()
        print("Imported Know enums")

        if self.AddEnum():
            return -1
        print("Done caching Enum definitions to IDB")

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def sanityCheck(self):

        if platform.system() == 'Windows':
            self.enum_cache = os.path.join(os.environ['APPDATA'], 'IdaAutoLibcFlags')
            self.os = 'win'
        elif platform.system() == 'Linux':
            self.enum_cache = os.path.join(os.environ['HOME'] ,".idalibcflags")
            self.os = 'linux'
        else:
            print('[IDAAutoLibcFlags] Unknown operating system | Not supported')
            return

        if os.path.exists(self.enum_cache):
            return 0
        else:
            print("Didn't not find enum cache ! Please install plugins with ./install.py script")
            return 1

    def registerFunctionSupported(self):
        filepath = os.path.join(self.enum_cache, "functions.json")
        with open(filepath) as fd:
            self.functions_libc_supported = json.load(fd)


    def ImportEnum(self):
      
        enumdir = os.path.join(self.enum_cache, "enum")
        for filename in os.listdir(enumdir):
            filepath = os.path.join(enumdir, filename)
            if os.path.isfile(filepath):
                with open(filepath, 'r') as f:
                    self.enum[filename] = []
                    for line in f:
                        enum_value, enum_name = line.strip().split(" ")
                        self.enum[filename].append([enum_value, enum_name])

    def AddEnum(self):
        root = idaapi.get_std_dirtree(idaapi.DIRTREE_ENUMS)

        #todo : check we are IDA > 7.4 , if not we can't use get_std_dirtree()
        err = root.mkdir("IDAAutoLibcFlags")
        if err not in (idaapi.DTE_OK, idaapi.DTE_ALREADY_EXISTS):
            print(f'Could not create IDAAutoLibcFlags structures directory: "{root.errstr(err)}"')
            return 1
        
        for enum in self.enum:

            if ida_enum.get_enum(enum) != idaapi.BADNODE: #cleanup old enums
                ida_enum.del_enum(ida_enum.get_enum(enum))
                    
            idaenum = idc.add_enum(-1, enum, 0)
            if idaenum == idaapi.BADADDR:
                print("Couldn't not create enum")
                return 1

            number_added = 0
            
            err = root.rename(enum, f"IDAAutoLibcFlags/{enum}")
            if err not in (idaapi.DTE_OK, idaapi.DTE_ALREADY_EXISTS, idaapi.DTE_NOT_FOUND):
                print(f'Could not moove {enum} into IDAAutoLibcflags directory: "{root.errstr(err)}"')
                return 1

            for enum_value, enum_member_name in self.enum[enum]:
                                 
                err = idc.add_enum_member(idaenum, enum_member_name, int(enum_value, 10), -1)
                if err:
                    print(f'Could not populate {enum} on {enum_member_name} = {enum_value}  : err = {err}')
                
                number_added += 1


    def lookupFunction(self, ea, name):
        
        addr = []
        for instr in idautils.FuncItems(ea):
            r = idaapi.is_call_insn(instr)
            if r:
                for r in idautils.XrefsFrom(instr, 1):
                    f_name = idc.get_func_name(r.to).replace(".","") #remove plt prefix
                    if f_name == name:
                        addr.append(instr)
        return addr

    def find_args_with_index(self, index, call_address):
        call_args_x64 = {0 : "rdi", 1 : "rsi", 2 : "rdx", 3 : "rcx", 4 : "r8d"}

        mnem_search = call_args_x64[index]
        mnem_search_32b = mnem_search.replace("r","e")
        max_reverse_limit = 5

        cpt=1
        ea = call_address
        
        for i in range(max_reverse_limit):
            ea = idc.prev_head(ea, 0)
            mnem = idc.print_insn_mnem(ea)
            if mnem == "mov":
                op = idc.print_operand(ea,0)
                if op == mnem_search or op == mnem_search_32b:
                    return ea

        return None


    def applyEnum(self):

        for fun in self.functions_libc_supported:
            for i in range(len(self.functions_libc_supported[fun])):

                index = self.functions_libc_supported[fun][i][0]
                enum_name = "{0}".format(self.functions_libc_supported[fun][i][1])
                               
                for ea in idautils.Functions():
                    addr_patch = self.lookupFunction(ea, fun)
                    if len(addr_patch) == 0:
                        continue
                    
                    for address in addr_patch:
                        arg_addr = self.find_args_with_index(index,address)
                        if arg_addr:

                            enumida = idaapi.get_enum(enum_name)
                            value = idc.get_operand_value(arg_addr, 1)
                            if enum_name == idc.BADADDR:
                                print("Enum : {0} don't exist !".format(enum_name))
                                continue

                            r = idc.op_enum(arg_addr, 1, enumida, 0)
                            if (r == -1): #sometime IDA just fail.
                                r = idc.op_enum(arg_addr, 1, enumida, 0)



    def main(self):
        self.applyEnum()
        print("Done Applying Enum DATA. Use F5 to refresh page cache")
        
    def run(self, arg):
        self.main()


def PLUGIN_ENTRY():
    return AutoLibcFlags()


repository: Titvt/BetterIDA
file: BetterIDA.py
description: Make IDA Better!
url: https://github.com/Titvt/BetterIDA/blob/0e583593721bc7f4f7f63980c8e9e0c89ac8ea73/BetterIDA.py
wanted_name: BetterIDA
created_at: 2023-11-08 06:37:45+00:00
pushed_at: 2024-11-21 09:25:28+00:00
content:
from ida_hexrays import decompile
from ida_idaapi import PLUGIN_KEEP, plugin_t
from ida_kernwin import (
    AST_ENABLE_ALWAYS,
    BWN_PSEUDOCODE,
    LROEF_CPS_RANGE,
    SETMENU_APP,
    UI_Hooks,
    action_desc_t,
    action_handler_t,
    attach_action_to_menu,
    get_widget_type,
    line_rendering_output_entry_t,
    refresh_custom_viewer,
    register_action,
)
from ida_lines import tag_remove
from ida_segment import get_segm_class, getseg
from idautils import Functions, Segments

INDENT_RAINBOW_COLORS = [0x40FF9E40, 0x403AC267, 0x403CA2E6, 0x406C6CF5]


class IndentRainbowUIHooks(UI_Hooks):
    def __init__(self):
        UI_Hooks.__init__(self)
        self.enabled = True

    def get_lines_rendering_info(self, out, widget, info):
        if not self.enabled or get_widget_type(widget) != BWN_PSEUDOCODE:
            return

        for i in info.sections_lines:
            for j in i:
                line = tag_remove(j.line)
                length = len(line) - len(line.lstrip())
                indent = 0

                while indent * 2 < length:
                    color = INDENT_RAINBOW_COLORS[indent % len(INDENT_RAINBOW_COLORS)]
                    entry = line_rendering_output_entry_t(
                        j,
                        indent * 2,
                        min(length - indent * 2, 2),
                        LROEF_CPS_RANGE,
                        color,
                    )
                    out.entries.push_back(entry)
                    indent += 1


class IndentRainbowActionHandler(action_handler_t):
    def __init__(self):
        action_handler_t.__init__(self)
        self.ui_hooks = IndentRainbowUIHooks()
        self.ui_hooks.hook()

    def activate(self, ctx):
        self.ui_hooks.enabled = not self.ui_hooks.enabled

        if ctx.widget_type == BWN_PSEUDOCODE:
            refresh_custom_viewer(ctx.widget)

        return 0

    def update(self, ctx):
        return AST_ENABLE_ALWAYS


PSEUDO_SEARCH_DECOMPILATIONS = []


def redecompile():
    PSEUDO_SEARCH_DECOMPILATIONS.clear()
    functions = []

    for i in Segments():
        segment = getseg(i)

        if get_segm_class(segment) == "CODE":
            for j in Functions(segment.start_ea, segment.end_ea):
                functions.append(j)

    while True:
        done = True

        for i in functions[:]:
            decompilation = decompile(i)

            if decompilation is None:
                continue

            PSEUDO_SEARCH_DECOMPILATIONS.append((i, str(decompilation)))
            functions.remove(i)
            done = False

        if done:
            break


class PseudoSearchActionHandler(action_handler_t):
    def __init__(self):
        action_handler_t.__init__(self)

    def activate(self, ctx):
        redecompile()
        print(PSEUDO_SEARCH_DECOMPILATIONS)
        return 0

    def update(self, ctx):
        return AST_ENABLE_ALWAYS


def register_menu(name, label, handler, shortcut):
    register_action(action_desc_t(name, label, handler, shortcut))
    attach_action_to_menu(
        f"Edit/BetterIDA/{name}",
        name,
        SETMENU_APP,
    )


class BetterIDA(plugin_t):
    wanted_name = "BetterIDA"
    flags = PLUGIN_KEEP

    def init(self):
        register_menu(
            "IndentRainbow",
            "Indent Rainbow",
            IndentRainbowActionHandler(),
            "Shift+R",
        )
        register_menu(
            "PseudoSearch",
            "Pseudo Search",
            PseudoSearchActionHandler(),
            "Shift+F",
        )
        print("[+] BetterIDA Plugin Loaded Successfully!")
        print("[+] Version: 1.3")
        print("[+] Author: @Titvt")
        return PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        pass


def PLUGIN_ENTRY():
    return BetterIDA()


repository: junjiexing/XCopy
file: XCopy.py
description: A plug-in to enhance IDA's copy function
url: https://github.com/junjiexing/XCopy/blob/2920d26289e7bf9ddf89e3e8c824dd68e25ea342/XCopy.py
wanted_name: XCopy
comment: A plug-in to enhance IDA's copy function
created_at: 2022-12-28 13:22:24+00:00
pushed_at: 2022-12-28 13:23:12+00:00
content:
import idaapi
import ida_kernwin
from PyQt5.QtWidgets import QApplication


class XMenuContext(ida_kernwin.action_handler_t):
	@classmethod
	def get_name(self):
		return self.__name__

	@classmethod
	def register(self, label):
		return idaapi.register_action(idaapi.action_desc_t(
				self.get_name(),
				label,
				self()
			))

	@classmethod
	def unregister(self):
		idaapi.unregister_action(self.get_name())

	@classmethod
	def update(self, ctx):
		return ida_kernwin.AST_ENABLE_ALWAYS

	@classmethod
	def set_clipboard(self, txt):
		cb = QApplication.clipboard()
		cb.setText(txt, mode=cb.Clipboard)
		print("copied", txt)


class CopyVA(XMenuContext):
	def activate(self, ctx):
		cur = ida_kernwin.get_screen_ea()
		self.set_clipboard(hex(cur))


class CopyRVA(XMenuContext):
	def activate(self, ctx):
		cur = ida_kernwin.get_screen_ea()
		base = idaapi.get_imagebase()
		self.set_clipboard(hex(cur - base))


class CopyRVAX96(XMenuContext):
	def activate(self, ctx):
		cur = ida_kernwin.get_screen_ea()
		base = idaapi.get_imagebase()
		self.set_clipboard(":$" + hex(cur - base))


class UIHooks(idaapi.UI_Hooks):
	def finish_populating_widget_popup(self, form, popup):
		wt = idaapi.get_widget_type(form)
		if wt == idaapi.BWN_DISASM or wt == idaapi.BWN_PSEUDOCODE:
			try:
				idaapi.attach_action_to_popup(form, popup, CopyVA.get_name(), 'XCopy/')
				idaapi.attach_action_to_popup(form, popup, CopyRVA.get_name(), 'XCopy/')
				idaapi.attach_action_to_popup(form, popup, CopyRVAX96.get_name(), 'XCopy/')
			except:
				pass

class XCopy(idaapi.plugin_t):
	flags = idaapi.PLUGIN_KEEP
	wanted_name = "XCopy"
	comment = "A plug-in to enhance IDA's copy function"

	def init(self):
		CopyVA.register("Copy VA")
		CopyRVA.register("Copy RVA")
		CopyRVAX96.register("Copy RVA(x96dbg)")

		self.hooks = UIHooks()
		self.hooks.hook()

		# idaapi.attach_action_to_menu("Edit", CopyVA.get_name(), idaapi.SETMENU_APP)
		return idaapi.PLUGIN_KEEP

	def run(self, arg):
		print("-"*86)
		print("XCopy by Xing")
		print("https://blog.xing.re")
		# print(self.comment)
		print("-"*86)


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

def PLUGIN_ENTRY():
	return XCopy()


repository: Russinovich/IDA7.5_SDK
file: python/examples/hexrays/vds11.py
description: SDK and IDAPython for IDA Pro 7.5
wanted_name: Optimize goto chains (IDAPython)
comment: Sample plugin11 for Hex-Rays decompiler
created_at: 2021-09-01 02:18:48+00:00
pushed_at: 2021-07-23 08:19:09+00:00
url:
https://github.com/Russinovich/IDA7.5_SDK/blob/5244a2b5059113e421b4da4ecf771013a6fde2f8/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: NL2Code/CodeS
file: cleaned_repos/LazyIDA-master/LazyIDA.py
wanted_name: LazyIDA
comment: LazyIDA
created_at: 2024-03-25 04:20:11+00:00
pushed_at: 2024-07-20 08:01:55+00:00
url:
https://github.com/NL2Code/CodeS/blob/0b624ab4ef22b0d9d223f274a986eb27fe090c88/cleaned_repos/LazyIDA-master/LazyIDA.py

content:
from __future__ import division, print_function

from struct import unpack

import idaapi
import idautils
import idc
from PyQt5.Qt import QApplication

ACTION_CONVERT = ["lazyida:convert%d" % i for i in range(10)]
ACTION_SCANVUL = "lazyida:scanvul"
ACTION_COPYEA = "lazyida:copyea"
ACTION_GOTOCLIP = "lazyida:gotoclip"
ACTION_XORDATA = "lazyida:xordata"
ACTION_FILLNOP = "lazyida:fillnop"

ACTION_HX_REMOVERETTYPE = "lazyida:hx_removerettype"
ACTION_HX_COPYEA = "lazyida:hx_copyea"
ACTION_HX_COPYNAME = "lazyida:hx_copyname"
ACTION_HX_GOTOCLIP = "lazyida:hx_gotoclip"

u16 = lambda x: unpack("<H", x)[0]
u32 = lambda x: unpack("<I", x)[0]
u64 = lambda x: unpack("<Q", x)[0]

ARCH = 0
BITS = 0


def copy_to_clip(data):
    QApplication.clipboard().setText(data)


def clip_text():
    return QApplication.clipboard().text()


def parse_location(loc):
    try:
        loc = int(loc, 16)
    except ValueError:
        try:
            loc = idc.get_name_ea_simple(loc.encode().strip())
        except:
            return idaapi.BADADDR
    return loc


class VulnChoose(idaapi.Choose):
    """
    Chooser class to display result of format string vuln scan
    """

    def __init__(self, title, items, icon, embedded=False):
        idaapi.Choose.__init__(
            self,
            title,
            [["Address", 20], ["Function", 30], ["Format", 30]],
            embedded=embedded,
        )
        self.items = items
        self.icon = 45

    def GetItems(self):
        return self.items

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

    def OnClose(self):
        pass

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

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

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


class hotkey_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for hotkey actions
    """

    def __init__(self, action):
        idaapi.action_handler_t.__init__(self)
        self.action = action

    def activate(self, ctx):
        if self.action == ACTION_COPYEA:
            ea = idc.get_screen_ea()
            if ea != idaapi.BADADDR:
                copy_to_clip("0x%X" % ea)
                print("Address 0x%X has been copied to clipboard" % ea)
        elif self.action == ACTION_GOTOCLIP:
            loc = parse_location(clip_text())
            if loc != idaapi.BADADDR:
                print("Goto location 0x%x" % loc)
                idc.jumpto(loc)
        return 1

    def update(self, ctx):
        if idaapi.IDA_SDK_VERSION >= 770:
            target_attr = "widget_type"
        else:
            target_attr = "form_type"
        if ctx.__getattribute__(target_attr) in (idaapi.BWN_DISASM, idaapi.BWN_DUMP):
            return idaapi.AST_ENABLE_FOR_WIDGET
        else:
            return idaapi.AST_DISABLE_FOR_WIDGET


class menu_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for menu actions
    """

    def __init__(self, action):
        idaapi.action_handler_t.__init__(self)
        self.action = action

    def activate(self, ctx):
        if self.action in ACTION_CONVERT:
            t0, t1, view = (
                idaapi.twinpos_t(),
                idaapi.twinpos_t(),
                idaapi.get_current_viewer(),
            )
            if idaapi.read_selection(view, t0, t1):
                start, end = t0.place(view).toea(), t1.place(view).toea()
                size = end - start + 1
            elif idc.get_item_size(idc.get_screen_ea()) > 1:
                start = idc.get_screen_ea()
                size = idc.get_item_size(start)
                end = start + size
            else:
                return False

            data = idc.get_bytes(start, size)
            if isinstance(data, str):  # python2 compatibility
                data = bytearray(data)
            name = idc.get_name(start, idc.GN_VISIBLE)
            if not name:
                name = "data"
            if data:
                print("\n[+] Dump 0x%X - 0x%X (%u bytes) :" % (start, end, size))
                if self.action == ACTION_CONVERT[0]:
                    # escaped string
                    output = '"%s"' % "".join("\\x%02X" % b for b in data)
                elif self.action == ACTION_CONVERT[1]:
                    # hex string
                    output = "".join("%02X" % b for b in data)
                elif self.action == ACTION_CONVERT[2]:
                    # C array
                    output = "unsigned char %s[%d] = {" % (name, size)
                    for i in range(size):
                        if i % 16 == 0:
                            output += "\n    "
                        output += "0x%02X, " % data[i]
                    output = output[:-2] + "\n};"
                elif self.action == ACTION_CONVERT[3]:
                    # C array word
                    data += b"\x00"
                    array_size = (size + 1) // 2
                    output = "unsigned short %s[%d] = {" % (name, array_size)
                    for i in range(0, size, 2):
                        if i % 16 == 0:
                            output += "\n    "
                        output += "0x%04X, " % u16(data[i : i + 2])
                    output = output[:-2] + "\n};"
                elif self.action == ACTION_CONVERT[4]:
                    # C array dword
                    data += b"\x00" * 3
                    array_size = (size + 3) // 4
                    output = "unsigned int %s[%d] = {" % (name, array_size)
                    for i in range(0, size, 4):
                        if i % 32 == 0:
                            output += "\n    "
                        output += "0x%08X, " % u32(data[i : i + 4])
                    output = output[:-2] + "\n};"
                elif self.action == ACTION_CONVERT[5]:
                    # C array qword
                    data += b"\x00" * 7
                    array_size = (size + 7) // 8
                    output = "unsigned long %s[%d] = {" % (name, array_size)
                    for i in range(0, size, 8):
                        if i % 32 == 0:
                            output += "\n    "
                        output += "0x%016X, " % u64(data[i : i + 8])
                    output = output[:-2] + "\n};"
                elif self.action == ACTION_CONVERT[6]:
                    # python list
                    output = "[%s]" % ", ".join("0x%02X" % b for b in data)
                elif self.action == ACTION_CONVERT[7]:
                    # python list word
                    data += b"\x00"
                    output = "[%s]" % ", ".join(
                        "0x%04X" % u16(data[i : i + 2]) for i in range(0, size, 2)
                    )
                elif self.action == ACTION_CONVERT[8]:
                    # python list dword
                    data += b"\x00" * 3
                    output = "[%s]" % ", ".join(
                        "0x%08X" % u32(data[i : i + 4]) for i in range(0, size, 4)
                    )
                elif self.action == ACTION_CONVERT[9]:
                    # python list qword
                    data += b"\x00" * 7
                    output = "[%s]" % ", ".join(
                        "%#018X" % u64(data[i : i + 8]) for i in range(0, size, 8)
                    ).replace("0X", "0x")
                copy_to_clip(output)
                print(output)
        elif self.action == ACTION_XORDATA:
            t0, t1, view = (
                idaapi.twinpos_t(),
                idaapi.twinpos_t(),
                idaapi.get_current_viewer(),
            )
            if idaapi.read_selection(view, t0, t1):
                start, end = t0.place(view).toea(), t1.place(view).toea()
            else:
                if idc.get_item_size(idc.get_screen_ea()) > 1:
                    start = idc.get_screen_ea()
                    end = start + idc.get_item_size(start)
                else:
                    return False

            data = idc.get_bytes(start, end - start)
            if isinstance(data, str):  # python2 compatibility
                data = bytearray(data)
            x = idaapi.ask_long(0, "Xor with...")
            if x:
                x &= 0xFF
                print(
                    "\n[+] Xor 0x%X - 0x%X (%u bytes) with 0x%02X:"
                    % (start, end, end - start, x)
                )
                print(repr("".join(chr(b ^ x) for b in data)))
        elif self.action == ACTION_FILLNOP:
            t0, t1, view = (
                idaapi.twinpos_t(),
                idaapi.twinpos_t(),
                idaapi.get_current_viewer(),
            )
            if idaapi.read_selection(view, t0, t1):
                start, end = t0.place(view).toea(), t1.place(view).toea()
                idaapi.patch_bytes(start, b"\x90" * (end - start))
                print(
                    "\n[+] Fill 0x%X - 0x%X (%u bytes) with NOPs"
                    % (start, end, end - start)
                )
        elif self.action == ACTION_SCANVUL:
            print("\n[+] Finding Format String Vulnerability...")
            found = []
            for addr in idautils.Functions():
                name = idc.get_func_name(addr)
                if (
                    "printf" in name
                    and "v" not in name
                    and idc.get_segm_name(addr)
                    in (".text", ".plt", ".idata", ".plt.got")
                ):
                    xrefs = idautils.CodeRefsTo(addr, False)
                    for xref in xrefs:
                        vul = self.check_fmt_function(name, xref)
                        if vul:
                            found.append(vul)
            if found:
                print("[!] Done! %d possible vulnerabilities found." % len(found))
                ch = VulnChoose("Vulnerability", found, None, False)
                ch.Show()
            else:
                print("[-] No format string vulnerabilities found.")
        else:
            return 0

        return 1

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

    @staticmethod
    def check_fmt_function(name, addr):
        """
        Check if the format string argument is not valid
        """
        function_head = idc.get_func_attr(addr, idc.FUNCATTR_START)

        while True:
            addr = idc.prev_head(addr)
            op = idc.print_insn_mnem(addr).lower()
            dst = idc.print_operand(addr, 0)

            if op in ("ret", "retn", "jmp", "b") or addr < function_head:
                return

            c = idc.get_cmt(addr, 0)
            if c and c.lower() == "format":
                break
            elif name.endswith(("snprintf_chk",)):
                if op in ("mov", "lea") and dst.endswith(("r8", "r8d", "[esp+10h]")):
                    break
            elif name.endswith(("sprintf_chk",)):
                if op in ("mov", "lea") and (
                    dst.endswith(("rcx", "[esp+0Ch]", "R3"))
                    or dst.endswith("ecx")
                    and BITS == 64
                ):
                    break
            elif name.endswith(("snprintf", "fnprintf")):
                if op in ("mov", "lea") and (
                    dst.endswith(("rdx", "[esp+8]", "R2"))
                    or dst.endswith("edx")
                    and BITS == 64
                ):
                    break
            elif name.endswith(("sprintf", "fprintf", "dprintf", "printf_chk")):
                if op in ("mov", "lea") and (
                    dst.endswith(("rsi", "[esp+4]", "R1"))
                    or dst.endswith("esi")
                    and BITS == 64
                ):
                    break
            elif name.endswith("printf"):
                if op in ("mov", "lea") and (
                    dst.endswith(("rdi", "[esp]", "R0"))
                    or dst.endswith("edi")
                    and BITS == 64
                ):
                    break

        # format arg found, check its type and value
        # get last oprend
        op_index = idc.generate_disasm_line(addr, 0).count(",")
        op_type = idc.get_operand_type(addr, op_index)
        opnd = idc.print_operand(addr, op_index)

        if op_type == idc.o_reg:
            # format is in register, try to track back and get the source
            _addr = addr
            while True:
                _addr = idc.prev_head(_addr)
                _op = idc.print_insn_mnem(_addr).lower()
                if _op in ("ret", "retn", "jmp", "b") or _addr < function_head:
                    break
                elif (
                    _op in ("mov", "lea", "ldr") and idc.print_operand(_addr, 0) == opnd
                ):
                    op_type = idc.get_operand_type(_addr, 1)
                    opnd = idc.print_operand(_addr, 1)
                    addr = _addr
                    break

        if op_type == idc.o_imm or op_type == idc.o_mem:
            # format is a memory address, check if it's in writable segment
            op_addr = idc.get_operand_value(addr, op_index)
            seg = idaapi.getseg(op_addr)
            if seg:
                if not seg.perm & idaapi.SEGPERM_WRITE:
                    # format is in read-only segment
                    return

        print("0x%X: Possible Vulnerability: %s, format = %s" % (addr, name, opnd))
        return ["0x%X" % addr, name, opnd]


class hexrays_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for hexrays actions
    """

    def __init__(self, action):
        idaapi.action_handler_t.__init__(self)
        self.action = action
        self.ret_type = {}

    def activate(self, ctx):
        if self.action == ACTION_HX_REMOVERETTYPE:
            vdui = idaapi.get_widget_vdui(ctx.widget)
            self.remove_rettype(vdui)
            vdui.refresh_ctext()
        elif self.action == ACTION_HX_COPYEA:
            ea = idaapi.get_screen_ea()
            if ea != idaapi.BADADDR:
                copy_to_clip("0x%X" % ea)
                print("Address 0x%X has been copied to clipboard" % ea)
        elif self.action == ACTION_HX_COPYNAME:
            highlight = idaapi.get_highlight(idaapi.get_current_viewer())
            name = highlight[0] if highlight else None
            if name:
                copy_to_clip(name)
                print("%s has been copied to clipboard" % name)
        elif self.action == ACTION_HX_GOTOCLIP:
            loc = parse_location(clip_text())
            print("Goto location 0x%x" % loc)
            idc.jumpto(loc)
        else:
            return 0

        return 1

    def update(self, ctx):
        vdui = idaapi.get_widget_vdui(ctx.widget)
        return idaapi.AST_ENABLE_FOR_WIDGET if vdui else idaapi.AST_DISABLE_FOR_WIDGET

    def remove_rettype(self, vu):
        if vu.item.citype == idaapi.VDI_FUNC:
            # current function
            ea = vu.cfunc.entry_ea
            old_func_type = idaapi.tinfo_t()
            if not vu.cfunc.get_func_type(old_func_type):
                return False
        elif (
            vu.item.citype == idaapi.VDI_EXPR
            and vu.item.e.is_expr()
            and vu.item.e.type.is_funcptr()
        ):
            # call xxx
            ea = vu.item.get_ea()
            old_func_type = idaapi.tinfo_t()

            func = idaapi.get_func(ea)
            if func:
                try:
                    cfunc = idaapi.decompile(func)
                except idaapi.DecompilationFailure:
                    return False

                if not cfunc.get_func_type(old_func_type):
                    return False
            else:
                return False
        else:
            return False

        fi = idaapi.func_type_data_t()
        if ea != idaapi.BADADDR and old_func_type.get_func_details(fi):
            # Return type is already void
            if fi.rettype.is_decl_void():
                # Restore ret type
                if ea not in self.ret_type:
                    return True
                ret = self.ret_type[ea]
            else:
                # Save ret type and change it to void
                self.ret_type[ea] = fi.rettype
                ret = idaapi.BT_VOID

            # Create new function info with new rettype
            fi.rettype = idaapi.tinfo_t(ret)

            # Create new function type with function info
            new_func_type = idaapi.tinfo_t()
            new_func_type.create_func(fi)

            # Apply new function type
            if idaapi.apply_tinfo(ea, new_func_type, idaapi.TINFO_DEFINITE):
                return vu.refresh_view(True)

        return False


class UI_Hook(idaapi.UI_Hooks):
    def __init__(self):
        idaapi.UI_Hooks.__init__(self)

    def finish_populating_widget_popup(self, form, popup):
        form_type = idaapi.get_widget_type(form)

        if form_type == idaapi.BWN_DISASM or form_type == idaapi.BWN_DUMP:
            t0, t1, view = (
                idaapi.twinpos_t(),
                idaapi.twinpos_t(),
                idaapi.get_current_viewer(),
            )
            if (
                idaapi.read_selection(view, t0, t1)
                or idc.get_item_size(idc.get_screen_ea()) > 1
            ):
                idaapi.attach_action_to_popup(form, popup, ACTION_XORDATA, None)
                idaapi.attach_action_to_popup(form, popup, ACTION_FILLNOP, None)
                for action in ACTION_CONVERT:
                    idaapi.attach_action_to_popup(form, popup, action, "Convert/")

        if form_type == idaapi.BWN_DISASM and (ARCH, BITS) in [
            (idaapi.PLFM_386, 32),
            (idaapi.PLFM_386, 64),
            (idaapi.PLFM_ARM, 32),
        ]:
            idaapi.attach_action_to_popup(form, popup, ACTION_SCANVUL, None)


class HexRays_Hook(object):
    def callback(self, event, *args):
        if event == idaapi.hxe_populating_popup:
            form, phandle, vu = args
            if vu.item.citype == idaapi.VDI_FUNC or (
                vu.item.citype == idaapi.VDI_EXPR
                and vu.item.e.is_expr()
                and vu.item.e.type.is_funcptr()
            ):
                idaapi.attach_action_to_popup(
                    form, phandle, ACTION_HX_REMOVERETTYPE, None
                )
        elif event == idaapi.hxe_double_click:
            vu, shift_state = args
            # auto jump to target if clicked item is xxx->func();
            if vu.item.citype == idaapi.VDI_EXPR and vu.item.e.is_expr():
                expr = idaapi.tag_remove(vu.item.e.print1(None))
                if "->" in expr:
                    # find target function
                    name = expr.split("->")[-1]
                    addr = idc.get_name_ea_simple(name)
                    if addr == idaapi.BADADDR:
                        # try class::function
                        e = vu.item.e
                        while e.x:
                            e = e.x
                        addr = idc.get_name_ea_simple(
                            "%s::%s" % (str(e.type).split()[0], name)
                        )

                    if addr != idaapi.BADADDR:
                        idc.jumpto(addr)
                        return 1
        return 0


class LazyIDA_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = "LazyIDA"
    help = ""
    wanted_name = "LazyIDA"
    wanted_hotkey = ""

    def init(self):
        self.hexrays_inited = False
        self.registered_actions = []
        self.registered_hx_actions = []

        global ARCH
        global BITS
        ARCH = idaapi.ph_get_id()
        info = idaapi.get_inf_structure()
        if info.is_64bit():
            BITS = 64
        elif info.is_32bit():
            BITS = 32
        else:
            BITS = 16

        print("LazyIDA (v1.0.0.3) plugin has been loaded.")

        # Register menu actions
        menu_actions = (
            idaapi.action_desc_t(
                ACTION_CONVERT[0],
                "Convert to string",
                menu_action_handler_t(ACTION_CONVERT[0]),
                None,
                None,
                80,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[1],
                "Convert to hex string",
                menu_action_handler_t(ACTION_CONVERT[1]),
                None,
                None,
                8,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[2],
                "Convert to C/C++ array (BYTE)",
                menu_action_handler_t(ACTION_CONVERT[2]),
                None,
                None,
                38,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[3],
                "Convert to C/C++ array (WORD)",
                menu_action_handler_t(ACTION_CONVERT[3]),
                None,
                None,
                38,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[4],
                "Convert to C/C++ array (DWORD)",
                menu_action_handler_t(ACTION_CONVERT[4]),
                None,
                None,
                38,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[5],
                "Convert to C/C++ array (QWORD)",
                menu_action_handler_t(ACTION_CONVERT[5]),
                None,
                None,
                38,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[6],
                "Convert to python list (BYTE)",
                menu_action_handler_t(ACTION_CONVERT[6]),
                None,
                None,
                201,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[7],
                "Convert to python list (WORD)",
                menu_action_handler_t(ACTION_CONVERT[7]),
                None,
                None,
                201,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[8],
                "Convert to python list (DWORD)",
                menu_action_handler_t(ACTION_CONVERT[8]),
                None,
                None,
                201,
            ),
            idaapi.action_desc_t(
                ACTION_CONVERT[9],
                "Convert to python list (QWORD)",
                menu_action_handler_t(ACTION_CONVERT[9]),
                None,
                None,
                201,
            ),
            idaapi.action_desc_t(
                ACTION_XORDATA,
                "Get xored data",
                menu_action_handler_t(ACTION_XORDATA),
                None,
                None,
                9,
            ),
            idaapi.action_desc_t(
                ACTION_FILLNOP,
                "Fill with NOPs",
                menu_action_handler_t(ACTION_FILLNOP),
                None,
                None,
                9,
            ),
            idaapi.action_desc_t(
                ACTION_SCANVUL,
                "Scan format string vulnerabilities",
                menu_action_handler_t(ACTION_SCANVUL),
                None,
                None,
                160,
            ),
        )
        for action in menu_actions:
            idaapi.register_action(action)
            self.registered_actions.append(action.name)

        # Register hotkey actions
        hotkey_actions = (
            idaapi.action_desc_t(
                ACTION_COPYEA,
                "Copy EA",
                hotkey_action_handler_t(ACTION_COPYEA),
                "w",
                "Copy current EA",
                0,
            ),
            idaapi.action_desc_t(
                ACTION_GOTOCLIP,
                "Goto clip EA",
                hotkey_action_handler_t(ACTION_GOTOCLIP),
                "Shift-G",
                "Goto clipboard EA",
                0,
            ),
        )
        for action in hotkey_actions:
            idaapi.register_action(action)
            self.registered_actions.append(action.name)

        # Add ui hook
        self.ui_hook = UI_Hook()
        self.ui_hook.hook()

        # Add hexrays ui callback
        if idaapi.init_hexrays_plugin():
            addon = idaapi.addon_info_t()
            addon.id = "tw.l4ys.lazyida"
            addon.name = "LazyIDA"
            addon.producer = "Lays"
            addon.url = "https://github.com/L4ys/LazyIDA"
            addon.version = "1.0.0.3"
            idaapi.register_addon(addon)

            hx_actions = (
                idaapi.action_desc_t(
                    ACTION_HX_REMOVERETTYPE,
                    "Remove return type",
                    hexrays_action_handler_t(ACTION_HX_REMOVERETTYPE),
                    "v",
                ),
                idaapi.action_desc_t(
                    ACTION_HX_COPYEA,
                    "Copy ea",
                    hexrays_action_handler_t(ACTION_HX_COPYEA),
                    "w",
                ),
                idaapi.action_desc_t(
                    ACTION_HX_COPYNAME,
                    "Copy name",
                    hexrays_action_handler_t(ACTION_HX_COPYNAME),
                    "c",
                ),
                idaapi.action_desc_t(
                    ACTION_HX_GOTOCLIP,
                    "Goto clipboard ea",
                    hexrays_action_handler_t(ACTION_HX_GOTOCLIP),
                    "Shift-G",
                ),
            )
            for action in hx_actions:
                idaapi.register_action(action)
                self.registered_hx_actions.append(action.name)

            self.hx_hook = HexRays_Hook()
            idaapi.install_hexrays_callback(self.hx_hook.callback)
            self.hexrays_inited = True

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        if hasattr(self, "ui_hook"):
            self.ui_hook.unhook()

        # Unregister actions
        for action in self.registered_actions:
            idaapi.unregister_action(action)

        if self.hexrays_inited:
            # Unregister hexrays actions
            for action in self.registered_hx_actions:
                idaapi.unregister_action(action)
            if self.hx_hook:
                idaapi.remove_hexrays_callback(self.hx_hook.callback)
            idaapi.term_hexrays_plugin()


def PLUGIN_ENTRY():
    return LazyIDA_t()


