# plugins

repository: PyEGo/exp-gist
file: hard-gists-me/b3143dde185cecda7c1dee7ffbce5d2c/snippet.py
description: Experiment results on hard-gists dataset
comment: Decode IOCTL codes!
created_at: 2021-04-18 11:16:06+00:00
pushed_at: 2022-11-18 17:22:11+00:00
url:
https://github.com/PyEGo/exp-gist/blob/bb8212e435ab525bb8ebdee2f620f850133a9b99/hard-gists-me/b3143dde185cecda7c1dee7ffbce5d2c/snippet.py


repository: PyEGo/exp-gist
file: hard-gists-ego/b3143dde185cecda7c1dee7ffbce5d2c/snippet.py
description: Experiment results on hard-gists dataset
comment: Decode IOCTL codes!
created_at: 2021-04-18 11:16:06+00:00
pushed_at: 2022-11-18 17:22:11+00:00
url:
https://github.com/PyEGo/exp-gist/blob/bb8212e435ab525bb8ebdee2f620f850133a9b99/hard-gists-ego/b3143dde185cecda7c1dee7ffbce5d2c/snippet.py


repository: r0ysue/FridaAndrioidNativeBeginnersBook
file: Chap14/AES/样本/mt/idaplugins/trace_natives-main/trace_natives-main/traceNatives.py
description: 《安卓Frida SO逆向分析实战》
wanted_name: traceNatives
comment: traceNatives
created_at: 2023-05-31 07:37:25+00:00
pushed_at: 2023-10-22 13:17:28+00:00
url:
https://github.com/r0ysue/FridaAndrioidNativeBeginnersBook/blob/6be19b07269fb91f5e054e5485ce327a1bf07a60/Chap14/AES/%E6%A0%B7%E6%9C%AC/mt/idaplugins/trace_natives-main/trace_natives-main/traceNatives.py

content:
# -*- coding:utf-8 -*-
import os
from idaapi import plugin_t
from idaapi import PLUGIN_PROC
from idaapi import PLUGIN_OK
import ida_nalt
import idaapi
import idautils
import idc
import time


# 获取SO文件名和路径
def getSoPathAndName():
    fullpath = ida_nalt.get_input_file_path()
    filepath,filename = os.path.split(fullpath)
    return filepath,filename

# 获取代码段的范围
def getSegAddr():
    textStart = []
    textEnd = []

    for seg in idautils.Segments():
        if (idc.get_segm_name(seg)).lower() == '.text' or (
        idc.get_segm_name(seg)).lower() == 'text':
            tempStart = idc.get_segm_start(seg)
            tempEnd = idc.get_segm_end(seg)

            textStart.append(tempStart)
            textEnd.append(tempEnd)

    return min(textStart), max(textEnd)


class traceNatives(plugin_t):
    flags = PLUGIN_PROC
    comment = "traceNatives"
    help = ""
    wanted_name = "traceNatives"
    wanted_hotkey = ""

    def init(self):
        print("traceNatives(v0.1) plugin has been loaded.")
        return PLUGIN_OK

    def run(self, arg):
        # 查找需要的函数
        ea, ed = getSegAddr()
        search_result = []
        for func in idautils.Functions(ea, ed):
            try:
                functionName = str(idaapi.ida_funcs.get_func_name(func))
                if len(list(idautils.FuncItems(func))) > 10:
                    # 如果是thumb模式，地址+1
                    arm_or_thumb = idc.get_sreg(func, "T")
                    if arm_or_thumb:
                        func += 1
                    search_result.append(hex(func))
            except:
                pass

        so_path, so_name = getSoPathAndName()
        search_result = [f"-a '{so_name}!{offset}'" for offset in search_result]
        search_result = " ".join(search_result)

        script_name = so_name.split(".")[0] + "_" + str(int(time.time())) +".txt"
        save_path = os.path.join(so_path, script_name)
        with open(save_path, "w", encoding="utf-8")as F:
            F.write(search_result)

        print("使用方法如下：")
        print(f"frida-trace -UF -O {save_path}")

    def term(self):
        pass


def PLUGIN_ENTRY():
    return traceNatives()


repository: REMnux/distro
file: ppasrc/ghidra/ghidra-10.1-bionic/ghidra/Extensions/IDAPro/Python/7xx/plugins/xml_importer.py
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2014-09-12 19:41:54+00:00
pushed_at: 2025-02-26 18:45:43+00:00
description:
This repository contains supplemental files for the REMnux distro and the source files for the Debian packages that the distro installs from the REMnux package repository on Launchpad.

url:
https://github.com/REMnux/distro/blob/9789bbc22b57362bc26e7025059bf72ac2f7bcd9/ppasrc/ghidra/ghidra-10.1-bionic/ghidra/Extensions/IDAPro/Python/7xx/plugins/xml_importer.py

content:
## ###
#  IP: GHIDRA
# 
#  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.
##
#---------------------------------------------------------------------
# xmlimp.py - IDA XML Importer plugin
#---------------------------------------------------------------------
"""
Plugin for IDA to import a XML PROGRAM file into an existing open database.
This file must be placed in the IDA plugins directory.
The file idaxml.py must be placed in the IDA python directory.
"""

import ida_idaapi
import ida_pro
import idaxml
import idc
import sys

class XmlImporterPlugin(ida_idaapi.plugin_t):
    """
    XML Importer plugin class
    """
    flags = 0
    comment = "Import XML PROGRAM file"
    help = "Import XML <PROGRAM> document to database"
    wanted_name = "XML Importer"
    wanted_hotkey = "Ctrl-Alt-l"

    def init(self):
        """
        init function for XML Importer plugin.
        
        Returns:
            Constant PLUGIN_OK if this IDA version supports the plugin,
            else returns PLUGIN_SKIP if this IDA is older than the supported
            baseline version.
        """
        if idaxml.is_ida_version_supported():
            return ida_idaapi.PLUGIN_OK 
        else:
            return ida_idaapi.PLUGIN_SKIP


    def run(self, arg):
        """
        run function for XML Importer plugin.
        
        Args:
            arg: Integer, a non-zero value enables auto-run feature for
                 IDA batch (no gui) processing mode. Default is 0.
        """
        st = idc.set_ida_state(idc.IDA_STATUS_WORK)
        xml = idaxml.XmlImporter(idaxml.PLUGIN, arg)
        try:
            try:
                xml.import_xml()
            except idaxml.Cancelled:
                msg = "XML Import cancelled!"
                print "\n" + msg
                idc.warning(msg)
            except idaxml.MultipleAddressSpacesNotSupported:
                msg  = "XML Import cancelled!"
                msg += "\n\nXML Import does not currently support"
                msg += "\nimporting multiple address spaces."
                print "\n" + msg
                idc.warning(msg)
            except:
                msg = "***** Exception occurred: XML Importer failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idc.warning(msg)
        finally:
            xml.cleanup()
            idc.set_ida_state(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


repository: ImpaleAskewArdently/ghidra_hex_bin
file: GhidraBuild/IDAPro/Python/7xx/plugins/xml_importer.py
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2024-08-01 03:15:40+00:00
pushed_at: 2024-08-01 04:30:29+00:00
url:
https://github.com/ImpaleAskewArdently/ghidra_hex_bin/blob/06ca3190d1d9519df8513489ce4f539f7be2af03/GhidraBuild/IDAPro/Python/7xx/plugins/xml_importer.py

content:
## ###
#  IP: GHIDRA
# 
#  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.
##
#---------------------------------------------------------------------
# xmlimp.py - IDA XML Importer plugin
#---------------------------------------------------------------------
"""
Plugin for IDA to import a XML PROGRAM file into an existing open database.
This file must be placed in the IDA plugins directory.
The file idaxml.py must be placed in the IDA python directory.
"""

import ida_idaapi
import ida_pro
import idaxml
import idc
import sys

class XmlImporterPlugin(ida_idaapi.plugin_t):
    """
    XML Importer plugin class
    """
    flags = 0
    comment = "Import XML PROGRAM file"
    help = "Import XML <PROGRAM> document to database"
    wanted_name = "XML Importer"
    wanted_hotkey = "Ctrl-Alt-l"

    def init(self):
        """
        init function for XML Importer plugin.
        
        Returns:
            Constant PLUGIN_OK if this IDA version supports the plugin,
            else returns PLUGIN_SKIP if this IDA is older than the supported
            baseline version.
        """
        if idaxml.is_ida_version_supported():
            return ida_idaapi.PLUGIN_OK 
        else:
            return ida_idaapi.PLUGIN_SKIP


    def run(self, arg):
        """
        run function for XML Importer plugin.
        
        Args:
            arg: Integer, a non-zero value enables auto-run feature for
                 IDA batch (no gui) processing mode. Default is 0.
        """
        st = idc.set_ida_state(idc.IDA_STATUS_WORK)
        xml = idaxml.XmlImporter(idaxml.PLUGIN, arg)
        try:
            try:
                xml.import_xml()
            except idaxml.Cancelled:
                msg = "XML Import cancelled!"
                print "\n" + msg
                idc.warning(msg)
            except idaxml.MultipleAddressSpacesNotSupported:
                msg  = "XML Import cancelled!"
                msg += "\n\nXML Import does not currently support"
                msg += "\nimporting multiple address spaces."
                print "\n" + msg
                idc.warning(msg)
            except:
                msg = "***** Exception occurred: XML Importer failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idc.warning(msg)
        finally:
            xml.cleanup()
            idc.set_ida_state(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


repository: OSUSecLab/PaymentScope
file: Tools/ghidra_10.0.3_PUBLIC/Extensions/IDAPro/Python/7xx/plugins/xml_importer.py
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2022-08-12 01:57:01+00:00
pushed_at: 2022-08-12 02:19:05+00:00
url:
https://github.com/OSUSecLab/PaymentScope/blob/ef8f29c79834c0a6e5d08e8420bd9d04b1a78cbe/Tools/ghidra_10.0.3_PUBLIC/Extensions/IDAPro/Python/7xx/plugins/xml_importer.py


repository: m223744/MalwareAnalysisCapstone
file: mainCode/ghidraApi/ghidra/Extensions/IDAPro/Python/7xx/plugins/xml_importer.py
description: Automated Static Malware Analysis Capstone
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2020-09-28 13:36:25+00:00
pushed_at: 2022-04-01 17:19:05+00:00
url:
https://github.com/m223744/MalwareAnalysisCapstone/blob/78f293a154b2ae8b3c7dc0fe01626b440f2ee69f/mainCode/ghidraApi/ghidra/Extensions/IDAPro/Python/7xx/plugins/xml_importer.py


repository: williballenthin/idawilli
file: plugins/navband_visited.py
description: IDA Pro resources, scripts, and configurations
wanted_name: Navband visited
comment: Tracks and records all visited addresses, showing in the navband
created_at: 2016-12-14 02:23:46+00:00
pushed_at: 2025-05-12 10:45:42+00:00
url:
https://github.com/williballenthin/idawilli/blob/063417207eebd061fbd58f6441e6b5e9ffc43f8e/plugins/navband_visited.py

content:
import idaapi
import ida_bytes
import ida_kernwin


class navband_visited_plugmod_t(idaapi.plugmod_t, ida_kernwin.UI_Hooks):
    def __init__(self):
        super().__init__()
        self.addresses = set()
        self.orig_colorizer = None

    def screen_ea_changed(self, ea, prev_ea):
        if ea == idaapi.BADADDR:
            return

        for i in range(ea, ea + ida_bytes.get_item_size(ea)):
            self.addresses.add(i)

    def colorizer(self, ea, nbytes):
        # when we've visited the address, color it black
        if ea in self.addresses:
            return 0x000000  # black

        # otherwise, show the original colors
        return ida_kernwin.call_nav_colorizer(self.orig_colorizer, ea, nbytes)

    def run(self, arg):
        self.hook()
        self.orig_colorizer = ida_kernwin.set_nav_colorizer(self.colorizer)

    def term(self):
        self.unhook()
        idaapi.set_nav_colorizer(self.orig_colorizer)


class navband_visited_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_MULTI
    comment = "Tracks and records all visited addresses, showing in the navband"
    help = "This plugin records all the addresses you visit during your analysis, showing the visited addresses in the navband"
    wanted_name = "Navband visited"
    wanted_hotkey = ""

    def init(self):
        return navband_visited_plugmod_t()


def PLUGIN_ENTRY():
    return navband_visited_plugin_t()


repository: cedricp/ddt4all
file: ddtplugins/card_programming.py
description: OBD tool
created_at: 2016-11-03 16:21:57+00:00
pushed_at: 2025-05-11 13:32:37+00:00
url:
https://github.com/cedricp/ddt4all/blob/21ed3920fdeab2fe1a874cb13929afa71a62a3ac/ddtplugins/card_programming.py

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

# Plugin to program keys on Megane II
# (c) 2017
# This is an example plugin


import PyQt5.QtCore as core
import PyQt5.QtWidgets as gui

import ecu
import options

_ = options.translator('ddt4all')

plugin_name = _("Megane/Scenic II card programming")
category = _("Keys")
need_hw = True
ecufile = "UCH_84_J84_03_60"

a8_mess = [19, 29, 39, 4, 24, 46, 31, 12,
           16, 28, 7, 10, 15, 0, 40, 26,
           18, 20, 3, 8, 34, 8, 47, 21,
           2, 25, 38, 6, 4, 22, 14, 33,
           17, 1, 44, 11, 32, 41, 36, 42,
           5, 45, 5, 43, 30, 37, 13, 27]

a8_xor = [0x10, 0x20, 0x20, 0x10, 0x81, 0x88]
a8_xor_2 = [0x10, 0x20, 0x28, 0x10, 0x81, 0x88]


def a8(isk_c):
    '''MC9S12DG256 (Megane 2, Scenic 2)'''
    apv = ''
    isk = isk_c.replace(' ', '').decode('hex')
    if len(isk) != 6:
        return None
    bit_i = ''
    for c in isk:
        bit_i += bin(ord(c))[2:].zfill(8)

    # if not (bit_i[ 4]!=bit_i[23] and bit_i[12]==bit_i[43] and
    #        bit_i[ 4]==bit_i[35] and bit_i[ 8]==bit_i[ 9] and bit_i[ 4]==bit_i[ 5]):
    #  return 'Invalid ISK'

    base = ''
    for b in a8_xor:
        base = base + bin(b)[2:].zfill(8)

    i = 0
    for b in base:
        z = bit_i[a8_mess[i]]
        if b != z:
            z = '1'
        else:
            z = '0'
        base = base[:i] + z + base[i + 1:]
        i += 1

    for i in range(0, 6):
        apv += hex(int(base[i * 8:i * 8 + 8], 2))[2:].zfill(2).upper() + ' '

    return apv


def a8_2(isk_c):
    '''MC9S12DG256 (algo2)(Megane 2, Scenic 2)'''
    apv = ''
    isk = isk_c.replace(' ', '').decode('hex')
    if len(isk) != 6:
        return None
    bit_i = ''
    for c in isk:
        bit_i += bin(ord(c))[2:].zfill(8)

    base = ''
    for b in a8_xor_2:
        base = base + bin(b)[2:].zfill(8)

    i = 0
    for b in base:
        z = bit_i[a8_mess[i]]
        if b != z:
            z = '1'
        else:
            z = '0'
        base = base[:i] + z + base[i + 1:]
        i += 1

    for i in range(0, 6):
        apv += hex(int(base[i * 8:i * 8 + 8], 2))[2:].zfill(2).upper() + ' '

    return apv


def get_isk(ecu_response):
    ecu_response = ecu_response.replace(' ', '').strip()
    return ecu_response[19 * 2:25 * 2]


class CardProg(gui.QDialog):
    def __init__(self):
        super(CardProg, self).__init__()
        options.debug = True
        self.apvok = False
        self.megane_ecu = ecu.Ecu_file(ecufile, True)
        self.start_session_request = self.megane_ecu.requests[u'Start Diagnostic Session']
        self.after_sale_request = self.megane_ecu.requests[u'ACCEDER AU MODE APRES-VENTE']
        self.learn_key_request = self.megane_ecu.requests[u'APPRENDRE BADGE']
        self.key_bits_status = self.megane_ecu.requests[u'Status général des opérations badges Bits']
        self.key_bytes_status = self.megane_ecu.requests[u'Status général des opérations badges Octets']
        self.reserved_frame_request = self.megane_ecu.requests[u'Trame AB: Trame réservée']
        self.learn_validate_request = self.megane_ecu.requests[u'SORTIE DU MODE APV : VALIDATION']
        self.learn_cancel_request = self.megane_ecu.requests[u'SORTIE DU MODE APV : ABANDON']

        self.ecu_connect()
        self.start_diag_session()

        self.apv_access_button = gui.QPushButton(_("ENTER AFTER SALE MODE"))
        self.apv_access_button.clicked.connect(self.set_apv_from_input)
        self.pininput = gui.QLineEdit()
        self.pininput.setText("0" * 12)
        self.pininput.setInputMask("H" * 12)
        self.iskoutput = gui.QLineEdit()
        self.iskoutput.setInputMask("H" * 12)
        self.numkeys_label = gui.QLabel("0")
        self.currentcardide = gui.QLineEdit()
        self.currentcardide.setReadOnly(True)
        self.apv_status = gui.QLabel("")
        self.apv_status.setAlignment(core.Qt.AlignHCenter)
        label = gui.QLabel(_("<font color='red'>MEGANE II CARD PROGRAMMING<br>EXPERIMENTAL : NOT TESTED YET</font>"))
        label.setAlignment(core.Qt.AlignHCenter)
        self.learnbutton = gui.QPushButton(_("LEARN"))
        self.validatebutton = gui.QPushButton(_("VALIDATE"))
        self.cancelbutton = gui.QPushButton(_("CANCEL"))

        self.algocheck = gui.QCheckBox(_("Algo 2"))
        self.algocheck.setChecked(False)

        self.learnbutton.clicked.connect(self.learn_action)
        self.validatebutton.clicked.connect(self.validate_action)
        self.cancelbutton.clicked.connect(self.cancel_action)
        self.iskoutput.textChanged.connect(self.calculate_pin)
        self.algocheck.toggled.connect(self.calculate_pin)

        layout = gui.QGridLayout()
        layout.addWidget(label, 0, 0, 1, 0)
        layout.addWidget(gui.QLabel(_("ISK")), 1, 0)
        layout.addWidget(self.iskoutput, 1, 1)
        layout.addWidget(gui.QLabel(_("PIN")), 2, 0)
        layout.addWidget(self.pininput, 2, 1)
        layout.addWidget(self.algocheck, 3, 1)
        layout.addWidget(self.apv_access_button, 4, 0, 1, 0)
        layout.addWidget(self.apv_status, 5, 0, 1, 0)
        layout.addWidget(gui.QLabel(_("Num key learnt")), 6, 0)
        layout.addWidget(self.numkeys_label, 6, 1)
        layout.addWidget(gui.QLabel(_("CARD IDE")), 7, 0)
        layout.addWidget(self.currentcardide, 7, 1)
        layout.addWidget(self.learnbutton, 8, 0, 1, 0)
        layout.addWidget(self.validatebutton, 9, 0)
        layout.addWidget(self.cancelbutton, 9, 1)

        self.setLayout(layout)

        # I have to use a timer here to avoid ECU from escaping session mode
        # A shot every 1.5 seconds is enough, normally
        self.tester_timer = core.QTimer()
        self.tester_timer.setSingleShot(False)
        self.tester_timer.setInterval(1500)
        self.tester_timer.timeout.connect(self.check_all)

        isk = self.get_isk()
        self.iskoutput.setText(isk)

        self.tester_timer.start()

    def calculate_pin(self):
        ISK = ''.join(str(ord(c)) for c in str(self.iskoutput.text()))
        if len(ISK) == 12:
            if self.algocheck.checkState():
                PIN = a8_2(ISK)
            else:
                PIN = a8(ISK)
            if PIN is not None:
                self.pininput.setText(PIN)
                self.iskoutput.setStyleSheet("color: green;")
                return

        self.pininput.setText("000000000000")
        self.iskoutput.setStyleSheet("color: red;")

    def learn_action(self):
        self.learn_key_request.send_request()

    def validate_action(self):
        self.learn_validate_request.send_request()

    def cancel_action(self):
        self.learn_cancel_request.send_request()

    def get_isk(self):
        # No dataitem for ISK :(
        # Need to decode by myself
        if options.simulation_mode:
            reply = "61 AB 02 FC 0D 08 51 4C 86 55 54 00 00 00 00 00 00 00 00 8D E8 EE 16 79 D3 C9 A7 A7 CC " \
                    "F6 AC 00 00 00 00 00 00 2A"
        else:
            reply = options.elm.request(self.reserved_frame_request.sentbytes)

        if reply.startswith('7F') or reply.startswith('WRONG'):
            options.main_window.logview.append(_("Cannot get ISK, check connections and UCH compatibility"))
            return

        isk = reply.replace(' ', '').strip()
        return get_isk(isk)

    def ecu_connect(self):
        connection = self.megane_ecu.connect_to_hardware()
        if not connection:
            options.main_window.logview.append(_("Cannot connect to ECU"))
            self.finished()

    def check_all(self):
        self.check_apv_status()
        self.check_num_key_learnt()

    def set_apv_from_input(self):
        apv = ''.join(str(ord(c)) for c in str(self.pininput.text()))
        if len(apv) != 12:
            return
        self.after_sale_request.send_request({u'Code APV': apv})

    def enable_buttons(self, enable):
        self.learnbutton.setEnabled(enable)
        self.validatebutton.setEnabled(enable)
        self.cancelbutton.setEnabled(enable)

    def check_apv_status(self):
        key_bits_status_values = self.key_bits_status.send_request({}, "61 06 21 21 20 00 00 00 08 00 00 80 03")

        value_apv_ok = key_bits_status_values[u'VSC Code APV_Reconnu']
        value_apv_reaff = key_bits_status_values[u'VSC ModeAPV_ReaffArmé']
        value_apv_uch_reaff = key_bits_status_values[u'VSC ModeAPV_AppUCH_Armé']

        if value_apv_ok == '0':
            self.apv_status.setText(_("<font color='red'>PIN CODE NOT RECOGNIZED</font>"))
            self.apvok = False
            self.enable_buttons(False)
            return

        if value_apv_reaff == '1':
            self.apv_status.setText(_("<font color='green'>PIN CODE OK / KEY LEARNING</font>"))
            self.enable_buttons(True)
            self.apvok = True
            return

        elif value_apv_uch_reaff == '1':
            self.apv_status.setText(_("<font color='green'>PIN CODE OK / UCH LEARNING</font>"))
            self.enable_buttons(True)
            self.apvok = True
            return

        self.apvok = False
        self.enable_buttons(False)

        self.apv_status.setText("<font color='red'>UCH NOT READY</font>")

    def check_num_key_learnt(self):
        key_bytes_status_value = self.key_bytes_status.send_request({}, "61 08 00 00 00 00 00 00 0A 00 00 00 00 00 00 "
                                                                        "00 00 00 00 00 00 00 00 00 00 02 00 00 01 2C "
                                                                        "00 00 01 2C 00 00 2C 00 00 01 2C")

        num_key_learnt = key_bytes_status_value[u'VSC NbTotalDeBadgeAppris']
        current_key_ide = key_bytes_status_value[u'VSC Code_IDE']

        self.numkeys_label.setText(num_key_learnt)
        self.currentcardide.setText(current_key_ide)

    def start_diag_session(self):
        sds_request = self.megane_ecu.requests[u"Start Diagnostic Session"]
        sds_stream = " ".join(sds_request.build_data_stream({}))
        if options.simulation_mode:
            print("SdS stream", sds_stream)
            return
        options.elm.start_session_can(sds_stream)


def plugin_entry():
    cp = CardProg()
    cp.exec_()


repository: patois/dsync
file: dsync.py
description: IDAPython plugin that synchronizes disassembler and decompiler views
url: https://github.com/patois/dsync/blob/af8467e6eeea60b0a2e4fcc133e8d5e198d4ab9f/dsync.py
wanted_name: dsync
created_at: 2019-05-09 11:33:29+00:00
pushed_at: 2021-01-20 10:12:24+00:00
content:
from idaapi import *

"""
This plugin for IDA synchronizes Hexrays decompiler views and disassembly views (from
decompiled code to disassembly by default - use the TAB key for synchronizing from
disassembly to decompiled code).

It also highlights all addresses (code and data definitions) involved. Pressing the
hotkey Ctrl-Shift-S switches synchronization on and off. Hovering over pseudocode
items will display corresponding disassembled code in a hint window. The item that
belongs to the item that is located under the cursor will be highlighted and pointed
to by an arrow.

The plugin requires IDA 7.3.
"""

__author__ = 'https://github.com/patois'

HL_COLOR = 0xAD8044

# -----------------------------------------------------------------------
class idb_hook_t(IDB_Hooks):
    def __init__(self, hxehook):
        self.hxehook = hxehook
        IDB_Hooks.__init__(self)

    def savebase(self):
        self.hxehook._reset_all_colors()
        return 0

# -----------------------------------------------------------------------
class hxe_hook_t(Hexrays_Hooks):
    def __init__(self):
        Hexrays_Hooks.__init__(self)
        self.idbhook = idb_hook_t(self)
        self.idbhook.hook()
        self.pseudocode_instances = {}
        self.n_spaces = 40

    def close_pseudocode(self, vd):
        self._reset_colors(vd.view_idx, ignore_vd=True)
        refresh_idaview_anyway()
        return 0

    def create_hint(self, vd):
        result = self._get_vd_context(vd)
        if result:
            _, _, _, item_ea_list = result

            if len(item_ea_list):
                if vd.get_current_item(USE_MOUSE):
                    cur_item_ea = vd.item.it.ea
                else:
                    cur_item_ea = BADADDR

                lines = []
                for ea in item_ea_list:
                    disasm_line = generate_disasm_line(ea, 0)
                    if disasm_line:
                        addr = "0x%x: " % ea
                        
                        if cur_item_ea == ea:
                            prefix = COLSTR("==> %s" % addr, SCOLOR_INSN)
                        else:
                            prefix = "    " + addr

                        lines.append(prefix+disasm_line)

                lines.append("")
                lines.append(self.n_spaces * "-")
                lines.append("")
                custom_hints = "\n".join(lines)
                # ask decompiler to append default hints
                return (2, custom_hints, len(lines))
        return 0

    def curpos(self, vd):
        self._reset_all_colors()
        self._apply_colors(vd)
        refresh_idaview_anyway()
        return 0

    def refresh_pseudocode(self, vd):
        self._reset_all_colors(ignore_vd=True)
        return 0

    def cleanup(self):
        self._reset_all_colors()
        refresh_idaview_anyway()

        if self.idbhook:
            self.idbhook.unhook()
            self.idbhook = None
        return
 
    def _reset_colors(self, idx, ignore_vd=False):
        v = self.pseudocode_instances[idx]
        if v:
            pseudocode, lineno, color, disasm_lines = v
            if not ignore_vd and pseudocode:
                try:
                    pseudocode[lineno].bgcolor = color
                except: # wtf
                    pass
            for ea, color in disasm_lines:
                set_item_color(ea, color)
        self.pseudocode_instances.pop(idx)
        return

    def _reset_all_colors(self, ignore_vd=False):
        # restore colors
        if self.pseudocode_instances:
            pi = list(self.pseudocode_instances)
            for k in pi:
                self._reset_colors(k, ignore_vd)
            self.pseudocode_instances = {}
        return

    def _apply_colors(self, vd):
        result = self._get_vd_context(vd)
        if result:
            pseudocode, lineno, col, item_ea_list = result
            disasm_lines = [(ea, get_item_color(ea)) for ea in item_ea_list]
            if len(item_ea_list):
                jumpto(item_ea_list[0], -1, UIJMP_IDAVIEW | UIJMP_DONTPUSH)
            self.pseudocode_instances[vd.view_idx] = (pseudocode, lineno, col, disasm_lines)

            if pseudocode:
                try:
                    pseudocode[lineno].bgcolor = HL_COLOR
                except: # wtf
                    pass
            for ea, _ in disasm_lines:
                set_item_color(ea, HL_COLOR)
        return

    def _get_item_indexes(self, line):
        indexes = []
        tag = COLOR_ON + chr(COLOR_ADDR)
        pos = line.find(tag)
        while pos != -1 and len(line[pos+len(tag):]) >= COLOR_ADDR_SIZE:
            addr = line[pos+len(tag):pos+len(tag)+COLOR_ADDR_SIZE]
            idx = int(addr, 16)
            a = ctree_anchor_t()
            a.value = idx
            if a.is_valid_anchor() and a.is_citem_anchor():
                """
                print "a.value %s %d lvar %s citem %s itp %s blkcmt %s" % (
                    a.is_valid_anchor(),
                    a.get_index(),
                    a.is_lvar_anchor(),
                    a.is_citem_anchor(),
                    a.is_itp_anchor(),
                    a.is_blkcmt_anchor())
                """
                indexes.append(a.get_index())
            pos = line.find(tag, pos+len(tag)+COLOR_ADDR_SIZE)
        return indexes

    def _get_vd_context(self, vd):
        if vd:
            lineno = vd.cpos.lnnum
            pseudocode = vd.cfunc.get_pseudocode()

            if pseudocode and lineno != -1:
                try:
                    color = pseudocode[lineno].bgcolor
                    line = pseudocode[lineno].line
                   
                    item_idxs = self._get_item_indexes(line)
                    ea_list = {}
                    for i in item_idxs:
                        try:
                            item = vd.cfunc.treeitems.at(i)
                            if item and item.ea != BADADDR:
                                ea_list[item.ea] = None
                        except:
                            pass
                    return (pseudocode, lineno, color, sorted(ea_list.keys()))
                except:
                    pass
        return None

# -----------------------------------------------------------------------
def is_ida_version(min_ver_required):
    return IDA_SDK_VERSION >= min_ver_required

# -----------------------------------------------------------------------
class Dsync(ida_idaapi.plugin_t):
    comment = ''
    help = ''
    flags = PLUGIN_MOD
    wanted_name = 'dsync'
    wanted_hotkey = 'Ctrl-Shift-S'
    hxehook = None

    def init(self):
        required_ver = 730
        if not is_ida_version(required_ver) or not init_hexrays_plugin():
            msg ("[!] '%s' is inactive (IDA v%d and decompiler required).\n" % (Dsync.wanted_name, required_ver))
            return PLUGIN_SKIP

        msg("[+] '%s' loaded. %s activates/deactivates synchronization.\n" % (Dsync.wanted_name, Dsync.wanted_hotkey))
        return PLUGIN_KEEP

    def run(self, arg):
        if not Dsync.hxehook:
            Dsync.hxehook = hxe_hook_t()
            Dsync.hxehook.hook()
        else:
            Dsync.hxehook.unhook()
            Dsync.hxehook.cleanup()
            Dsync.hxehook = None

        msg("[+] %s is %sabled now.\n" % (Dsync.wanted_name, "en" if Dsync.hxehook else "dis"))

    def term(self):
        msg("[+] %s unloaded.\n" % (Dsync.wanted_name))
        if Dsync.hxehook:
            Dsync.hxehook.unhook()
            Dsync.hxehook.cleanup()
            Dsync.hxehook = None

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


repository: WPeace-HcH/WPeChatGPT
file: WPeChatGPT.py
created_at: 2023-02-28 06:30:37+00:00
pushed_at: 2025-03-28 02:58:29+00:00
description:
A plugin for IDA that can help to analyze binary file, it can be based on commonly used AI big models such as OpenAI and DeepSeek.

url:
https://github.com/WPeace-HcH/WPeChatGPT/blob/16d07972f678b7c35fa1a974b482009aea4361fe/WPeChatGPT.py

content:
import functools
import idaapi
import ida_hexrays
import ida_kernwin
import idc
import openai
import re
import threading
import json
import httpx
import sys, os
# Windows
path = os.path.dirname(os.path.abspath(__file__)) + "\\Auto-WPeGPT_WPeace\\"
# MacOS
#path = os.path.dirname(os.path.abspath(__file__)) + "/Auto-WPeGPT_WPeace/"
sys.path.append(path)
import Auto_WPeGPT
# Whether to use Chinese explanation code.
ZH_CN = True

# Use ChatGPT
PLUGIN_NAME = 'WPeChat-GPT'
# Use DeepSeek
#PLUGIN_NAME = 'WPeChat-DeepSeek'

# Set your API key here, or put in in the model_api_key environment variable.
model_api_key = "ENTER_API_KEY_HERE"

# Set your forward-proxy if necessary. (e.g. Clash = http://127.0.0.1:7890)
proxy = ""
# Set reverse-proxy-URL or custom-api-URL if you need. (e.g. Azure OpenAI)
proxy_address = ""


# Plugin information, you can change the model here.
if PLUGIN_NAME == "WPeChat-GPT":
    PROD_NAME = 'ChatGPT'
    MODEL = 'gpt-4'
    print("WPeChatGPT is using ChatGPT.")
elif PLUGIN_NAME == "WPeChat-DeepSeek":
    PROD_NAME = 'DeepSeek'
    MODEL = 'deepseek-chat'
    print("WPeChatGPT is using DeepSeek.")
# Create openai client (python openai package version > 1.2)
if PROD_NAME == "DeepSeek":
    client = openai.OpenAI(base_url="https://api.deepseek.com", api_key=model_api_key)
elif proxy:
    client = openai.OpenAI(http_client=httpx.Client(proxies=proxy, transport=httpx.HTTPTransport(local_address="0.0.0.0")), api_key=model_api_key)
    print("WPeChatGPT has appointed the forward-proxy.")
elif proxy_address:
    client = openai.OpenAI(base_url=proxy_address, api_key=model_api_key)
    print("WPeChatGPT has appointed the reverse-proxy.")
else:
    client = openai.OpenAI(api_key=model_api_key)


# WPeChatGPT 分析解释函数
class ExplainHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        funcComment = getFuncComment(idaapi.get_screen_ea())
        if "---WPeChat_START---" in funcComment:
            if ZH_CN:
                print("当前函数已经完成过 %s:Explain 分析，请查看注释或删除注释重新分析。@WPeace"%(PLUGIN_NAME))
            else:
                print("The current function has been analyzed by %s:Explain, please check the comment or delete the comment to re-analyze. @WPeace"%(PLUGIN_NAME))
            return 0
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        # 中文
        if ZH_CN:
            query_model_async("下面是一个C语言伪代码函数，分别分析该函数的预期目的、参数的作用、详细功能，最后取一个新的函数名字。（用简体中文回答我，并且回答开始前加上'---WPeChat_START---'字符串结束后加上'---WPeChat_END---'字符串）\n"
                + str(decompiler_output),
                functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v, cmtFlag=0, printFlag=0),
                0)
        # English
        else:
            query_model_async("Analyze the following C language pseudo-code function, respectively speculate on the use environment, expected purpose, and detailed function of the function, and finally choose a new name for this function. (add '---WPeChat_START---' string before the beginning of the answer and add '---WPeChat_END---' string after the end)\n" + str(decompiler_output), functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v, cmtFlag=0, printFlag=0), 0)
        return 1

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


# WPeChatGPT 重命名变量函数
class RenameHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async("Analyze the following C function:\n" + str(decompiler_output) +
                            "\nSuggest better variable names, reply with a JSON array where keys are the original names"
                            "and values are the proposed names. Do not explain anything, only print the JSON "
                            "dictionary.",
                          functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=v),
                          0)
        return 1

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


# WPeChatGPT 使用python3对函数进行还原
class PythonHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        # lastAddr 为函数的最后一行汇编代码地址
        lastAddr = idc.prev_head(idc.get_func_attr(idaapi.get_screen_ea(), idc.FUNCATTR_END))
        # 获取对应注释
        addrComment = getAddrComment(lastAddr)
        if "---WPeChat_Python_START---" in str(addrComment):
            if ZH_CN:
                print("当前函数已经完成过 %s:Python 分析，请查看注释或删除注释重新分析。@WPeace"%(PLUGIN_NAME))
            else:
                print("The current function has been analyzed by %s:Python, please check the comment or delete the comment to re-analyze. @WPeace"%(PLUGIN_NAME))
            return 0
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        # 中文
        if ZH_CN:
            query_model_async("分析下面的C语言伪代码并用python3代码进行还原。（回答开始前加上'---WPeChat_Python_START---'字符串结束后加上'---WPeChat_Python_END---'字符串）\n"
                + str(decompiler_output),
                functools.partial(comment_callback, address=lastAddr, view=v, cmtFlag=1, printFlag=1),
                0)
        # English
        else:
            query_model_async("Analyze the following C language pseudocode and restore it with python3 code. (Add '---WPeChat_Python_START---' string before the beginning of the answer and add '---WPeChat_Python_END---' string after the end)\n"
                              + str(decompiler_output),
                              functools.partial(comment_callback, address=lastAddr, view=v, cmtFlag=1, printFlag=1),
                              0)
        return 1

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


# WPeChatGPT 尝试寻找函数漏洞
class FindVulnHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        funcComment = getFuncComment(idaapi.get_screen_ea())
        if "---WPeChat_VulnFinder_START---" in funcComment:
            if ZH_CN:
                print("当前函数已经完成过 %s:VulnFinder 分析，请查看注释或删除注释重新分析。@WPeace"%(PLUGIN_NAME))
            else:
                print("The current function has been analyzed by %s:VulnFinder, please check the comment or delete the comment to re-analyze. @WPeace"%(PLUGIN_NAME))
            return 0
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        # 中文
        if ZH_CN:
            query_model_async("查找下面这个C语言伪代码函数的漏洞并提出可能的利用方法。（用简体中文回答我，并且回答开始前加上'---WPeChat_VulnFinder_START---'字符串结束后加上'---WPeChat_VulnFinder_END---'字符串）\n"
                + str(decompiler_output),
                functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v, cmtFlag=0, printFlag=2),
                0)
        # English
        else:
            query_model_async("Find the following C function vulnerabilty and suggest a possible way to exploit it.(Use English to answer me, and answer before plus '---WPeChat_VulnFinder_START---' the end of the string plus '---WPeChat_VulnFinder_END---' string)\n"
                              + str(decompiler_output),
                              functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v, cmtFlag=0, printFlag=2),
                              0)
        return 1

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


# WPeChatGPT 尝试对漏洞函数生成EXP
class expCreateHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        funcComment = getFuncComment(idaapi.get_screen_ea())
        if "---WPeChat_VulnPython_START---" in funcComment:
            if ZH_CN:
                print("当前函数已经完成过 %s:ExpCreater 分析，请查看注释或删除注释重新分析。@WPeace"%(PLUGIN_NAME))
            else:
                print("The current function has been analyzed by %s:ExpCreater, please check the comment or delete the comment to re-analyze. @WPeace"%(PLUGIN_NAME))
            return 0
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        # 中文
        if ZH_CN:
            query_model_async("使用Python构造代码来利用下面函数中的漏洞。（用简体中文回答我，并且回答开始前加上'---WPeChat_VulnPython_START---'字符串结束后加上'---WPeChat_VulnPython_END---'字符串）\n"
                + str(decompiler_output),
                functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v, cmtFlag=0, printFlag=3),
                0)
        # English
        else:
            query_model_async("Use Python to construct code to exploit the vulnerabilities in the following functions.(Answer before plus '---WPeChat_VulnPython_START---' the end of the string plus '---WPeChat_VulnPython_END---' string)\n"
                              + str(decompiler_output),
                              functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v, cmtFlag=0, printFlag=3),
                              0)
        return 1

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


def autoChatFunc(funcTree:str, strings:str, callback):
    messages = []
    input_funcTree = funcTree
    messages.append({"role": "user", "content": input_funcTree})
    input_strings = strings
    messages.append({"role": "user", "content": input_strings})
    if ZH_CN:
        messages.append({"role": "user", "content": "结合该程序的函数调用结构及其所包含的字符串，猜测其运行目的及功能。"})
        messages.append({"role": "user", "content": "请再仔细分析后告诉我该程序的运行目的及大概功能。"})
    else:
        messages.append({"role": "user", "content": "Combining the function call structure of the program and the strings it contains, guess its purpose and function."})
        messages.append({"role": "user", "content": "Please tell me the purpose and general function of the program after careful analysis."})
    t = threading.Thread(target=chat_api_worker, args=(messages, MODEL, callback))
    t.start()


def chat_api_worker(messages, model, callback):
    try:
        response = client.chat.completions.create(messages=messages, model=model)
    except Exception as e:
        if "maximum context length" in str(e):
            print("此二进制文件的分析数据超过了 GPT-3.5-API 的最大长度！请期待后续版本 :)@WPeace")
            return 0
        elif "Cannot connect to proxy" in str(e):
            print("代理出现问题，请稍后重试或检查代理。@WPeace")
            return 0
        else:
            print(f"General exception encountered while running the query: {str(e)}")
            return 0
    callback(response)


def handle_response(autoGptfolder, response):
    message = response.choices[0].message
    if ZH_CN:
        print("GPT 分析完毕，已将结果输出到文件夹：" + autoGptfolder + " 当中！")
    else:
        print("The GPT analysis is complete and the result has been output to the folder: " + autoGptfolder)
    fp = open(autoGptfolder + "GPT-Result.txt", "w")
    fp.write(message.content)
    fp.close()
    print("Auto-WPeGPT finished! :)@WPeace\n")


# Auto-WPeGPT 自动化分析
class autoHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        Auto_WPeGPT.main()
        idb_path = idc.get_idb_path()
        idb_name = 'WPe_' + os.path.basename(idb_path)
        autoGptfolder = os.path.join(os.getcwd(), idb_name) + '\\'
        functreeFilepath = autoGptfolder + "funcTree.txt"
        mainFunctreeFilepath = autoGptfolder + "mainFuncTree.txt"
        stringsFilepath = autoGptfolder + "effectiveStrings.txt"
        file = open(functreeFilepath, "r")
        functreeData = file.read()
        file.close()
        file = open(mainFunctreeFilepath, "r")
        mainFunctreeData = file.read()
        file.close()
        file = open(stringsFilepath, "r")
        stringsData = file.read()
        file.close()
        funcNumber = idaapi.get_func_qty()
        print("There are %d functions in total in this binary file." %funcNumber)
        if funcNumber < 150:
            callback_autogpt = functools.partial(handle_response, autoGptfolder)
            autoChatFunc(functreeData, stringsData, callback_autogpt)
        else:
            callback_autogpt = functools.partial(handle_response, autoGptfolder)
            autoChatFunc(mainFunctreeData, stringsData, callback_autogpt)
        print("Auto-WPeGPT v0.2 start to analysis...")
        return 1

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


# Gepetto query_model Method
def query_model(query, cb, max_tokens=2500):
    """
    向 gpt-3.5-turbo 发送查询的函数。
    :param query: The request to send to gpt-3.5-turbo
    :param cb: Tu function to which the response will be passed to.
    """
    try:
        response = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "user", "content": query}
            ]
        )
        ida_kernwin.execute_sync(functools.partial(cb, response=response.choices[0].message.content), ida_kernwin.MFF_WRITE)
    except openai.BadRequestError as e:
        # Context length exceeded. Determine the max number of tokens we can ask for and retry.
        m = re.search(r'maximum context length is (\d+) tokens, however you requested \d+ tokens \((\d+) in your '
                      r'prompt;', str(e))
        if not m:
            print(f"{MODEL} could not complete the request: {str(e)}")
            return
        (hard_limit, prompt_tokens) = (int(m.group(1)), int(m.group(2)))
        max_tokens = hard_limit - prompt_tokens
        if max_tokens >= 750:
            print(f"{PLUGIN_NAME}-Warning: Context length too long! Try reducing tokens to {max_tokens}...")
            print("Request to %s sent retried..."%(MODEL))
            query_model(query, cb, max_tokens)
        else:
            print(f"Unfortunately, this function is too large to be analyzed using {PROD_NAME}-{MODEL} API. @WPeace")
    except openai.OpenAIError as e:
        if "That model is currently overloaded with other requests" in str(e) or "Request timed out" in str(e):
            print(f"{PROD_NAME}-{MODEL} API 繁忙，请稍后重试或检查代理。@WPeace")
        elif "Cannot connect to proxy" in str(e):
            print("代理出现问题，请稍后重试或检查代理。@WPeace")
        else:
            print(f"OpenAI server could not complete the request: {str(e)}")
    except Exception as e:
        print(f"General exception encountered while running the query: {str(e)}")


# Gepetto query_model_async Method
def query_model_async(query, cb, time):
    """
    创建线程调用 query_model 函数。
    :param query: The request to send to gpt-3.5-turbo
    :param cb: Tu function to which the response will be passed to.
    :param time: whether it is a retry.
    """
    if time == 0:
        if ZH_CN:
            print(f"正在发送 {PROD_NAME}:{MODEL} API 请求，完成后将输出提示。@WPeace")
        else:
            print(f"Sending {PROD_NAME}-{MODEL} API request, will output a prompt when completed. @WPeace")
        print("Request to %s sent..."%(MODEL))
    else:
        if ZH_CN:
            print(f"正在重新发送 {PROD_NAME}-{MODEL} API 请求。@WPeace")
        else:
            print(f"Resending {PROD_NAME}-{MODEL} API request. @WPeace")
    t = threading.Thread(target=query_model, args=[query, cb])
    t.start()


# Gepetto comment_callback Method
def comment_callback(address, view, response, cmtFlag, printFlag):
    """
    在对应地址处设置注释的回调函数。
    :param address: The address of the function to comment
    :param view: A handle to the decompiler window
    :param response: The comment to add
    """
    # Add the response as a comment in IDA.
    # 通过参数控制不同形式添加注释
    if cmtFlag == 0:
        idc.set_func_cmt(address, response, 0)
    elif cmtFlag == 1:
        idc.set_cmt(address, response, 1)
    # Refresh the window so the comment is displayed properly
    if view:
        view.refresh_view(False)
    print("%s query finished!"%(MODEL))
    if printFlag == 0:
        if ZH_CN:
            print("%s:Explain 完成分析，已对函数 %s 进行注释。@WPeace" %(PLUGIN_NAME, idc.get_func_name(address)))
        else:
            print("%s:Explain finished analyzing, function %s has been commented. @WPeace" %(PLUGIN_NAME, idc.get_func_name(address)))
    elif printFlag == 1:
        if ZH_CN:
            print("%s:Python 完成分析，已在函数末尾地址 %s 汇编处进行注释。@WPeace"%(PLUGIN_NAME, hex(address)))
        else:
            print("%s:Python finished parsing, commented at assembly at address %s at end of function. @WPeace" %(PLUGIN_NAME, hex(address)))
    elif printFlag == 2:
        if ZH_CN:
            print("%s:VulnFinder 完成分析，已对函数 %s 进行注释。@WPeace" %(PLUGIN_NAME, idc.get_func_name(address)))
        else:
            print("%s: VulnFinder finished analyzing, function %s has been annotated. @WPeace" %(PLUGIN_NAME, idc.get_func_name(address)))
    elif printFlag == 3:
        if ZH_CN:
            print("%s:ExpCreater 完成分析，已对函数 %s 进行注释。@WPeace" %(PLUGIN_NAME, idc.get_func_name(address)))
        else:
            print("%s:ExpCreater finished analyzing, commented on function %s. @WPeace" %(PLUGIN_NAME, idc.get_func_name(address)))


# Gepetto rename_callback Method
def rename_callback(address, view, response, retries=0):
    """
    重命名函数变量的回调函数。
    :param address: The address of the function to work on
    :param view: A handle to the decompiler window
    :param response: The response from gpt-3.5-turbo
    :param retries: The number of times that we received invalid JSON
    """
    j = re.search(r"\{[^}]*?\}", response)
    if not j:
        if retries >= 3:  # Give up obtaining the JSON after 3 times.
            print(f"{PROD_NAME}-{MODEL} API has no valid response, please try again later. @WPeace")
            return
        print(f"Cannot extract valid JSON from the response. Asking the model to fix it...")
        query_model_async("The JSON document provided in this response is invalid. Can you fix it?\n" + response,
                          functools.partial(rename_callback,
                                            address=address,
                                            view=view,
                                            retries=retries + 1), 
                                            1)
        return
    try:
        names = json.loads(j.group(0))
    except json.decoder.JSONDecodeError:
        if retries >= 3:  # Give up fixing the JSON after 3 times.
            print(f"{PROD_NAME}-{MODEL} API has no valid response, please try again later. @WPeace")
            return
        print(f"The JSON document returned is invalid. Asking the model to fix it...")
        query_model_async("Please fix the following JSON document:\n" + j.group(0),
                          functools.partial(rename_callback,
                                            address=address,
                                            view=view,
                                            retries=retries + 1), 
                                            1)
        return
    # The rename function needs the start address of the function
    function_addr = idaapi.get_func(address).start_ea
    replaced = []
    for n in names:
        if ida_hexrays.rename_lvar(function_addr, n, names[n]):
            replaced.append(n)

    # Update possible names left in the function comment
    comment = idc.get_func_cmt(address, 0)
    if comment and len(replaced) > 0:
        for n in replaced:
            comment = re.sub(r'\b%s\b' % n, names[n], comment)
        idc.set_func_cmt(address, comment, 0)
    # Refresh the window to show the new names
    if view:
        view.refresh_view(True)
    print("%s query finished!"%(MODEL))
    if ZH_CN:
        print(f"{PLUGIN_NAME}:RenameVariable 完成分析，已重命名{len(replaced)}个变量。@WPeace")
    else:
        print(f"{PLUGIN_NAME}:RenameVariable Completed analysis, renamed {len(replaced)} variables. @WPeace")


# 获取函数注释
def getFuncComment(address):
    cmt = idc.get_func_cmt(address, 0)
    if not cmt:
        cmt = idc.get_func_cmt(address, 1)
    return cmt


# 获取地址注释
def getAddrComment(address):
    cmt = idc.get_cmt(address, 0)
    if not cmt:
        cmt = idc.get_cmt(address, 1)
    return cmt


# Add context menu actions
class ContextMenuHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        idaapi.attach_action_to_popup(form, popup, myplugin_WPeChatGPT.explain_action_name, "%s/"%(PLUGIN_NAME))
        idaapi.attach_action_to_popup(form, popup, myplugin_WPeChatGPT.rename_action_name, "%s/"%(PLUGIN_NAME))
        idaapi.attach_action_to_popup(form, popup, myplugin_WPeChatGPT.python_action_name, "%s/"%(PLUGIN_NAME))
        idaapi.attach_action_to_popup(form, popup, myplugin_WPeChatGPT.vulnFinder_action_name, "%s/"%(PLUGIN_NAME))
        idaapi.attach_action_to_popup(form, popup, myplugin_WPeChatGPT.expPython_action_name, "%s/"%(PLUGIN_NAME))


class myplugin_WPeChatGPT(idaapi.plugin_t):
    autoWPeGPT_action_name = "%s:Auto-WPeGPT"%(PLUGIN_NAME)
    autoWPeGPT_menu_path = "Edit/%s/Auto-WPeGPT/Auto-WPeGPT v0.2"%(PLUGIN_NAME)
    explain_action_name = "%s:Explain_Function"%(PLUGIN_NAME)
    explain_menu_path = "Edit/%s/函数分析"%(PLUGIN_NAME)
    rename_action_name = "%s:Rename_Function"%(PLUGIN_NAME)
    rename_menu_path = "Edit/%s/重命名函数变量"%(PLUGIN_NAME)
    python_action_name = "%s:Python_Function"%(PLUGIN_NAME)
    python_menu_path = "Edit/%s/Python还原此函数"%(PLUGIN_NAME)
    vulnFinder_action_name = "%s:VulnFinder_Function"%(PLUGIN_NAME)
    vulnFinder_menu_path = "Edit/%s/二进制漏洞查找"%(PLUGIN_NAME)
    expPython_action_name = "%s:VulnPython_Function"%(PLUGIN_NAME)
    expPython_menu_path = "Edit/%s/尝试生成Exploit"%(PLUGIN_NAME)
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ''
    comment = "%s Plugin for IDA"%(PLUGIN_NAME)
    help = "Find more information at https://github.com/wpeace-hch"
    menu = None
    flags = 0
    def init(self):
        # Check whether the decompiler is available
        if not ida_hexrays.init_hexrays_plugin():
            return idaapi.PLUGIN_SKIP
        if ZH_CN:
            # create Auto-WPeGPT action
            autoWPeGPT_action = idaapi.action_desc_t(self.autoWPeGPT_action_name,
                                                  '二进制文件自动化分析 v0.2',
                                                  autoHandler(),
                                                  "",
                                                  "使用 %s 对二进制文件进行自动化分析"%(MODEL),
                                                  199)
            idaapi.register_action(autoWPeGPT_action)
            idaapi.attach_action_to_menu(self.autoWPeGPT_menu_path, self.autoWPeGPT_action_name, idaapi.SETMENU_APP)
            # Function explaining action
            explain_action = idaapi.action_desc_t(self.explain_action_name,
                                                  '函数分析',
                                                  ExplainHandler(),
                                                  "Ctrl+Alt+G",
                                                  "使用 %s 分析当前函数"%(MODEL),
                                                  199)
            idaapi.register_action(explain_action)
            idaapi.attach_action_to_menu(self.explain_menu_path, self.explain_action_name, idaapi.SETMENU_APP)
            # Variable renaming action
            rename_action = idaapi.action_desc_t(self.rename_action_name,
                                                 '重命名函数变量',
                                                 RenameHandler(),
                                                 "Ctrl+Alt+R",
                                                 "使用 %s 重命名当前函数的变量"%(MODEL),
                                                 199)
            idaapi.register_action(rename_action)
            idaapi.attach_action_to_menu(self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)
            # python function action
            python_action = idaapi.action_desc_t(self.python_action_name,
                                                 'Python还原此函数',
                                                 PythonHandler(),
                                                 "",
                                                 "使用 %s 分析当前函数并用python3还原"%(MODEL),
                                                 199)
            idaapi.register_action(python_action)
            idaapi.attach_action_to_menu(self.python_menu_path, self.python_action_name, idaapi.SETMENU_APP)
            # find vulnerabilty action
            vulnFinder_action = idaapi.action_desc_t(self.vulnFinder_action_name,
                                                  '二进制漏洞查找',
                                                  FindVulnHandler(),
                                                  "Ctrl+Alt+E",
                                                  "使用 %s 在当前函数中查找漏洞"%(MODEL),
                                                  199)
            idaapi.register_action(vulnFinder_action)
            idaapi.attach_action_to_menu(self.vulnFinder_menu_path, self.vulnFinder_action_name, idaapi.SETMENU_APP)
            # create exploit action
            expPython_action = idaapi.action_desc_t(self.expPython_action_name,
                                                  '尝试生成Exploit',
                                                  expCreateHandler(),
                                                  "",
                                                  "使用 %s 尝试对漏洞函数生成EXP"%(MODEL),
                                                  199)
            idaapi.register_action(expPython_action)
            idaapi.attach_action_to_menu(self.expPython_menu_path, self.expPython_action_name, idaapi.SETMENU_APP)
            # Register context menu actions
            self.menu = ContextMenuHooks()
            self.menu.hook()
            print("Auto-WPeGPT v0.2 is ready.")
            print("%s v2.6 works fine! :)@WPeace\n"%(PLUGIN_NAME))
        else:
            # create Auto-WPeGPT action
            autoWPeGPT_action = idaapi.action_desc_t(self.autoWPeGPT_action_name,
                                                  'Automated analysis v0.2',
                                                  autoHandler(),
                                                  "",
                                                  "使用 %s 对二进制文件进行自动化分析"%(MODEL),
                                                  199)
            idaapi.register_action(autoWPeGPT_action)
            idaapi.attach_action_to_menu(self.autoWPeGPT_menu_path, self.autoWPeGPT_action_name, idaapi.SETMENU_APP)
            # Function explaining action
            explain_action = idaapi.action_desc_t(self.explain_action_name,
                                                  'Function analysis',
                                                  ExplainHandler(),
                                                  "Ctrl+Alt+G",
                                                  "使用 %s 分析当前函数"%(MODEL),
                                                  199)
            idaapi.register_action(explain_action)
            idaapi.attach_action_to_menu(self.explain_menu_path, self.explain_action_name, idaapi.SETMENU_APP)
            # Variable renaming action
            rename_action = idaapi.action_desc_t(self.rename_action_name,
                                                 'Rename function variables',
                                                 RenameHandler(),
                                                 "Ctrl+Alt+R",
                                                 "使用 %s 重命名当前函数的变量"%(MODEL),
                                                 199)
            idaapi.register_action(rename_action)
            idaapi.attach_action_to_menu(self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)
            # python function action
            python_action = idaapi.action_desc_t(self.python_action_name,
                                                 'Python restores this function',
                                                 PythonHandler(),
                                                 "",
                                                 "使用 %s 分析当前函数并用python3还原"%(MODEL),
                                                 199)
            idaapi.register_action(python_action)
            idaapi.attach_action_to_menu(self.python_menu_path, self.python_action_name, idaapi.SETMENU_APP)
            # find vulnerabilty action
            vulnFinder_action = idaapi.action_desc_t(self.vulnFinder_action_name,
                                                  'Vulnerability finding',
                                                  FindVulnHandler(),
                                                  "Ctrl+Alt+E",
                                                  "使用 %s 在当前函数中查找漏洞"%(MODEL),
                                                  199)
            idaapi.register_action(vulnFinder_action)
            idaapi.attach_action_to_menu(self.vulnFinder_menu_path, self.vulnFinder_action_name, idaapi.SETMENU_APP)
            # create exploit action
            expPython_action = idaapi.action_desc_t(self.expPython_action_name,
                                                  'Try to generate Exploit',
                                                  expCreateHandler(),
                                                  "",
                                                  "使用 %s 尝试对漏洞函数生成EXP"%(MODEL),
                                                  199)
            idaapi.register_action(expPython_action)
            idaapi.attach_action_to_menu(self.expPython_menu_path, self.expPython_action_name, idaapi.SETMENU_APP)
            # Register context menu actions
            self.menu = ContextMenuHooks()
            self.menu.hook()
            print("Auto-WPeGPT v0.2 is ready.")
            print("%s v2.6 works fine! :)@WPeace\n"%(PLUGIN_NAME))
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        idaapi.detach_action_from_menu(self.autoWPeGPT_menu_path, self.autoWPeGPT_action_name)
        idaapi.detach_action_from_menu(self.explain_menu_path, self.explain_action_name)
        idaapi.detach_action_from_menu(self.rename_menu_path, self.rename_action_name)
        idaapi.detach_action_from_menu(self.python_menu_path, self.python_action_name)
        idaapi.detach_action_from_menu(self.vulnFinder_menu_path, self.vulnFinder_action_name)
        idaapi.detach_action_from_menu(self.expPython_menu_path, self.expPython_action_name)
        if self.menu:
            self.menu.unhook()
        return


def PLUGIN_ENTRY():
    global model_api_key
    if model_api_key == "ENTER_API_KEY_HERE":
        model_api_key = os.getenv("model_api_key")
        if not model_api_key:
            print("未找到 API_KEY，请在脚本中填写 model_api_key! :(@WPeace")
            raise ValueError("No valid OpenAI API key found!")
    return type(f"myplugin_{PLUGIN_NAME}", (myplugin_WPeChatGPT, ), dict())()


repository: wmliang/pe-afl
file: lighthouse/lighthouse/ida_loader.py
wanted_name: Lighthouse
comment: Code Coverage Explorer
created_at: 2019-02-15 08:14:36+00:00
pushed_at: 2020-09-26 08:30:04+00:00
url:
https://github.com/wmliang/pe-afl/blob/4036d2f41da20ff12ecac43a076de5d60ce68bd9/lighthouse/lighthouse/ida_loader.py

content:
import logging

import idaapi
from lighthouse.util.log import lmsg
from lighthouse.ida_integration import LighthouseIDA

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

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

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

class LighthouseIDAPlugin(idaapi.plugin_t):
    """
    The IDA plugin stub for Lighthouse.
    """

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

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_MOD | idaapi.PLUGIN_HIDE
    comment = "Code Coverage Explorer"
    help = ""
    wanted_name = "Lighthouse"
    wanted_hotkey = ""

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

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        try:
            self._lighthouse = LighthouseIDA()
            self._lighthouse.load()
        except Exception as e:
            lmsg("Failed to initialize Lighthouse")
            logger.exception("Exception details:")
            return idaapi.PLUGIN_SKIP

        # tell IDA to keep the plugin loaded (everything is okay)
        return idaapi.PLUGIN_KEEP

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

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """
        try:
            self._lighthouse.unload()
            self._lighthouse = None
        except Exception as e:
            logger.exception("Failed to cleanly unload Lighthouse from IDA.")


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

content:
"""Loads all possible BAP IDA Python plugins."""
import os
import bap.plugins
import bap.utils.run
import idaapi


class bap_loader(idaapi.plugin_t):
    """Loads plugins from the bap/plugins directory."""

    flags = idaapi.PLUGIN_FIX
    comment = "BAP Plugin Loader"
    help = "BAP Plugin Loader"
    wanted_name = "BAP_Plugin_Loader"
    wanted_hotkey = ""

    def init(self):
        """Read directory and load as many plugins as possible."""
        self.plugins = []

        idaapi.msg("BAP Loader activated\n")

        bap.utils.run.check_and_configure_bap()

        plugin_path = os.path.dirname(bap.plugins.__file__)
        idaapi.msg("BAP> Loading plugins from {}\n".format(plugin_path))

        for plugin in sorted(os.listdir(plugin_path)):
            path = os.path.join(plugin_path, plugin)
            if not plugin.endswith('.py') or plugin.startswith('__'):
                continue  # Skip non-plugins
            idaapi.msg('BAP> Loading {}\n'.format(plugin))
            self.plugins.append(idaapi.load_plugin(path))
        return idaapi.PLUGIN_KEEP

    def term(self):
        """Ignored."""
        pass

    def run(self, arg):
        """Ignored."""
        pass


def PLUGIN_ENTRY():
    """Load the bap_loader."""
    return bap_loader()


repository: cra0/ida-scripts
file: plugins/cvutils-gotooffset.py
description: Various IDA scripts I've created for Reverse engineering.
comment: Go to an offset.
created_at: 2022-07-23 02:56:21+00:00
pushed_at: 2024-10-29 02:59:25+00:00
url:
https://github.com/cra0/ida-scripts/blob/ee4ce429b3c059d0441c63e2476a9c0c7d726ccc/plugins/cvutils-gotooffset.py

content:
#------------------------------------------------------------------------------
# IDA Plugin to jump to an offset from the Imagebase.
# Copy the 'cvutils-getoffset.py' into the plugins directory of IDA
#------------------------------------------------------------------------------

VERSION = '1.1.0'
__AUTHOR__ = 'cra0'

PLUGIN_NAME = "Go To Offset"
PLUGIN_HOTKEY = "Shift+G"


  
import os
import sys
import idc
import idaapi
import idautils
import string


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 PLUGIN_ENTRY():
    """
    Required plugin entry point for IDAPython Plugins.
    """
    return cvutils_gotooffset()

class cvutils_gotooffset(idaapi.plugin_t):

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Go to an offset."
    help = "Use the shortcut key to open the goto dialog box."
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY

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

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

        # initialize the menu actions our plugin will inject
        self._init_action_goto_offset()

        # done
        idaapi.msg("%s %s initialized...\n" % (self.wanted_name, VERSION))
        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.
        """


        # unregister our actions & free their resources
        self._del_action_goto_offset()


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


    #--------------------------------------------------------------------------
    # IDA Actions
    #--------------------------------------------------------------------------

    ACTION_GET_OFFSET  = "prefix:goto_offset"


    def _init_action_goto_offset(self):
        """
        Register the copy bytes action with IDA.
        """   
        vaction_desc = "Go to offset."
        if (sys.version_info > (3, 0)):
            # Describe the action using python3 copy
            action_desc = idaapi.action_desc_t(
                self.ACTION_GET_OFFSET,                                     # The action name.
                "Go to offset",                                             # The action text.
                IDACtxEntry(goto_offset),                                # The action handler.
                PLUGIN_HOTKEY,                                              # Optional: action shortcut
                vaction_desc,                                               # Optional: tooltip
                31                                                          # Copy icon
            )
        else:
            # Describe the action using python2 copy
            action_desc = idaapi.action_desc_t(
                self.ACTION_GET_OFFSET,                                 # The action name.
                "Go to offset",                                         # The action text.
                IDACtxEntry(goto_offset),                            # The action handler.
                PLUGIN_HOTKEY,                                          # Optional: action shortcut
                vaction_desc,                                           # Optional: tooltip
                31                                                      # Copy icon
            )


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


    def _del_action_goto_offset(self):
        """
        Delete the bulk prefix action from IDA.
        """
        idaapi.unregister_action(self.ACTION_GET_OFFSET)

#------------------------------------------------------------------------------
# Display a warning message box
#------------------------------------------------------------------------------
def display_warning(message): 
    if newer_version_compatible:
        return idaapi.warning(message)
    else:
        return idc.Warning(message)
   
#------------------------------------------------------------------------------
# Jump to a certain address
#------------------------------------------------------------------------------
def jump_to_address(jump_address): 
    if newer_version_compatible:
        ida_kernwin.jumpto(jump_address)
    else:
        idc.Jump(jump_address)


#------------------------------------------------------------------------------
# Image Min EA
#------------------------------------------------------------------------------
def get_minEA(): 
    if newer_version_compatible:
        return ida_ida.inf_get_min_ea()
    else:
        return idc.MinEA()

#------------------------------------------------------------------------------
# IsBadAddress
#------------------------------------------------------------------------------

def is_hex(s):
     hex_digits = set(string.hexdigits)
     # if s is long, then it is faster to check against a set
     return all(c in hex_digits for c in s)

def is_hex(s):
    try:
        int(s, 16)
        return True
    except ValueError:
        return False

def isvalid_address(ea):
    """Check if the given address is valid
    
    Arguments:
        ea: The linear address to check.
    """   
    if (ea == idaapi.BADADDR):
        print("[%x] BADADDR" % ea)
        return 0
        
    pe_min_ea = get_minEA()
       
    if (ea < pe_min_ea):
        print("[%x] is lower than MinEA [%x]" % (ea, pe_min_ea))
        return 0
               
    if not idaapi.getseg(ea):
        print("[%x] getseg failed" % ea)
        return 0
    
    return 1


#------------------------------------------------------------------------------
# Go to offset
#------------------------------------------------------------------------------

def goto_offset():
    hint_min_offset = get_minEA() - idaapi.get_imagebase()
    
    #Use string for now, force them to use hex as the offset
    offset_value=None
    if newer_version_compatible:
        offset_value = ida_kernwin.ask_str("0x%x" % hint_min_offset, 0, "To Offset[HEX]:")    
    else:
        offset_value = idc.AskStr("0x%x" % hint_min_offset, "To Offset[HEX]:")

    if not offset_value:
        print("No value was provided. Ignoring")
        return
    
    if not is_hex(offset_value):
        print("Bad input; It doesn't contain valid hex")
        display_warning("Bad Input!")
        return
        
    offset_hex = int(offset_value, 16)
    if offset_hex == 0:
        display_warning("Input is Invalid!")
        return
       
    image_base = idaapi.get_imagebase()
    jump_address = image_base + offset_hex
    
    if (isvalid_address(jump_address)):
        print ("Offset [%x] =-> Address [0x%x]" % (offset_hex, jump_address))
        jump_to_address(jump_address)
    else:
        display_warning("Bad Offset!")
    
    return

#------------------------------------------------------------------------------
# IDA ctxt
#------------------------------------------------------------------------------

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: grayhatacademy/ida
file: plugins/alleycat/alleycat.py
wanted_name: AlleyCat
created_at: 2013-10-16 18:52:21+00:00
pushed_at: 2023-10-23 10:38:19+00:00
url:
https://github.com/grayhatacademy/ida/blob/faf13a13c4f5070fe8f0666c1508f928ffe40e23/plugins/alleycat/alleycat.py

content:
import sys
import idc
import time
import idaapi
import idautils

from shims import ida_shims

# This limits the depth of any individual path, as well as the maximum
# number of paths that will be searched; this is needed for practical
# reasons, as IDBs with tens of thousands of functions take a long time
# to exhaust all possible paths without some practical limitation.
#
# This is global so it's easy to change from the IDAPython prompt.
ALLEYCAT_LIMIT = 10000


def add_to_namespace(namespace, source, name, variable):
    '''
    Add a variable to a different namespace, likely __main__.
    '''
    import importlib

    importer_module = sys.modules[namespace]
    if source in list(sys.modules.keys()):
        if not (sys.version_info.major == 3 and sys.version_info.minor >= 4):
            import imp
            imp.reload(sys.modules[source])
        else:
            importlib.reload(sys.modules[source])
    else:
        m = importlib.import_module(source, None)
        sys.modules[source] = m

    setattr(importer_module, name, variable)


class AlleyCatException(Exception):
    pass


class AlleyCat(object):
    '''
    Class which resolves code paths. This is where most of the work is done.
    '''

    def __init__(self, start, end, quiet=False):
        '''
        Class constructor.

        @start - The start address.
        @end   - The end address.

        Returns None.
        '''
        global ALLEYCAT_LIMIT
        self.limit = ALLEYCAT_LIMIT
        self.paths = []
        self.quiet = quiet

        # We work backwards via xrefs, so we start at the end and end at the
        # start
        if not self.quiet:
            print("Generating call paths from %s to %s..." % (self._name(end),
                                                              self._name(start)))
        self._build_paths(start, end)

    def _name(self, ea):
        name = ida_shims.get_name(ea)
        if not name:
            name = ida_shims.get_func_off_str(ea)
            if not name:
                name = '0x%X' % ea
        return name

    def _add_path(self, path):
        if path not in self.paths:
            self.paths.append(path)

    def _build_paths(self, start, end=idc.BADADDR):
        partial_paths = [[start]]

        # Loop while there are still unresolve paths and while all path sizes
        # have not exceeded ALLEYCAT_LIMIT
        while partial_paths and \
                len(self.paths) < self.limit and \
                len(partial_paths) < self.limit:
            callers = set()

            # Callee is the last entry of the first path in partial paths.
            # The first path list will change as paths are completed and
            # popped from the list.
            callee = partial_paths[0][-1]

            # Find all unique functions that reference the callee, assuming this
            # path has not exceeded ALLEYCAT_LIMIT.
            if len(partial_paths[0]) < self.limit:
                for xref in idautils.XrefsTo(callee):
                    caller = self._get_code_block(xref.frm)
                    if caller:
                        start_ea = ida_shims.start_ea(caller)

                        if start_ea not in callers:
                            callers.add(start_ea)

            # If there are callers to the callee, remove the callee's current
            # path and insert new ones with the new callers appended.
            if callers:
                base_path = partial_paths.pop(0)
                for caller in callers:

                    # Don't want to loop back on ourselves in the same path
                    if caller in base_path:
                        continue

                    # If we've reached the desired end node, don't go any
                    # further down this path
                    if caller == end:
                        self._add_path((base_path + [caller])[::-1])
                    else:
                        partial_paths.append(base_path + [caller])
            # Else, our end node is not in this path, so don't include it in the
            # finished path list.
            elif end not in partial_paths[0]:
                partial_paths.pop(0)
            # If there were no callers then this path has been exhausted and
            # should be popped from the partial path list into the finished path list.
            elif end in partial_paths[0]:
                # Paths start with the end function and end with the start
                # function; reverse it.
                self._add_path(partial_paths.pop(0)[::-1])

    def _get_code_block(self, ea):
        return idaapi.get_func(ea)


class AlleyCatFunctionPaths(AlleyCat):

    def __init__(self, start_ea, end_ea, quiet=False):

        # We work backwards via xrefs, so we start at the end and end at the start
        try:
            func = idaapi.get_func(end_ea)
            start = ida_shims.start_ea(func)
        except:
            raise AlleyCatException("Address 0x%X is not part of a function!" %
                                    end_ea)
        try:
            func = idaapi.get_func(start_ea)
            end = ida_shims.start_ea(func)
        except:
            end = idc.BADADDR

        super(AlleyCatFunctionPaths, self).__init__(start, end, quiet)


class AlleyCatCodePaths(AlleyCat):

    def __init__(self, start_ea, end_ea, quiet=False):
        end_func = idaapi.get_func(end_ea)
        start_func = idaapi.get_func(start_ea)

        if not start_func:
            raise AlleyCatException("Address 0x%X is not part of a function!" %
                                    start_ea)
        if not end_func:
            raise AlleyCatException("Address 0x%X is not part of a function!" %
                                    end_ea)

        start_func_ea = ida_shims.start_ea(start_func)
        end_func_ea = ida_shims.start_ea(end_func)

        if start_func_ea != end_func_ea:
            raise AlleyCatException("The start and end addresses are not part "
                                    "of the same function!")

        self.func = start_func
        self.blocks = [block for block in idaapi.FlowChart(self.func)]

        # We work backwards via xrefs, so we start at the end and end at the start
        end_block = self._get_code_block(start_ea)
        start_block = self._get_code_block(end_ea)

        if not end_block:
            raise AlleyCatException("Failed to find the code block associated "
                                    "with address 0x%X" % start_ea)
        if not start_block:
            raise AlleyCatException("Failed to find the code block associated "
                                    "with address 0x%X" % end_ea)

        start_block_ea = ida_shims.start_ea(start_block)
        end_block_ea = ida_shims.start_ea(end_block)

        super(AlleyCatCodePaths, self).__init__(
            start_block_ea, end_block_ea, quiet)

    def _get_code_block(self, ea):
        for block in self.blocks:
            start_ea = ida_shims.start_ea(block)
            end_ea = ida_shims.end_ea(block)
            if start_ea <= ea and end_ea > ea:
                return block
        return None


# Everything below here is just IDA UI/Plugin stuff


class AlleyCatGraphHistory(object):
    '''
    Manages include/exclude graph history.
    '''

    INCLUDE_ACTION = 0
    EXCLUDE_ACTION = 1

    def __init__(self):
        self.history = []
        self.includes = []
        self.excludes = []
        self.history_index = 0
        self.include_index = 0
        self.exclude_index = 0

    def reset(self):
        self.history = []
        self.includes = []
        self.excludes = []
        self.history_index = 0
        self.include_index = 0
        self.exclude_index = 0

    def update_history(self, action):
        if self.excludes and len(self.history)-1 != self.history_index:
            self.history = self.history[0:self.history_index+1]
        self.history.append(action)
        self.history_index = len(self.history)-1

    def add_include(self, obj):
        if self.includes and len(self.includes)-1 != self.include_index:
            self.includes = self.includes[0:self.include_index+1]
        self.includes.append(obj)
        self.include_index = len(self.includes)-1
        self.update_history(self.INCLUDE_ACTION)

    def add_exclude(self, obj):
        if len(self.excludes)-1 != self.exclude_index:
            self.excludes = self.excludes[0:self.exclude_index+1]
        self.excludes.append(obj)
        self.exclude_index  = len(self.excludes)-1
        self.update_history(self.EXCLUDE_ACTION)

    def get_includes(self):
        return set(self.includes[0:self.include_index+1])

    def get_excludes(self):
        return set(self.excludes[0:self.exclude_index+1])

    def undo(self):
        if self.history:
            if self.history[self.history_index] == self.INCLUDE_ACTION:
                if self.include_index >= 0:
                    self.include_index -= 1
            elif self.history[self.history_index] == self.EXCLUDE_ACTION:
                if self.exclude_index >= 0:
                    self.exclude_index -= 1

            self.history_index -= 1
            if self.history_index < 0:
                self.history_index = 0

    def redo(self):
        self.history_index += 1
        if self.history_index >= len(self.history):
            self.history_index = len(self.history)-1

        if self.history[self.history_index] == self.INCLUDE_ACTION:
            if self.include_index < len(self.includes)-1:
                self.include_index += 1
        elif self.history[self.history_index] == self.EXCLUDE_ACTION:
            if self.exclude_index < len(self.excludes)-1:
                self.exclude_index += 1


class AlleyCatGraph(idaapi.GraphViewer):
    '''
    Displays the graph and manages graph actions.
    '''
    def __init__(self, results, title="AlleyCat Graph"):
        idaapi.GraphViewer.__init__(self, title)
        self.results = results

        self.nodes_ea2id = {}
        self.nodes_id2ea = {}
        self.edges = {}
        self.end_nodes = []
        self.edge_nodes = []
        self.start_nodes = []

        self.history = AlleyCatGraphHistory()
        self.include_on_click = False
        self.exclude_on_click = False

        self.cmd_undo = None
        self.cmd_redo = None
        self.cmd_reset = None
        self.cmd_exclude = None
        self.cmd_include = None
        self.cmd_unhighlight = None

    def Show(self):
        '''
        Display the graph.

        Returns True on success, False on failure.
        '''
        if not idaapi.GraphViewer.Show(self):
            return False
        else:
            self.cmd_undo = self.AddCommand("Undo", "")
            self.cmd_redo = self.AddCommand("Redo", "")
            self.cmd_reset = self.AddCommand("Reset graph", "")
            self.cmd_exclude = self.AddCommand("Exclude node", "")
            self.cmd_include = self.AddCommand("Include node", "")
            self.cmd_unhighlight = self.AddCommand(
                "Temporarily un-highlight all paths", "")
            return True

    def OnRefresh(self):
        # Clear the graph before refreshing
        self.clear()
        self.nodes_ea2id = {}
        self.nodes_id2ea = {}
        self.edges = {}
        self.end_nodes = []
        self.edge_nodes = []
        self.start_nodes = []

        includes = self.history.get_includes()
        excludes = self.history.get_excludes()

        for path in self.results:
            parent_node = None

            # Check to see if this path contains all nodes marked for explicit
            # inclusion
            if (set(path) & includes) != includes:
                continue

            # Check to see if this path contains any nodes marked for explicit
            # exclusion
            if (set(path) & excludes) != set():
                continue

            for ea in path:
                # If this node already exists, use its existing node ID
                if ea in self.nodes_ea2id:
                    this_node = self.nodes_ea2id[ea]
                # Else, add this node to the graph
                else:
                    this_node = self.AddNode(self.get_name_by_ea(ea))
                    self.nodes_ea2id[ea] = this_node
                    self.nodes_id2ea[this_node] = ea

                # If there is a parent node, add an edge between the parent node
                # and this one
                if parent_node is not None:
                    self.AddEdge(parent_node, this_node)
                    if this_node not in self.edges[parent_node]:
                        self.edges[parent_node].append(this_node)

                # Update the parent node for the next loop
                parent_node = this_node
                if parent_node not in self.edges:
                    self.edges[parent_node] = []

                # Highlight this node in the disassembly window
                self.highlight(ea)

            try:
                # Track the first, last, and next to last nodes in each path for
                # proper colorization in self.OnGetText.
                self.start_nodes.append(self.nodes_ea2id[path[0]])
                self.end_nodes.append(self.nodes_ea2id[path[-1]])
                self.edge_nodes.append(self.nodes_ea2id[path[-2]])
            except:
                pass

        return True

    def OnGetText(self, node_id):
        color = idc.DEFCOLOR

        if node_id in self.edge_nodes:
            color = 0x00ffff
        elif node_id in self.start_nodes:
            color = 0x00ff00
        elif node_id in self.end_nodes:
            color = 0x0000ff

        return self[node_id], color

    def OnHint(self, node_id):
        hint = ""

        try:
            for edge_node in self.edges[node_id]:
                hint += "%s\n" % self[edge_node]
        except:
            pass

        return hint

    def OnCommand(self, cmd_id):
        if self.cmd_undo == cmd_id:
            if self.include_on_click or self.exclude_on_click:
                self.include_on_click = False
                self.exclude_on_click = False
            else:
                self.history.undo()
            self.Refresh()
        elif self.cmd_redo == cmd_id:
            self.history.redo()
            self.Refresh()
        elif self.cmd_include == cmd_id:
            self.include_on_click = True
        elif self.cmd_exclude == cmd_id:
            self.exclude_on_click = True
        elif self.cmd_reset == cmd_id:
            self.include_on_click = False
            self.exclude_on_click = False
            self.history.reset()
            self.Refresh()
        elif self.cmd_unhighlight == cmd_id:
            self.unhighlight_all()

    def OnClick(self, node_id):
        if self.include_on_click:
            self.history.add_include(self.nodes_id2ea[node_id])
            self.include_on_click = False
        elif self.exclude_on_click:
            self.history.add_exclude(self.nodes_id2ea[node_id])
            self.exclude_on_click = False
        self.Refresh()

    def OnDblClick(self, node_id):
        xref_locations = []
        node_ea = self.get_ea_by_name(self[node_id])

        if node_id in self.edges:
            for edge_node_id in self.edges[node_id]:

                edge_node_name = self[edge_node_id]
                edge_node_ea = self.get_ea_by_name(edge_node_name)

                if edge_node_ea != idc.BADADDR:
                    for xref in idautils.XrefsTo(edge_node_ea):
                        # Is the specified node_id the source of this xref?
                        if self.match_xref_source(xref, node_ea):
                            xref_locations.append((xref.frm, edge_node_ea))

        if xref_locations:
            xref_locations.sort()

            print("")
            print("Path Xrefs from %s:" % self[node_id])
            print("-" * 100)
            for (xref_ea, dst_ea) in xref_locations:
                print("%-50s  =>  %s" % (self.get_name_by_ea(xref_ea),
                                         self.get_name_by_ea(dst_ea)))
            print("-" * 100)
            print("")

            ida_shims.jumpto(xref_locations[0][0])
        else:
            ida_shims.jumpto(node_ea)

    def OnClose(self):
        if ida_shims.ask_yn(1, "Path nodes have been highlighted in the "
                               "disassembly window. Undo highlighting?") == 1:
            self.unhighlight_all()

    def match_xref_source(self, xref, source):
        return ((xref.type != idc.fl_F) and
                (ida_shims.get_func_attr(xref.frm, idc.FUNCATTR_START) == source))

    def get_ea_by_name(self, name):
        '''
        Get the address of a location by name.

        @name - Location name

        Returns the address of the named location, or idc.BADADDR on failure.
        '''
        # This allows support of the function offset style names (e.g., main+0C)
        ea = 0
        if '+' in name:
            (func_name, offset) = name.split('+')
            base_ea = ida_shims.get_name_ea_simple(func_name)
            if base_ea != idc.BADADDR:
                try:
                    ea = base_ea + int(offset, 16)
                except:
                    ea = idc.BADADDR
        else:
            ea = ida_shims.get_name_ea_simple(name)
            if ea == idc.BADADDR:
                try:
                    ea = int(name, 0)
                except:
                    ea = idc.BADADDR

        return ea

    def clear(self):
        # Clears the graph and unhighlights the disassembly
        self.Clear()
        self.unhighlight_all()

    def get_name_by_ea(self, ea):
        '''
        Get the name of the specified address.

        @ea - Address.

        Returns a name for the address, one of idc.Name, idc.GetFuncOffset or
        0xXXXXXXXX.
        '''
        name = ida_shims.get_name(ea)
        if not name:
            name = ida_shims.get_func_off_str(ea)
            if not name:
                name = "0x%X" % ea
        return name

    def colorize_node(self, ea, color):
        # Colorizes an entire code block
        func = idaapi.get_func(ea)
        if func:
            for block in idaapi.FlowChart(func):
                block_start_ea = ida_shims.start_ea(block)
                block_end_ea = ida_shims.end_ea(block)

                if block_start_ea <= ea and block_end_ea > ea:
                    ea = block_start_ea
                    while ea < block_end_ea:
                        idaapi.set_item_color(ea, color)
                        ea = ida_shims.next_head(ea)
                    break

    def highlight(self, ea):
        # Highlights an entire code block
        self.colorize_node(ea, 0x00FF00)

    def unhighlight(self, ea):
        # Unhighlights an entire code block
        self.colorize_node(ea, idc.DEFCOLOR)

    def unhighlight_all(self):
        # Unhighlights all code blocks
        for path in self.results:
            for ea in path:
                self.unhighlight(ea)


class AlleyCatPaths(object):
    def _current_function(self):
        function = idaapi.get_func(ida_shims.get_screen_ea())
        return ida_shims.start_ea(function)

    def _find_and_plot_paths(self, sources, targets, klass=AlleyCatFunctionPaths):
        results = []

        for target in targets:
            for source in sources:
                s = time.time()
                r = klass(source, target).paths
                e = time.time()
                print("Found %d paths in %f seconds." % (len(r), (e-s)))

                if r:
                    results += r
                else:
                    name = ida_shims.get_name(target)
                    if not name:
                        name = "0x%X" % target
                    print("No paths found to", name)

        if results:
            # Be sure to close any previous graph before creating a new one.
            # Failure to do so may crash IDA.
            try:
                self.graph.Close()
            except:
                pass

            self.graph = AlleyCatGraph(results, 'Path Graph')
            self.graph.Show()

    def _get_user_selected_functions(self, many=False):
        functions = []
        ea = ida_shims.get_screen_ea()
        try:
            current_function = ida_shims.get_func_attr(ea, idc.FUNCATTR_START)
        except:
            current_function = None

        while True:
            function = ida_shims.choose_func(
                "Select a function and click 'OK' until all functions have "
                "been selected. When finished, click 'Cancel' to display the "
                "graph.")

            if ida_shims.get_screen_ea() != ea:
                ida_shims.jumpto(ea)

            if not function or \
                    function == idc.BADADDR or function == current_function:
                break
            elif function not in functions:
                functions.append(function)

            if not many:
                break

        return functions

    def FindPathsToCodeBlock(self):
        target = ida_shims.get_screen_ea()
        source = self._current_function()

        if source:
            self._find_and_plot_paths(
                [source], [target], klass=AlleyCatCodePaths)

    def FindPathsToMany(self):
        source = self._current_function()

        if source:
            targets = self._get_user_selected_functions(many=True)
            if targets:
                self._find_and_plot_paths([source], targets)

    def FindPathsFromMany(self):
        target = self._current_function()

        if target:
            sources = self._get_user_selected_functions(many=True)
            if sources:
                self._find_and_plot_paths(sources, [target])


# Helper functions to execute commands selected from dropdown menus.
# args parameter is required for IDA version < 7.0
def find_paths_from_many(arg=None):
    AlleyCatPaths().FindPathsFromMany()


def find_paths_to_many(arg=None):
    AlleyCatPaths().FindPathsToMany()


def find_paths_to_code_block(args=None):
    AlleyCatPaths().FindPathsToCodeBlock()


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

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

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


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

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

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

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

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

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


class idapathfinder_t(idaapi.plugin_t):
    flags = 0
    comment = ''
    help = ''
    wanted_name = 'AlleyCat'
    wanted_hotkey = ''
    menu_name = 'View/Graphs/'
    menu_context = []

    to_from_action_name = 'tocurrfrom:action'
    from_to_action_name = 'fromcurrto:action'
    curr_func_curr_block_action_name = 'currfunccurrblock:action'

    to_from_menu_name = 'Find paths to the current function from...'
    from_to_menu_name = 'Find paths from the current function to...'
    curr_func_curr_block_menu_name = 'Find paths in the current function to ' \
                                     'the current code block'

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            # Add ALLEYCAT_LIMIT variable to the global namespace so it can be
            # accessed from the IDA python terminal.
            global ALLEYCAT_LIMIT
            add_to_namespace(
                '__main__', 'alleycat', 'ALLEYCAT_LIMIT', ALLEYCAT_LIMIT)

            # Add functions to global namespace.
            add_to_namespace(
                '__main__', 'alleycat', 'AlleyCatFunctionPaths',
                AlleyCatFunctionPaths)
            add_to_namespace(
                '__main__', 'alleycat', 'AlleyCatCodePaths', AlleyCatCodePaths)
            add_to_namespace(
                '__main__', 'alleycat', 'AlleyCatGraph', AlleyCatGraph)

            to_curr_from_desc = idaapi.action_desc_t(
                self.to_from_action_name, self.to_from_menu_name,
                ToCurrentFromAction(), self.wanted_hotkey,
                'Find paths to the current function from...', 199)

            from_curr_to_desc = idaapi.action_desc_t(
                self.from_to_action_name, self.from_to_menu_name,
                FromCurrentToAction(), self.wanted_hotkey,
                'Find paths from the current function to...', 199)

            curr_func_to_block_desc = idaapi.action_desc_t(
                self.curr_func_curr_block_action_name,
                self.curr_func_curr_block_menu_name,
                InCurrentFunctionToCurrentCodeBlockAction(),
                self.wanted_hotkey,
                'Find paths in the current function to the current code block',
                199)

            idaapi.register_action(to_curr_from_desc)
            idaapi.register_action(from_curr_to_desc)
            idaapi.register_action(curr_func_to_block_desc)

            idaapi.attach_action_to_menu(
                self.menu_name, self.to_from_action_name, idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu(
                self.menu_name, self.from_to_action_name, idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu(
                self.menu_name, self.curr_func_curr_block_action_name,
                idaapi.SETMENU_APP)

        else:
            self.menu_context.append(
                idaapi.add_menu_item(
                    self.menu_name, self.to_from_menu_name, "", 0,
                    find_paths_from_many, (None,)))

            self.menu_context.append(
                idaapi.add_menu_item(
                    self.menu_name, self.from_to_menu_name, "", 0,
                    find_paths_to_many, (None,)))

            self.menu_context.append(
                idaapi.add_menu_item(
                    self.menu_name, self.curr_func_curr_block_menu_name, "", 0,
                    find_paths_to_code_block, (None,)))

        return idaapi.PLUGIN_KEEP

    def term(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.detach_action_from_menu(
                self.menu_name, self.to_from_action_name)
            idaapi.detach_action_from_menu(
                self.menu_name, self.from_to_action_name)
            idaapi.detach_action_from_menu(
                self.menu_name, self.curr_func_curr_block_action_name)
        else:
            for context in self.menu_context:
                idaapi.del_menu_item(context)
        return None

    def run(self, arg):
        pass


def PLUGIN_ENTRY():
    return idapathfinder_t()


repository: LeadroyaL/IDA_ARM_Unwind
file: arm_unwind_plugin.py
description: IDA plugin, unwind stack trace when debugging arm.
wanted_name: arm unwind plugin
comment: arm unwind plugin
created_at: 2020-08-13 12:11:48+00:00
pushed_at: 2020-08-20 09:51:06+00:00
url:
https://github.com/LeadroyaL/IDA_ARM_Unwind/blob/9a44a7da0ea1bee79c028107e0a72d031ecfd2ed/arm_unwind_plugin.py

content:
# coding=utf-8
from idc import *
from idaapi import *
from PyQt5 import QtWidgets, QtGui

from io import BytesIO
from elftools.elf.elffile import ELFFile
from elftools.elf.structs import ELFStructs
from elftools.ehabi.ehabiinfo import EHABIInfo


class ArmUnwinder:
    def __init__(self):
        self.unwind_caches = list()
        self.pid = get_event_pid()

    def check_pid(self, pid):
        if pid == self.pid:
            return
        print "Pid changed, refresh arm unwind caches."
        self.unwind_caches = list()
        self.pid = pid

    def unwind_all(self):
        self.check_pid(get_event_pid())
        _status = VRSStatus()
        _status.load_ida_registers()
        pc_list = list()
        pc_list.append(GetRegValue("PC"))
        while self.unwind_single_frame(_status):
            pc_value = _status.getRx(_status.PC_INDEX)
            pc_list.append(pc_value)
        return pc_list

    def unwind_single_frame(self, _status):
        # 1. 拿PC
        pc = _status.getRx(_status.PC_INDEX)
        # 2. 查cache
        unwind_item = None
        for unwind_cache in self.unwind_caches:
            if unwind_cache.startEA <= pc < unwind_cache.endEA:
                unwind_item = unwind_cache
                print "found cache", unwind_cache.name
                break
        if unwind_item is None:
            # 2.5 查不到cache就去解析内存
            cache = self.create_unwind_cache(pc)
            if cache is None:
                print "cannot unwind this pc", '%0.8x' % pc
            else:
                print "add", cache.name, "into cache"
                self.unwind_caches.append(cache)
            unwind_item = cache

        # 3. 使用cache
        if unwind_item is None:
            return False
        else:
            entries = unwind_item.entries
            relative_pc = pc - unwind_item.startEA
            for i in range(len(entries) - 1):
                if entries[i].function_offset <= relative_pc < entries[i + 1].function_offset:
                    print 'find entry', entries[i]
                    # 4. 如果entry不包含bytecode
                    if entries[i].bytecode_array is None:
                        print "cannot unwind", entries[i]
                        return False
                    # 4.5 进行unwind
                    if _interp(entries[i].bytecode_array, _status) == URC_FAILURE:
                        raise RuntimeError("interpret fail")
                    lr_value = _status.getRx(_status.LR_INDEX)
                    if lr_value == 0:
                        return False
                    # 5. lr -> pc
                    _status.setRx(_status.PC_INDEX, _lr2pc(lr_value))
                    print 'pc', '%0.8x' % _status.getRx(_status.PC_INDEX)
                    return True
        print "Cannot find entry for pc", '%0.8x' % pc
        return False

    @staticmethod
    def create_unwind_cache(_pc):
        """根据当前 pc，解析对应的 elf，拿到 exception handler entry 表"""
        # 找到 ELF 头
        _header_seg, _seg_name = _getELFHeader(_pc)
        if not _checkELFHeader(_header_seg):
            print "cannot determine ELF header of this module", _seg_name
            return None
        # 稍微加载一小段，用于获得 program headers，没必要加载完
        _header_data = DbgRead(_header_seg.start_ea, min(0x1000, _header_seg.size()))
        _stream = BytesIO(_header_data)

        _elf = MemoryELFFile(_stream)
        _exidx_segment = _elf.get_exidx_segment()
        if _exidx_segment is None:
            print 'cannot get PT_ARM_EXIDX of this module', _seg_name
            return None
        _exidx_offset = _exidx_segment['p_offset']
        _exidx_size = _exidx_segment['p_memsz']
        _load_size = _elf.get_load_segment()['p_memsz']
        _stream.close()

        # 加载整个内存，需要获得 .arm.exidx 和 .arm.extab 的数据
        _load_data = DbgRead(_header_seg.start_ea, _load_size)
        _stream = BytesIO(_load_data)
        _ehabi_info = EHABIInfo(FakeSection(_stream, _exidx_offset, _exidx_size), _elf.little_endian)
        _unwind_entries = [_ehabi_info.get_entry(i) for i in range(_ehabi_info.num_entry())]
        _stream.close()
        return UnwindItem(_header_seg.start_ea, _header_seg.start_ea + _load_size, _seg_name, _unwind_entries)

    @staticmethod
    def pc2frame(_pc):
        name = get_segm_name(getseg(_pc))
        if name == '.text':
            name = get_root_filename()
            base = get_imagebase()
        else:
            base = get_segm_start(_pc)
        relative_pc = _pc - base
        fullpath = get_module_name(base)
        if fullpath == 0:
            fullpath = ""
        return Frame(relative_pc, _pc, name, fullpath)

    @staticmethod
    def pclist2framelist(_pclist):
        return [ArmUnwinder.pc2frame(_pc) for _pc in _pclist]


def pc2func(_pc):
    _func = get_func(_pc)
    if _func is None:
        return ""
    return "%s + %d" % (get_func_name(_pc), _pc - _func.startEA)


class Frame:
    def __init__(self, relative_pc, absolute_pc, name, fullpath):
        self.relative_pc = relative_pc
        self.func_with_offset = pc2func(absolute_pc)
        self.absolute_pc = absolute_pc
        self.name = name
        self.fullpath = fullpath

    def __repr__(self):
        return "%s-%s-%0.8x-%0.8x" % (self.name, self.fullpath, self.relative_pc, self.absolute_pc)


class MemoryELFFile(ELFFile):
    """精简过的 ELFFile，去除里面多余的功能"""

    def __init__(self, stream):
        self.stream = stream
        self._identify_file()
        self.structs = ELFStructs(
            little_endian=self.little_endian,
            elfclass=self.elfclass)

        self.structs.create_basic_structs()
        self.header = self._parse_elf_header()
        self.structs.e_type = self['e_type']
        self.structs.e_machine = self['e_machine']
        self.structs.e_ident_osabi = self['e_ident']['EI_OSABI']
        self.structs._create_phdr()

    def get_load_segment(self):
        return self._get_first_segment_by_type('PT_LOAD')

    def get_exidx_segment(self):
        return self._get_first_segment_by_type('PT_ARM_EXIDX')

    def _get_first_segment_by_type(self, _type):
        for i in range(self.num_segments()):
            seg_header = self._get_segment_header(i)
            seg_type = seg_header['p_type']
            if seg_type == _type:
                return seg_header


class FakeSection:
    """因为无法从内存中恢复完整的 Section，只能手动创建虚假的 Section，填充关键数据"""

    def __init__(self, stream, offset, size):
        self.stream = stream
        self.offset = offset
        self.size = size

    def __getitem__(self, item):
        if item == 'sh_offset':
            return self.offset
        elif item == 'sh_size':
            return self.size


class UnwindItem:
    def __init__(self, _startEA, _endEA, _seg_name, _unwind_entries):
        self.startEA = _startEA
        self.endEA = _endEA
        self.name = _seg_name
        self.entries = _unwind_entries


def _getELFHeader(_pc):
    """
        parameter pc
        return (segment, module_name)
    """
    name = get_segm_name(getseg(_pc))
    if name == '.text':
        # 如果是.text，说明就是正在被分析的文件
        return getseg(get_imagebase()), get_root_filename()
    else:
        # 寻找同名module，通常第一个就是 R+X而且 7FELF
        return get_segm_by_name(name), name


def _checkELFHeader(_seg):
    DWORD_MAGIC = 0x464c457f
    return DWORD_MAGIC == Dword(_seg.startEA)


def _lr2pc(lr_value):
    """lr是下一条指令，arm是-4的位置，thumb可能是-2和-4，需要都尝试一遍"""
    if lr_value & 1 == 0:
        # ARM, 32byte
        return lr_value - 4
    else:
        # thumb, 16byte or 32byte
        test1 = lr_value - 1 - 4
        test2 = lr_value - 1 - 2
        hit1 = False
        hit2 = False
        pc = None
        if Word(test1) & 0b11110 << 11 == 0b11110 << 11 and Word(test1 + 2) & 0b11 << 14 == 0b11 << 14:
            hit1 = True
        if Word(test2) & 0b010001111 << 7 == 0b010001111 << 7:
            hit2 = True
        if hit1 and hit2:
            print "cannot determine -2 or -4"
        elif hit1:
            pc = test1 + 1
        elif hit2:
            pc = test2 + 1
        else:
            raise RuntimeError("disasm result is not BL/BLX", lr_value)
        return pc


class VRSStatus:
    """解释执行用于存放寄存器状态"""

    def __init__(self):
        self._regs = [0] * 16
        self._double_regs = [0] * 32

    VSP_INDEX = 13
    FP_INDEX = 11
    IP_INDEX = 12
    SP_INDEX = 12
    LR_INDEX = 14
    PC_INDEX = 15

    def load_ida_registers(self):
        """从 IDA 读取初始状态"""
        for i in range(16):
            self._regs[i] = GetRegValue("R%d" % i)

    def vsp_mov(self, data):
        self._regs[self.VSP_INDEX] = data

    def vsp_add(self, data):
        self._regs[self.VSP_INDEX] += data

    def vsp_sub(self, data):
        self._regs[self.VSP_INDEX] -= data

    def pop(self, registers):
        # pop uint32 register
        # pop {R0,R1,R2} -> pop R0;pop R1;pop R2
        for i in range(0, 16):
            if registers & (1 << i) != 0:
                self._regs[i] = Dword(self._regs[self.VSP_INDEX])
                self._regs[self.VSP_INDEX] += 4

    def popD(self, registers):
        # pop uint32 register
        # pop {D0,D1,D2} -> pop D0;pop D1;pop D2
        for i in range(0, 32):
            if registers & (1 << i) != 0:
                self._double_regs[i] = Dword(self._regs[self.VSP_INDEX])
                self._double_regs[self.VSP_INDEX] += 4

    def getRx(self, n):
        return self._regs[n]

    def setRx(self, n, value):
        self._regs[n] = value

    def __repr__(self):
        part1 = """R0: %0.8x, R1: %0.8x, R2: %0.8x, R3: %0.8x
R4: %0.8x, R5: %0.8x, R6: %0.8x, R7: %0.8x
R8: %0.8x, R9: %0.8x, R10:%0.8x, FP: %0.8x
IP: %0.8x, SP: %0.8x, LR: %0.8x, PC: %0.8x""" % (
            self._regs[0], self._regs[1], self._regs[2], self._regs[3],
            self._regs[4], self._regs[5], self._regs[6], self._regs[7],
            self._regs[8], self._regs[9], self._regs[10], self._regs[11],
            self._regs[12], self._regs[13], self._regs[14], self._regs[15],)
        if any(s != 0 for s in self._double_regs):
            part2 = """D0: %0.8x, D1: %0.8x, D2: %0.8x, D3: %0.8x
D4: %0.8x, D5: %0.8x, D6: %0.8x, D7: %0.8x
D8: %0.8x, D9: %0.8x, D10: %0.8x, D11: %0.8x
D12: %0.8x, D13: %0.8x, D14: %0.8x, D15: %0.8x
D16: %0.8x, D17: %0.8x, D18: %0.8x, D19: %0.8x
D20: %0.8x, D21: %0.8x, D22: %0.8x, D23: %0.8x
D24: %0.8x, D25: %0.8x, D26: %0.8x, D27: %0.8x
D28: %0.8x, D29: %0.8x, D30: %0.8x, D31: %0.8x""" % (
                self._double_regs[0], self._double_regs[1], self._double_regs[2], self._double_regs[3],
                self._double_regs[4], self._double_regs[5], self._double_regs[6], self._double_regs[7],
                self._double_regs[8], self._double_regs[9], self._double_regs[10], self._double_regs[11],
                self._double_regs[12], self._double_regs[13], self._double_regs[14], self._double_regs[15],
                self._double_regs[16], self._double_regs[17], self._double_regs[18], self._double_regs[19],
                self._double_regs[20], self._double_regs[21], self._double_regs[22], self._double_regs[23],
                self._double_regs[24], self._double_regs[25], self._double_regs[26], self._double_regs[27],
                self._double_regs[28], self._double_regs[29], self._double_regs[30], self._double_regs[31],
            )
            return "%s\n%s" % (part1, part2)
        else:
            return part1


def _register_mask(start, count_minus_one):
    return ((1 << (count_minus_one + 1)) - 1) << start


URC_SUCCESS = 0
URC_FAILURE = 1


def _interp(bytecode_array, status):
    """
    reference: https://github.com/llvm/llvm-project/blob/master/libunwind/src/Unwind-EHABI.cpp
        _LIBUNWIND_EXPORT _Unwind_Reason_Code 
        _Unwind_VRS_Interpret(_Unwind_Context *context, const uint32_t *data, size_t offset, size_t len)
    """
    offset = 0
    finish = False
    wrote_pc = False
    bc_len = len(bytecode_array)
    while offset < bc_len and not finish:
        b = int(bytecode_array[offset])
        offset += 1
        if b & 0x80 == 0:
            if b & 0x40 != 0:
                status.vsp_sub(((b & 0x3f) << 2) + 4)
            else:
                status.vsp_add((b << 2) + 4)
        else:
            if b & 0xf0 == 0x80:
                registers = ((b & 0xf) << 12) | (bytecode_array[offset] << 4)
                offset += 1
                if registers == 0:
                    return URC_FAILURE
                if registers & (1 << 15) != 0:
                    wrote_pc = True
                status.pop(registers)
            elif b & 0xf0 == 0x90:
                reg = b & 0x0f
                if reg == 13 or reg == 15:
                    return URC_FAILURE
                status.vsp_mov(status.getRx(reg))
            elif b & 0xf0 == 0xa0:
                registers = _register_mask(4, b & 0x07)
                if b & 0x08 != 0:
                    registers |= 1 << 14
                status.pop(registers)
            elif b & 0xf0 == 0xb0:
                if b == 0xb0:
                    finish = True
                elif b == 0xb1:
                    registers = bytecode_array[offset]
                    offset += 1
                    if registers & (1 << 15) != 0 or registers == 0:
                        return URC_FAILURE
                    status.pop(registers)
                elif b == 0xb2:
                    addend = 0
                    shift = 0
                    while True:
                        b = bytecode_array[offset]
                        offset += 1
                        addend |= (b & 0x7f) << shift
                        if b & 0x80 == 0:
                            break
                        shift += 7
                    status.vsp_add(0x204 + (addend << 2))
                elif b == 0xb3:
                    b = bytecode_array[offset]
                    offset += 1
                    start = ((b & 0xf0) >> 4)
                    count = ((b & 0x0f) >> 0)
                    status.popD(_register_mask(start, count))
                elif b == 0xb4 or b == 0xb5 or b == 0xb6 or b == 0xb7:
                    return URC_FAILURE
                else:
                    b = bytecode_array[offset]
                    offset += 1
                    status.popD(_register_mask(8, b & 0x07))
            elif b & 0xf0 == 0xc0:
                if b == 0xc0 or b == 0xc1 or b == 0xc2 or b == 0xc3 or b == 0xc4 or b == 0xc5:
                    raise RuntimeError("unsupport wmmxd")
                elif b == 0xc6:
                    raise RuntimeError("unsupport wmmxd")
                elif b == 0xc7:
                    raise RuntimeError("unsupport wmmxd")
                elif b == 0xc8 or b == 0xc9:
                    b = bytecode_array[offset]
                    offset += 1
                    start = 16 + ((b & 0xf0) >> 4)
                    count = ((b & 0x0f) >> 0)
                    status.popD(_register_mask(start, count))
                else:
                    return URC_FAILURE
            elif b & 0xf0 == 0xd0:
                if b & 0x80 != 0:
                    return URC_FAILURE
                status.popD(_register_mask(8, b & 0x07))
            else:
                return URC_FAILURE
    if not wrote_pc:
        status.setRx(status.IP_INDEX, status.getRx(status.LR_INDEX))
    print "interp done, now list registers"
    print status
    return URC_SUCCESS


class ArmUnwindView(PluginForm):
    """GUI container"""
    cols = ['Module', 'Path', 'Position', 'Address']

    def __init__(self, frames):
        super(ArmUnwindView, self).__init__()
        self.tree = None
        self.frames = frames

    def dblclick(self, item):
        """Handle double click event."""
        try:
            if get_imagebase() == 0:
                jumpto(int(item.text(2).split('(')[0], 16))
            else:
                jumpto(int(item.text(3), 16))
        except:
            pass

    def load_data(self):
        x = lambda ea: '0x%X' % ea

        for frame in self.frames:
            item = QtWidgets.QTreeWidgetItem(self.tree)
            item.setText(0, frame.name)
            item.setText(1, frame.fullpath)
            item.setText(2, x(frame.relative_pc) + '(' + frame.func_with_offset + ')')
            item.setText(3, x(frame.absolute_pc))

    def OnCreate(self, form):
        """Called when the plugin form is created"""

        self.parent = self.FormToPyQtWidget(form)

        self.tree = QtWidgets.QTreeWidget()
        self.tree.setColumnCount(len(self.cols))
        self.tree.setHeaderLabels(self.cols)
        self.tree.itemDoubleClicked.connect(self.dblclick)

        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.tree)
        layout.setSpacing(0)
        layout.setContentsMargins(0, 0, 0, 0)

        self.load_data()

        self.tree.setColumnWidth(0, 256)
        self.tree.setColumnWidth(1, 256)
        self.tree.setColumnWidth(2, 256)
        self.tree.setColumnWidth(3, 256)
        self.parent.setLayout(layout)

    def OnClose(self, form):
        """Called when the plugin form is closed."""
        del self

    def Show(self):
        """Creates the form is not created or focuses it if it was."""
        return PluginForm.Show(self, 'ArmUnwind 0x%0.8x' % GetRegValue('PC'))


g_arm_unwinder = ArmUnwinder()


def get_unwinder():
    """exported api"""
    return g_arm_unwinder


def unwind_now(show_gui):
    """exported api"""
    if not is_debugger_on():
        warning("Please run script after attach target process.")
        return None
    if get_inf_structure().procName != "ARM":
        warning("This plugin only supports ARM32 CPU.")
        return None
    _pcs = get_unwinder().unwind_all()
    _frames = ArmUnwinder.pclist2framelist(_pcs)
    if show_gui:
        _view = ArmUnwindView(_frames)
        _view.Show()
    return _frames


def unwind_with_gui():
    """exported api"""
    return unwind_now(True)


def unwind_without_gui():
    """exported api"""
    return unwind_now(False)


ARM_UNWIND_HOTKEY = "Ctrl-Shift-U"


def unwind_add_hotkey(hotkey=ARM_UNWIND_HOTKEY):
    """exported api"""
    add_hotkey(hotkey, unwind_with_gui)


class ArmUnwindPlugin(plugin_t):
    """Class that is required for the code to be recognized as
    a plugin by IDA."""
    flags = 0
    comment = "arm unwind plugin"
    help = comment
    wanted_name = "arm unwind plugin"
    wanted_hotkey = ARM_UNWIND_HOTKEY

    def init(self):
        print "ArmUnwindPlugin init success."
        return PLUGIN_OK

    def run(self, arg):
        unwind_with_gui()

    def term(self):
        pass


def PLUGIN_ENTRY():
    return ArmUnwindPlugin()


repository: ke0z/VulChatGPT
file: VulChatGPT.py
description: Use IDA PRO HexRays decompiler with OpenAI(ChatGPT) to find possible vulnerabilities in binaries
url: https://github.com/ke0z/VulChatGPT/blob/e9dfeb7a289878f33695d60c658172508e42a984/VulChatGPT.py
wanted_name: VulChat
comment: Uses davinci-003 to find vulnerabilites the decompiler's output
created_at: 2022-12-27 11:39:45+00:00
pushed_at: 2023-02-23 16:01:26+00:00
content:
import functools
import json
import idaapi
import ida_hexrays
import ida_kernwin
import idc
import openai
import os
import re
import textwrap
import threading

# Set your API key here, or put in in the OPENAI_API_KEY environment variable.
openai.api_key = "ENTER_OPEN_API_KEY_HERE"

# =============================================================================
# Setup the context menu and hotkey in IDA
# =============================================================================

class VulChatPlugin(idaapi.plugin_t):
    flags = 0 
    vuln_action_name = "vulchat:vuln_function"
    vuln_menu_path = "Edit/VulChat/Find Possible Vulnerability"
    expl_action_name = "vulchat:expl_function"
    expl_menu_path = "Edit/Vulchat/Write Python Exploit Sample Script"
    explain_action_name = "vulchat:explain_function"
    explain_menu_path = "Edit/Vulchat/Explain the following Code"
    rename_action_name = "vulchat:rename_function"
    rename_menu_path = "Edit/Vulchat/Rename Variables and Functions"
    wanted_name = 'VulChat'
    wanted_hotkey = ''
    comment = "Uses davinci-003 to find vulnerabilites the decompiler's output"
    help = "See usage instructions on GitHub"
    menu = None

    def init(self):
        # Check whether the decompiler is available
        if not ida_hexrays.init_hexrays_plugin():
            return idaapi.PLUGIN_SKIP
        
        # Function explaining action
        explain_action = idaapi.action_desc_t(self.explain_action_name,
                                              'Explain function',
                                              ExplainHandler(),
                                              "Ctrl+Alt+G",
                                              'Use davinci-003 to explain the currently selected function',
                                              199)
        idaapi.register_action(explain_action)
        idaapi.attach_action_to_menu(self.explain_menu_path, self.explain_action_name, idaapi.SETMENU_APP)

        # Variable renaming action
        rename_action = idaapi.action_desc_t(self.rename_action_name,
                                             'Rename variables',
                                             RenameHandler(),
                                             "Ctrl+Alt+R",
                                             "Use davinci-003 to rename this function's variables",
                                             199)
        idaapi.register_action(rename_action)
        idaapi.attach_action_to_menu(self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)

        #Function vulnerability Checker
        vuln_action = idaapi.action_desc_t(self.vuln_action_name,
                                           'Find possible vulnerability in function',
                                           VulnHandler(),
                                           "Ctrl+Alt+V",
                                           "Use davinci-003 to find possible vulnerability in decompiled function",
                                           199)
        idaapi.register_action(vuln_action)
        idaapi.attach_action_to_menu(self.vuln_menu_path, self.vuln_action_name, idaapi.SETMENU_APP)

        #Function Exploit Creator
        exploit_action = idaapi.action_desc_t(self.expl_action_name,
                                              'Create Sample Python Exploit',
                                              ExploitHandler(),
                                              "Ctrl+Alt+X",
                                              "Use davinci-003 to create a sample exploit script in python",
                                              199)
        idaapi.register_action(exploit_action)
        idaapi.attach_action_to_menu(self.expl_menu_path, self.expl_action_name, idaapi.SETMENU_APP)



        # Register context menu actions
        self.menu = ContextMenuHooks()
        self.menu.hook()

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        idaapi.detach_action_from_menu(self.explain_menu_path, self.explain_action_name)
        idaapi.detach_action_from_menu(self.rename_menu_path, self.rename_action_name)
        idaapi.detach_action_from_menu(self.vuln_menu_path, self.vuln_action_name)
        idaapi.detach_action_from_menu(self.expl_menu_path, self.expl_action_name)
        if self.menu:
            self.menu.unhook()
        return

# -----------------------------------------------------------------------------

class ContextMenuHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        # Add actions to the context menu of the Pseudocode view
        if idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE:
            idaapi.attach_action_to_popup(form, popup, VulChatPlugin.vuln_action_name, "VulChat/")
            idaapi.attach_action_to_popup(form, popup, VulChatPlugin.expl_action_name, "VulChat/Create_EXploit")
            idaapi.attach_action_to_popup(form, popup, VulChatPlugin.explain_action_name, "VulChat/Explain_Function")
            idaapi.attach_action_to_popup(form, popup, VulChatPlugin.rename_action_name, "VulChat/Rename_Variables")

# -----------------------------------------------------------------------------

def comment_callback(address, view, response):
    """
    Callback that sets a comment at the given address.
    :param address: The address of the function to comment
    :param view: A handle to the decompiler window
    :param response: The comment to add
    """
    # Add newlines at the end of each sentence.
    response = "\n".join(textwrap.wrap(response, 80, replace_whitespace=False))

    # Add the response as a comment in IDA.
    idc.set_func_cmt(address, response, 0)
    # Refresh the window so the comment is displayed properly
    if view:
        view.refresh_view(False)
    print("davinci-003 query finished!")
    
# -----------------------------------------------------------------------------

class ExplainHandler(idaapi.action_handler_t):
    """
    This handler is tasked with querying davinci-003 for an explanation of the
    given function. Once the reply is received, it is added as a function
    comment.
    """
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async("Can you explain what the following C function does and suggest a better name for it?\n"
                          + str(decompiler_output),
                          functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

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

# -----------------------------------------------------------------------------

def rename_callback(address, view, response):
    """
    Callback that extracts a JSON array of old names and new names from the
    response and sets them in the pseudocode.
    :param address: The address of the function to work on
    :param view: A handle to the decompiler window
    :param response: The response from davinci-003
    """
    j = re.search(r"\{[^}]*?\}", response)
    if not j:
        print(f"Cannot extract valid JSON from the response. Asking the model to fix it...")
        query_model_async("The JSON document provided in this response is invalid. Can you fix it?\n" + response,
                          functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=view))
        return
    try:
        names = json.loads(j.group(0))
    except json.decoder.JSONDecodeError:
        print(f"The JSON document returned is invalid. Asking the model to fix it...")
        query_model_async("Please fix the following JSON document:\n" + j.group(0),
                          functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=view))
        return

    # The rename function needs the start address of the function
    function_addr = idaapi.get_func(address).start_ea

    replaced = []
    for n in names:
        if ida_hexrays.rename_lvar(function_addr, n, names[n]):
            replaced.append(n)

    # Update possible names left in the function comment
    comment = idc.get_func_cmt(address, 0)
    if comment and len(replaced) > 0:
        for n in replaced:
            comment = re.sub(r'\b%s\b' % n, names[n], comment)
        idc.set_func_cmt(address, comment, 0)

    # Refresh the window to show the new names
    if view:
        view.refresh_view(True)
    print(f"davinci-003 query finished! {len(replaced)} variable(s) renamed.")

# -----------------------------------------------------------------------------

class RenameHandler(idaapi.action_handler_t):
    """
    This handler requests new variable names from davinci-003 and updates the
    decompiler's output.
    """
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async("Analyze the following C function:\n" + str(decompiler_output) +
                            "\nSuggest better variable names, reply with a JSON array where keys are the original names"
                            "and values are the proposed names. Do not explain anything, only print the JSON "
                            "dictionary.",
                          functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

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

# -----------------------------------------------------------------------------
class VulnHandler(idaapi.action_handler_t):
    """
    This handler is tasked with querying davinci-003 for a possible check of vulneranilities on a given function.
    Once the reply is received its added to the function as a comment.
    """
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async("Can you find the vulnerabilty in the following C function and suggest the possible way to exploit it?\n"
        + str(decompiler_output),
        functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS
# -----------------------------------------------------------------------------
class ExploitHandler(idaapi.action_handler_t):
    """
    This handler requests a python exploit for the vulnerable function
    """
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        decompiler_ouput = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async("Find the vulnerability in the following C function:\n" + str(decompiler_ouput) + 
        "\nWrite a python one liner to exploit the function",
        functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS   
# =============================================================================
# davinci-003 interaction
# =============================================================================

def query_model(query, cb, max_tokens=2500):
    """
    Function which sends a query to davinci-003 and calls a callback when the response is available.
    Blocks until the response is received
    :param query: The request to send to davinci-003
    :param cb: Tu function to which the response will be passed to.
    """
    try:
        response = openai.Completion.create(
            model="text-davinci-003",
            prompt=query,
            temperature=0.6,
            max_tokens=max_tokens,
            top_p=1,
            frequency_penalty=1,
            presence_penalty=1,
            #timeout=60  # Wait 60 seconds maximum
        )
        ida_kernwin.execute_sync(functools.partial(cb, response=response.choices[0].text), ida_kernwin.MFF_WRITE)
    except openai.InvalidRequestError as e:
        # Context length exceeded. Determine the max number of tokens we can ask for and retry.
        m = re.search(r'maximum context length is (\d+) tokens, however you requested \d+ tokens \((\d+) in your '
                      r'prompt;', str(e))
        if not m:
            print(f"davinci-003 could not complete the request: {str(e)}")
            return
        (hard_limit, prompt_tokens) = (int(m.group(1)), int(m.group(2)))
        max_tokens = hard_limit - prompt_tokens
        if max_tokens >= 750:
            print(f"Context length exceeded! Reducing the completion tokens to {max_tokens}...")
            query_model(query, cb, max_tokens)
        else:
            print("Unfortunately, this function is too big to be analyzed with the model's current API limits.")
        
    except openai.OpenAIError as e:
        print(f"davinci-003 could not complete the request: {str(e)}")
    except Exception as e:
        print(f"General exception encountered while running the query: {str(e)}")

# -----------------------------------------------------------------------------

def query_model_async(query, cb):
    """
    Function which sends a query to davinci-003 and calls a callback when the response is available.
    :param query: The request to send to davinci-003
    :param cb: Tu function to which the response will be passed to.
    """
    print("Request to davinci-003 sent...")
    t = threading.Thread(target=query_model, args=[query, cb])
    t.start()

# =============================================================================
# Main
# =============================================================================

def PLUGIN_ENTRY():
    if not openai.api_key:
        openai.api_key = os.getenv("OPENAI_API_KEY")
        if not openai.api_key:
            print("Please edit this script to insert your OpenAI API key!")
            raise ValueError("No valid OpenAI API key found")

    return VulChatPlugin()


repository: ZetaTwo/binja-experiments
file: script_flagyard_tables/flagyard_tables.py
description: Experiments, snippets and other things related to Binary Ninja
created_at: 2021-10-05 13:29:15+00:00
pushed_at: 2024-10-22 23:18:02+00:00
url:
https://github.com/ZetaTwo/binja-experiments/blob/ca5ba019f1346604aeaa21aff313f139b519db0f/script_flagyard_tables/flagyard_tables.py

content:
#!/usr/bin/env python3

import binaryninja
from binaryninja import HighLevelILOperation

# Manually extracted addresses and lookup tables
ENTRY_FUNC = 0x401350
DEST = 0x4037f0
LOOKUP = {
    0x00: 0x1e, 0x01: 0x1f, 0x02: 0x20, 0x03: 0x21, 0x04: 0x22,
    0x05: 0x23, 0x06: 0x24, 0x07: 0x25, 0x08: 0x26, 0x09: 0x27,
    0x16: 0x1c, 0x29: 0x1d, 0x2f: 0x1b, 0x31: 0x01, 0x32: 0x02,
    0x33: 0x03, 0x34: 0x04, 0x35: 0x05, 0x36: 0x06, 0x37: 0x07,
    0x38: 0x08, 0x39: 0x09, 0x3a: 0x0a, 0x3b: 0x0b, 0x3c: 0x0c,
    0x3d: 0x0d, 0x3e: 0x0e, 0x3f: 0x0f, 0x40: 0x10, 0x41: 0x11,
    0x42: 0x12, 0x43: 0x13, 0x44: 0x14, 0x45: 0x15, 0x46: 0x16,
    0x47: 0x17, 0x48: 0x18, 0x49: 0x19, 0x4a: 0x1a, 0x4b: 0x28,
    0x4d: 0x29
}
LOOKUP_INV = {v:k for k,v in LOOKUP.items()}


def parse_constructor(bv, addr):
    if addr == DEST:
        return -1
    func = bv.get_function_at(addr)
    table = None
    for ins in func.hlil.instructions:
        if ins.operation == HighLevelILOperation.HLIL_CALL:
            if str(ins.dest) == 'memset':
                continue
            table = parse_constructor(bv, ins.dest.value.value)
        elif ins.operation == HighLevelILOperation.HLIL_ASSIGN:
            val = ins.operands[1].value.value
            table = val
    assert table != None, f'{addr:x}'
    return table

def find_path(bv, entry):
    queue = [(parse_constructor(bv, ENTRY_FUNC), [])]
    tables = {}
    seen = set()
    while len(queue) > 0:
        (table, path), queue = queue[0], queue[1:]
        if table in seen or table == 0 or table == -1:
            continue
        seen.add(table)
        tables[table] = []
        for i in range(0x2A):
            func = bv.read_pointer(table + 8*i)
            child = parse_constructor(bv, func)
            tables[table].append(child)
            if child == -1:
                return path
            else:
                queue.append((child, path + [i]))
        
    return None

def plugin_entry(bv):
    flag_path = find_path(bv, ENTRY_FUNC)
    print(bytes(LOOKUP_INV[x]+0x30 for x in flag_path).decode())


binaryninja.PluginCommand.register("Flagyard Tables", "", plugin_entry)
# FlagY{vt4bl3s_and_vtabl3s_and_m0re_vt3bles}


repository: techbliss/Python_editor
file: 7.2/plugins/Python_editor.py
description: Better CodeEditor for Ida Pro.
wanted_name: Python Editor
comment: Run me
created_at: 2015-07-04 21:14:02+00:00
pushed_at: 2020-12-25 13:54:19+00:00
url:
https://github.com/techbliss/Python_editor/blob/4c2b67bd0b1a2b9b9403f5ae31cf8f13c247c981/7.2/plugins/Python_editor.py

content:
# Created by: Storm Shadow http://www.techbliss.org
import os, sys
import ida_idaapi, ida_kernwin, ida_idaapi, ida_loader
import idc
from idc import *
from ida_diskio import *
from ida_idaapi import *
from ida_loader import *
import sys
sys.path.insert(0 , ida_diskio.idadir("plugins\\Code editor\\icons"))
import ico
from ico import *

PLUGIN_VERSION = "1.6"
IDAVERISONS    = "IDA PRO 7.2+"
AUTHORS        = "Storm Shadow"
DATE           = "2019"
TWITTER        = "Twitter @zadow28"

def banner():
    banner_options = (PLUGIN_VERSION, AUTHORS, DATE, TWITTER, IDAVERISONS)
    banner_titles = "Python Editor v%s - (c) %s - %s - %s - %s" % banner_options

# print plugin banner
    print("---[" + banner_titles + "]---\n")

banner()

# 1) Create the handler class
class MyEditorHandler(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    # Run editor when invoked.
    def activate(self, ctx):
        g = globals()
        idahome = ida_diskio.idadir("plugins\\Code editor")
        tupac = str(idahome + "\\pyeditor.py")
        ida_kernwin.load_custom_icon(":/ico/python.png")
        load_and_run_plugin(tupac, True)
    def update(self, ctx):
        return ida_kernwin.AST_ENABLE_ALWAYS

class ripeye(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_FIX
    comment = "Run me"
    help = "Python Editor"
    wanted_name = "Python Editor"
    wanted_hotkey = "" #the tooltip horkey goes away when setting it here DONT DO it! and only is shown in File/Plugins menu


    def editor_menuaction(self):
        action_desc = ida_kernwin.action_desc_t(
            'my:editoraction',  # The action name. This acts like an ID and must be unique
            'Python Editor!',  # The action text.
            MyEditorHandler(),  # The action handler.
            'Ctrl+H',  # Optional: the action shortcut DO IT  HERE!
            'Script editor',  # Optional: the action tooltip (available in menus/toolbar)
            ida_kernwin.load_custom_icon(":/ico/python.png")  # hackish load action icon , if no custom icon use number from 1-150 from internal ida
        )

        # 3) Register the action
        ida_kernwin.register_action(action_desc)

        ida_kernwin.attach_action_to_menu(
            'Edit/Editor...',  # The relative path of where to add the action
            'my:editoraction',  # The action ID (see above)
            ida_kernwin.SETMENU_APP)  # We want to append the action after the 'Manual instruction...

        form = ida_kernwin.get_current_tform()
        ida_kernwin.attach_action_to_popup(form, None, "my:editoraction", None)

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        #self._icon_id_file = ida_diskio.BADADDR
        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            form = ida_kernwin.get_current_tform()
            pass

        return PLUGIN_KEEP


    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self.editor_menuaction()
        self._init()

    def term(self):
        pass

    def run(self, arg = 0):
        #we need the calls again if we wanna load it via File/Plugins/editor
        ida_kernwin.msg("Python Editor Loaded to menu \n use Alt+E hot key to quick load ")
        hackish = MyEditorHandler()
        hackish.activate(self)

def PLUGIN_ENTRY():
    return ripeye()


repository: marlkiller/dylib_dobby_hook
file: script/apps/IDA/plugins/patching.py
description: A macOS dylib project , aimed at enhancing and extending the functionality of target software.
wanted_name: Patching
comment: A plugin to enable binary patching in IDA
created_at: 2024-01-14 11:22:23+00:00
pushed_at: 2025-05-11 04:39:38+00:00
url:
https://github.com/marlkiller/dylib_dobby_hook/blob/688f5f4fe30ce1a1162d910e94d9ef9e11be99cb/script/apps/IDA/plugins/patching.py

content:
#------------------------------------------------------------------------------
# Plugin Preflight
#------------------------------------------------------------------------------
#
#    the purpose of this 'preflight' is to test if the plugin is compatible
#    with the environment it is being loaded in. specifically, these preflight
#    checks are designed to be compatible with IDA 7.0+ and Python 2/3
#
#    if the environment does not meet the specifications required by the
#    plugin, this file will gracefully decline to load the plugin without
#    throwing noisy errors (besides a simple print to the IDA console)
#
#    this makes it easy to install the plugin on machines with numerous
#    versions of IDA / Python / virtualenvs which employ a shared plugin
#    directory such as the 'preferred' IDAUSR plugin directory...
#

import sys

# this plugin requires Python 3
SUPPORTED_PYTHON = sys.version_info[0] == 3

# this plugin requires IDA 7.6 or newer
try:
    import ida_pro
    import ida_idaapi
    IDA_GLOBAL_SCOPE = sys.modules['__main__']
    SUPPORTED_IDA = ida_pro.IDA_SDK_VERSION >= 760
except:
    SUPPORTED_IDA = False

# is this deemed to be a compatible environment for the plugin to load?
SUPPORTED_ENVIRONMENT = bool(SUPPORTED_IDA and SUPPORTED_PYTHON)
if not SUPPORTED_ENVIRONMENT:
    print("Patching plugin is not compatible with this IDA/Python version")

#------------------------------------------------------------------------------
# IDA Plugin Stub
#------------------------------------------------------------------------------

if SUPPORTED_ENVIRONMENT:
    import patching
    from patching.util.python import reload_package

def PLUGIN_ENTRY():
    """
    Required plugin entry point for IDAPython plugins.
    """
    return PatchingPlugin()

class PatchingPlugin(ida_idaapi.plugin_t):
    """
    The IDA Patching plugin stub.
    """

    #
    # Plugin flags:
    # - PLUGIN_PROC: Load / unload this plugin when an IDB opens / closes
    # - PLUGIN_HIDE: Hide this plugin from the IDA plugin menu
    # - PLUGIN_UNL:  Unload the plugin after calling run()
    #

    flags = ida_idaapi.PLUGIN_PROC | ida_idaapi.PLUGIN_HIDE | ida_idaapi.PLUGIN_UNL
    comment = "A plugin to enable binary patching in IDA"
    help = ""
    wanted_name = "Patching"
    wanted_hotkey = ""

    def __init__(self):
        self.__updated = getattr(IDA_GLOBAL_SCOPE, 'RESTART_REQUIRED', False)

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

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        if not SUPPORTED_ENVIRONMENT or self.__updated:
            return ida_idaapi.PLUGIN_SKIP

        # load the plugin core
        self.core = patching.PatchingCore(defer_load=True)

        # inject a reference to the plugin context into the IDA console scope
        IDA_GLOBAL_SCOPE.patching = self

        # mark the plugin as loaded
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        pass

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """
        try:
            self.core.unload()
        except Exception as e:
            pass
        self.core = None

    #--------------------------------------------------------------------------
    # Development Helpers
    #--------------------------------------------------------------------------

    def reload(self):
        """
        Hot-reload the plugin.
        """
        if self.core:
            self.core.unload()
        reload_package(patching)
        self.core = patching.PatchingCore()


repository: jk45054/CTF-writeups
file: Flare-On 10/12_HVM/code/do_io_patching.py
description: Writeups for CTF challenges
created_at: 2021-02-03 20:17:28+00:00
pushed_at: 2023-11-12 20:38:55+00:00
url:
https://github.com/jk45054/CTF-writeups/blob/30eb50029d959a56908aae68f05b58afde0c03f7/Flare-On%2010/12_HVM/code/do_io_patching.py

content:
# Flare-On 10, Challenge 12, HVM
#
# IDA Pro plugin for code decryption
# Set cursor (ScreenEA) to an IN/OUT instruction
# - IN: decrypt the code based on R8/R9 values
#   - r8: RC4 key
#   - r9: size
# - OUT: NOP it away (code re-encryption after execution)
#

import idaapi
from ida_ua import insn_t, decode_insn, decode_prev_insn, o_imm, o_reg, create_insn
from ida_kernwin import get_screen_ea
from ida_bytes import get_bytes, patch_bytes, set_cmt, del_items

PLUGIN_COMMENT = "42"
PLUGIN_HELP = "42"
PLUGIN_NAME = "FlareOn10-Ch12-Patching"
PLUGIN_WANTED_HOTKEY = "Alt-9"


class RC4:
    S = None
    i = 0
    j = 0
    dropped = 0

    # constructor for instance of class RC4
    # if n > 0, drop first n key bytes - alias RC4-drop[n]
    def __init__(self, key: bytes, n: int = 0):
        self.i = 0
        self.j = 0
        self.S = bytearray(256)
        self.KSA(key)
        while n > 0:
            self.PRGA()
            n = n - 1

    # Key Scheduling Algorithm (KSA)
    def KSA(self, key: bytes):
        for i in range(256):
            self.S[i] = i
        j = 0
        for i in range(256):
            j = (j + self.S[i] + key[i % len(key)]) % 256
            self.S[i], self.S[j] = self.S[j], self.S[i]

    # Pseudo-random generation algorithm (PRGA)
    def PRGA(self):
        self.i = (self.i + 1) % 256
        self.j = (self.j + self.S[self.i]) % 256
        self.S[self.i], self.S[self.j] = self.S[self.j], self.S[self.i]
        return self.S[(self.S[self.i] + self.S[self.j]) % 256]

    # Encrypt/decrypt bytes
    def crypt(self, data: bytes) -> bytes:
        result = bytearray()
        for b in data:
            result.append(b ^ self.PRGA())
        return bytes(result)


class do_io_patching_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = PLUGIN_COMMENT
    help = PLUGIN_HELP
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_WANTED_HOTKEY

    def patch_in(self, ea):
        # Patch for case opcode: in
        # for IN opcode -> hvm.exe applies XOR decrypt with key r8, fixed key_len 8, buf at VA_alloc + rip + 2 and buf_len r9
        # Decode the two instructions before the IN opcode to get r8 and r9 immed values
        #
        # Decode first instruction ahead of in = prev1_insn
        # This should yield r9d value
        print(f"[*] Found in instruction @ {hex(ea)}")
        prev1_insn = insn_t()
        len_prev1_insn = decode_prev_insn(prev1_insn, ea)
        print(f"[*] Decoding prev1 insn")
        if len_prev1_insn == 0:
            print(f"[-] Couldn't decode prev1 insn @ {hex(prev1_insn.ea)}")
            return ()
        if prev1_insn.get_canon_mnem() != "mov":
            print(
                f"[-] Prev1 instruction is a {prev1_insn.get_canon_mnem()}, expected mov r9d, immed"
            )
            return ()
        if prev1_insn.ops[0].type != o_reg:
            print(
                f"[-] Prev1 mov instruction operand 0 is of type {prev1_insn.ops[1].type}, expected o_reg"
            )
            return ()
        if prev1_insn.ops[1].type != o_imm:
            print(
                f"[-] Prev1 mov instruction operand 1 is of type {prev1_insn.ops[1].type}, expected o_imm"
            )
            return ()
        print(
            f"[*] Found prev1 insn: mov ops[0] reg number {prev1_insn.ops[0].reg}, immed value {hex(prev1_insn.ops[1].value64)}"
        )
        # Decode second instruction ahead of in = prev2_insn
        # This should yield r8 value
        prev2_insn = insn_t()
        len_prev2_insn = decode_prev_insn(prev2_insn, prev1_insn.ea)
        print(f"[*] Decoding prev2 insn @ {hex(prev2_insn.ea)}")
        if len_prev2_insn == 0:
            print(f"[-] Couldn't decode prev2 insn")
            return ()
        if prev2_insn.ops[0].type != o_reg:
            print(
                f"[-] Prev2 mov instruction operand 0 is of type {prev2_insn.ops[1].type}, expected o_reg"
            )
            return ()
        if prev2_insn.get_canon_mnem() != "mov":
            print(
                f"[-] Prev2 instruction is a {prev2_insn.get_canon_mnem()}, expected mov r8, immed"
            )
            return ()
        if prev2_insn.ops[1].type != o_imm:
            print(
                f"[-] Prev2 mov instruction operand 1 is of type {prev2_insn.ops[1].type}, expected o_imm"
            )
            return ()
        print(
            f"[*] Found prev2 insn: mov ops[0] reg number {prev2_insn.ops[0].reg}, immed value {hex(prev2_insn.ops[1].value64)}"
        )
        # Grab the bytes to be decrypted
        # - start: ea + 2 (rip + 2)
        # - len: r9d
        r9d = prev1_insn.ops[1].value64
        r8 = prev2_insn.ops[1].value64
        enc_code_bytes = get_bytes(ea + 2, r9d)
        print(f"[*] Encrypted bytes beginning at {hex(ea+2)}: {enc_code_bytes.hex()}")
        # Initialize RC4 cipher with prev2 immed value as RC4 key
        cipher = RC4(int.to_bytes(r8, 8, "little"))
        dec_code_bytes = cipher.crypt(enc_code_bytes)
        print(f"[*] Decrypted bytes: {dec_code_bytes.hex()} using key: {hex(r8)}")
        # Del items in the range -> make them unexplored
        del_items(ea, r9d)
        # Patch the bytes and NOP the IN
        patch_bytes(ea, b"\x90\x90")
        patch_bytes(ea + 2, dec_code_bytes)
        # Create instruction at patched address
        create_insn(ea)
        # Set comment
        self.in_counter += 1
        # set_cmt(ea, f"Counter: {self.counter}, decrypted @ {hex(ea+2)}, len {hex(r9d)}, before: {enc_code_bytes.hex()}, after: {dec_code_bytes.hex()}", 0)
        set_cmt(
            ea,
            f"IN-Counter: {self.in_counter}, decrypted @ {hex(ea+2)}, len {hex(r9d)}",
            0,
        )
        return ()

    def patch_out(self, ea):
        # Patch for case opcode: out
        # for OUT opcode -> applies XOR key r8, fixed key_len 8 at VA_alloc + rip - 16 - r9, for len r9
        # Decode the two instructions before the IN opcode to get r8 and r9 immed values
        print(f"[*] Found out instruction @ {hex(ea)}")
        # Del the out
        del_items(ea, 2)
        # Patch it with NOP
        patch_bytes(ea, b"\x90\x90")
        # Create instruction at patched address
        create_insn(ea)
        # Set comment
        self.out_counter += 1
        set_cmt(ea, f"OUT-Counter: {self.out_counter}, patched with NOPs", 0)
        return ()

    def init(self):
        try:
            idaapi.msg("do_io_patching init() called!\n")
            self.in_counter = 0
            self.out_counter = 0
            return idaapi.PLUGIN_OK
        except Exception as err:
            idaapi.msg("Exception during init: %s\n" % str(err))
        return idaapi.PLUGIN_SKIP

    def run(self, arg):
        try:
            # idaapi.msg("do_io_patching run() called with %d!\n" % arg)
            # Get current cursor location VA
            ea = get_screen_ea()
            # Decode instruction at ea
            insn = insn_t()
            len_inst = decode_insn(insn, ea)
            if len_inst == 0:
                print(f"[-] Couldn't decode insn @ {hex(ea)}")
                return ()
            print(
                f"do_io_patching: decoded insn @ {hex(ea)}, canon mnem = {insn.get_canon_mnem()}"
            )
            if insn.get_canon_mnem() == "in":
                self.patch_in(ea)
            elif insn.get_canon_mnem() == "out":
                self.patch_out(ea)
            else:
                print(f"[-] Neither in nor out opcode @ {hex(ea)}")
                return ()
        except Exception as err:
            idaapi.msg("Exception during run: %s\n" % str(err))
            raise
        idaapi.msg("do_io_patching run() complete!\n")

    def term(self):
        idaapi.msg("do_io_patching term() called!\n")


def PLUGIN_ENTRY():
    return do_io_patching_plugin_t()


repository: anic/ida2pwntools
file: ida2pwntools.py
description: a IDA 7.0 plugins that helps to attach process created by pwntools and debug pwn
url: https://github.com/anic/ida2pwntools/blob/5a64c269a451808538f2767bafe99d31e2d9c3eb/ida2pwntools.py
created_at: 2018-12-14 05:07:29+00:00
pushed_at: 2023-03-29 14:16:37+00:00
content:
# -*- coding:utf-8 -*-

import idaapi
import ida_nalt
import ida_idd
import ida_dbg
import ida_kernwin

from PyQt5 import QtGui, QtCore, QtWidgets
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *

PLUGNAME = "Ida2pwntools"
WAIT_TIME_NO_UI = 60 # seconds


class MainCallable(object):
	def __init__(self, handler):
		self.handler = handler
	def __call__(self):
		return self.handler()


class WaitThread(QThread):
	sinOut = pyqtSignal(str)  
	sinOutEnd = pyqtSignal(str)
  
	def __init__(self, parent=None):
		super(WaitThread,self).__init__(parent)
		self.filename = ida_nalt.get_root_filename()
		self.target_pid = -1

	def run(self):  
		self.prepare_debug()
		self.sinOutEnd.emit("start debug (PID: %d)" % self.target_pid)
		
	def prepare_debug(self):
		def get_processes_list():
			self.pis = ida_idd.procinfo_vec_t()
			ida_dbg.get_processes(self.pis)
			return 1

		c = 0
		found_pid = False
		while not found_pid:
			c = (c+1) % 5
			self.sinOut.emit("finding process: [%s]" % self.filename+('.' * c).ljust(5, " "))

			ida_kernwin.execute_sync(MainCallable(get_processes_list), \
				ida_kernwin.MFF_FAST)

			for proc in self.pis:
				proc_name = proc.name.split(" ")[1]
				idx = proc_name.rfind("/")

				if idx != -1:
					proc_name = proc_name[idx+1:]
				
				if self.filename == proc_name:
					self.target_pid = proc.pid
					found_pid = True
					break

			if not found_pid:
				self.sleep(1)


class WaitDialog(QDialog):
	def __init__(self):
		super(WaitDialog,self).__init__()
		self.initUI()
		self.thread = WaitThread()
		self.thread.sinOut.connect(self.outText)
		self.thread.sinOutEnd.connect(self.start_debug)
		
	def outText(self, text):
		self.l2.setText(text)
		
	def close_debug(self):
		self.thread.terminate()
		self.hide()

	def get_target_pid(self):
		return self.thread.target_pid
	
	def start_debug(self, text):
		idaapi.msg("[%s] %s\n" % (PLUGNAME, text))
		self.hide()

	def initUI(self):
		self.setWindowTitle(PLUGNAME)
		self.setGeometry(400,400,300,260)
		self.setFixedWidth(320)
		self.setFixedHeight(240)
		vbox = QVBoxLayout()
		
		self.l2 = QLabel()
		self.l2.setAlignment(Qt.AlignCenter)
		vbox.addWidget(self.l2)
		
		self.closeButton = QPushButton()
		self.closeButton.setText("Stop Waiting")
		self.closeButton.setShortcut('Ctrl+D')
		self.closeButton.clicked.connect(self.close_debug)
		self.closeButton.setToolTip("Stop Waiting")
		self.closeButton.resize(200,100)
		vbox.addWidget(self.closeButton)
		
		self.setLayout(vbox)


class timer_debug_noui_t(object):
	def __init__(self):
		self.interval = 1000 # 1 second
		self.obj = idaapi.register_timer(self.interval, self)
		if self.obj is None:
			raise RuntimeError("Failed to register timer")
		self.times = WAIT_TIME_NO_UI

	def __call__(self):
		target_pid = -1

		if idaapi.is_debugger_on():
			idaapi.msg("[%s] the debugger is currently running\n" % PLUGNAME)
			return -1

		if not self.times%5:
			idaapi.msg("[%s] waiting for the process (%ds left)...\n" % \
				(PLUGNAME, self.times))

		filename = ida_nalt.get_root_filename()
		pis = ida_idd.procinfo_vec_t()
		ida_dbg.get_processes(pis)

		for proc in pis:
			if proc.pid==0   and "attach" in proc.name:
				target_pid=0
				break
			proc_name = proc.name.split(" ")[1]
			idx = proc_name.rfind("/")

			if idx != -1:
				proc_name = proc_name[idx+1:]

			if filename == proc_name:
				target_pid = proc.pid
				break

		if target_pid != -1:
			idaapi.msg("[%s] found. start debug (PID: %d)\n" % (PLUGNAME, target_pid))
			ida_dbg.attach_process(target_pid, -1)
			ida_dbg.wait_for_next_event(ida_dbg.WFNE_SUSP, -1)
			ida_dbg.continue_process()
			return -1

		self.times -= 1
		return -1 if self.times == 0 else self.interval

	def __del__(self):
		print("[%s] Timer disposed" % PLUGNAME)


class IDA_Pwntools_Plugin_t(idaapi.plugin_t):
	comment = ""
	help = "help"
	wanted_name = PLUGNAME
	wanted_hotkey = "f12"
	flags = idaapi.PLUGIN_KEEP
	
	def prepare_debug_ui(self):
		if idaapi.is_debugger_on():
			idaapi.warning("[%s] the debugger is currently running" % PLUGNAME)
			return

		wd = WaitDialog()
		idaapi.msg("[%s] waiting...\n" % (PLUGNAME))
		wd.thread.start()
		wd.exec_()

		target_pid = wd.get_target_pid()
		if target_pid != -1:
			ida_dbg.attach_process(target_pid,-1)
			ida_dbg.wait_for_next_event(ida_dbg.WFNE_SUSP, -1)
			ida_dbg.continue_process()
		else:
			idaapi.msg("[%s] exit waiting\n" % (PLUGNAME))

	def init(self):
		menu_bar = next(i for i in QtWidgets.qApp.allWidgets() if isinstance(i, QtWidgets.QMenuBar))
		self.menu = menu_bar.addMenu(PLUGNAME)
		self.menu.addAction("Connect to pwntools").triggered.connect(self.prepare_debug_ui)
		return idaapi.PLUGIN_KEEP

	def term(self):
		idaapi.msg("[%s] terminated\n" % (PLUGNAME))
		self.menu.deleteLater()
		
	def run(self, arg):
		timer_debug_noui_t()


def PLUGIN_ENTRY():
	return IDA_Pwntools_Plugin_t()


repository: Reodus/CBS
file: custom_breakpoint_setter.py
description: IDA Plugin to set custom breakpoints on instructions
wanted_name: Custom Breakpoint Setter
comment: Set breakpoints on specific instructions with enable/disable options
created_at: 2024-06-15 21:37:54+00:00
pushed_at: 2024-08-15 12:42:26+00:00
url:
https://github.com/Reodus/CBS/blob/a1237c164f55932032b1876f81ea97fe5b005f44/custom_breakpoint_setter.py

content:
import idaapi, idautils, idc
import re
from PyQt5 import QtCore, QtGui, QtWidgets

default_instructions = [
    'int3', 'lods', 'stos', 'scas', 'ins', 'rdtsc', 'rdtscp',
    'sysenter', 'sysexit', 'syscall', 'rep', 'repne', 'repe',
    'nop', 'ud2', 'cpuid', 'popf', 'xchg', 'cmpxchg', 'xadd',
    'xor', 'rol', 'ror', 'shl', 'shr', 'neg', 'sbb', 'cmp',
]

class CustomBreakPointPlugin(idaapi.PluginForm):
    def OnCreate(self, form):
        # Get parent widget
        self.parent = self.FormToPyQtWidget(form)
        self.PopulateForm()

    def PopulateForm(self):
        # Create layout
        layout = QtWidgets.QVBoxLayout()

        # Create a table widget
        self.table = QtWidgets.QTableWidget(len(default_instructions), 1)
        self.table.setHorizontalHeaderLabels(["Instructions"])
        self.table.horizontalHeader().setStretchLastSection(True)  # Stretch the instruction column to maximize width
        
        for index, instruction in enumerate(default_instructions):
            item = QtWidgets.QTableWidgetItem(instruction)
            self.table.setItem(index, 0, item)

        layout.addWidget(self.table)

        # Create input field for new instruction
        input_layout = QtWidgets.QHBoxLayout()
        self.instruction_input = QtWidgets.QLineEdit()
        self.instruction_input.setPlaceholderText("Enter new instruction (Regex supported)")
        self.instruction_input.returnPressed.connect(self.add_instruction)  # Add instruction on Enter key press

        # Create combobox for enable/disable/remove options
        self.status_combobox = QtWidgets.QComboBox()
        self.status_combobox.addItems(["Disable", "Enable", "Remove Breakpoints"])

        # Create button to add new instruction
        self.add_instruction_btn = QtWidgets.QPushButton("Add Instruction")
        self.add_instruction_btn.clicked.connect(self.add_instruction)

        layout.addLayout(input_layout)

        # Create button to delete selected instruction
        self.delete_instruction_btn = QtWidgets.QPushButton("Delete Selected")
        self.delete_instruction_btn.clicked.connect(self.delete_instruction)

        # Create set breakpoint button
        self.set_bp_btn = QtWidgets.QPushButton("Execute")
        self.set_bp_btn.clicked.connect(self.set_breakpoints)

        # Add keyboard shortcuts
        self.add_instruction_shortcut = QtWidgets.QShortcut(QtGui.QKeySequence("Return"), self.parent)
        self.add_instruction_shortcut.activated.connect(self.add_instruction)

        self.delete_instruction_shortcut = QtWidgets.QShortcut(QtGui.QKeySequence("Delete"), self.parent)
        self.delete_instruction_shortcut.activated.connect(self.delete_instruction)

        layout.addWidget(self.delete_instruction_btn)
        layout.addWidget(self.set_bp_btn)
        input_layout.addWidget(self.instruction_input)
        input_layout.addWidget(self.status_combobox)
        input_layout.addWidget(self.add_instruction_btn)

        self.parent.setLayout(layout)

    def add_instruction(self):
        new_instruction = self.instruction_input.text().strip()
        if new_instruction:
            row_count = self.table.rowCount()
            self.table.insertRow(row_count)
            item = QtWidgets.QTableWidgetItem(new_instruction)
            self.table.setItem(row_count, 0, item)
            self.instruction_input.clear()

    def delete_instruction(self):
        selected_rows = set()
        for item in self.table.selectedItems():
            selected_rows.add(item.row())

        for row in sorted(selected_rows, reverse=True):
            self.table.removeRow(row)

    def set_breakpoints(self):
        selected_instructions = []
        # Iterate over table rows to find selected instructions
        for row in range(self.table.rowCount()):
            item = self.table.item(row, 0)
            if item is not None:  # Check if item exists
                selected_instructions.append(self.normalize_instruction(item.text()))

        status = self.status_combobox.currentText()

        # Set/enable/disable breakpoints based on selected instructions
        for seg_ea in idautils.Segments():
            for func_ea in idautils.Functions(seg_ea, idc.get_segm_end(seg_ea)):
                func_end = idc.get_func_attr(func_ea, idc.FUNCATTR_END)

                for head in idautils.Heads(func_ea, func_end):
                    if idc.is_code(idc.get_full_flags(head)):
                        full_instruction = self.normalize_instruction(idc.generate_disasm_line(head, 0))  # Get the normalized full instruction

                        for pattern in selected_instructions:
                            if re.search(pattern, full_instruction):  # Use regex search on normalized strings
                                idaapi.add_bpt(head)
                                
                                msg = ""

                                if status == "Disable":
                                    idaapi.disable_bpt(head)
                                    msg = "Disabled"
                                elif status == "Enable":
                                    idaapi.enable_bpt(head)
                                    msg = "Enabled"
                                elif status == "Remove Breakpoints":
                                    idaapi.del_bpt(head)
                                    msg = "Removed"
                                print(f"[CBS] ({msg}) breakpoint -> {hex(head)}: {full_instruction}")

    def normalize_instruction(self, instruction):
        """Normalize instruction string by removing extra spaces, commas, and semicolons."""
        return re.sub(r'\s+', ' ', instruction.replace(',', '').replace(';', '')).strip()

    def OnClose(self, form):
        pass

class MyPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Set breakpoints on specific instructions with enable/disable options"
    help = "This plugin sets breakpoints on specified instructions with enable/disable options"
    wanted_name = "Custom Breakpoint Setter"

    def init(self):
        return idaapi.PLUGIN_OK

    def run(self, arg):
        self.form = CustomBreakPointPlugin()
        self.form.Show("Custom BreakPoint Plugin")
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

def PLUGIN_ENTRY():
    return MyPlugin()


repository: sean2077/hexpaste-ida
file: hexpaste.py
description: IDA plugin for quickly paste hex bytes into memory.
url: https://github.com/sean2077/hexpaste-ida/blob/0a680bc4eafd9fd8dd20fcfe0d123e55a9437b26/hexpaste.py
comment: Paste Hex Bytes
created_at: 2024-06-15 23:50:10+00:00
pushed_at: 2024-06-16 06:06:11+00:00
content:
############################################################################################
##
## One-Click Hex Paste!
##
## Available for IDA 7+ and Python 3.8+
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\<ida version>\plugins
##
## To run:
##      Right-click on an address in the disassembly and select "Hex Paste"
##      The hex encoded bytes from the clipboard will be written to the selected address
##
############################################################################################

__AUTHOR__ = "@sean2077"

PLUGIN_NAME = "Hex Paste"
PLUGIN_HOTKEY = "Ctrl+Shift+V"
VERSION = "1.0.0"

ACTION_PREFIX = "sean2077"

import re

import idaapi
import idc
import PyQt5.QtCore as QtCore
import PyQt5.QtGui as QtGui
import PyQt5.QtWidgets as QtWidgets
from PyQt5.Qt import QApplication


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


def validate_hex_string(hex_string):
    # Remove all whitespace and convert to uppercase
    cleaned = re.sub(r"\s+", "", hex_string).upper()
    # Check if it's a valid hex string
    if re.fullmatch(r"[0-9A-F]*", cleaned):
        return cleaned
    return None


def write_bytes_to_address(address, hex_string):
    bytes_data = bytes.fromhex(hex_string)
    idaapi.patch_bytes(address, bytes_data)


def show_message(title, message):
    msg_box = QtWidgets.QMessageBox()
    msg_box.setIcon(QtWidgets.QMessageBox.Information)
    msg_box.setWindowTitle(title)
    msg_box.setText(message)
    msg_box.setStandardButtons(QtWidgets.QMessageBox.Ok)
    msg_box.exec_()


def confirm_write(hex_string, address):
    msg_box = QtWidgets.QMessageBox()
    msg_box.setIcon(QtWidgets.QMessageBox.Question)
    msg_box.setWindowTitle("Confirm Hex Paste")
    msg_box.setText(f"Are you sure you want to write the following bytes to address {address:X}?\n\n{hex_string}")
    msg_box.setStandardButtons(QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No)
    return msg_box.exec_() == QtWidgets.QMessageBox.Yes


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


class HexPastePlugin(idaapi.plugin_t):
    """
    The IDA Plugin for hex paste.
    """

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Paste Hex Bytes"
    help = "Right-click an address and select 'Paste Hex'"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        self._init_action_paste_bytes()
        self._init_hooks()
        idaapi.msg("%s %s initialized...\n" % (self.wanted_name, VERSION))
        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.
        """
        self._hooks.unhook()
        self._del_action_paste_bytes()
        idaapi.msg("%s terminated...\n" % self.wanted_name)

    def _init_hooks(self):
        """
        Install plugin hooks into IDA.
        """
        self._hooks = Hooks()
        self._hooks.hook()

    ACTION_PASTE_BYTES = f"{ACTION_PREFIX}:paste_bytes"

    def _init_action_paste_bytes(self):
        """
        Register the paste bytes action with IDA.
        """
        action_desc = idaapi.action_desc_t(
            self.ACTION_PASTE_BYTES,  # The action name.
            "Paste Hex",  # The action text.
            IDACtxEntry(paste_bytes),  # The action handler.
            PLUGIN_HOTKEY,  # Optional: action shortcut
            "Paste hex bytes from clipboard",  # Optional: tooltip
            31,  # Icon
        )
        assert idaapi.register_action(action_desc), "Action registration failed"

    def _del_action_paste_bytes(self):
        """
        Delete the paste action from IDA.
        """
        idaapi.unregister_action(self.ACTION_PASTE_BYTES)


class Hooks(idaapi.UI_Hooks):

    def finish_populating_widget_popup(self, widget, popup):
        """
        A right-click menu is about to be shown. (IDA 7)
        """
        inject_hex_paste_actions(widget, popup, idaapi.get_widget_type(widget))
        return 0


def inject_hex_paste_actions(form, popup, form_type):
    """
    Inject paste actions to popup menu(s) based on context.
    """
    if form_type == idaapi.BWN_DISASMS:
        idaapi.attach_action_to_popup(form, popup, HexPastePlugin.ACTION_PASTE_BYTES, "Paste Hex", idaapi.SETMENU_APP)
    return 0


def paste_bytes():
    """
    Paste bytes from clipboard to selected address
    """
    ea = idc.get_screen_ea()
    if ea == idaapi.BADADDR:
        idaapi.msg("Invalid address selected\n")
        return

    clipboard_data = get_clipboard_data()
    valid_hex = validate_hex_string(clipboard_data)
    if not valid_hex:
        show_message("Invalid Hex String", f"Invalid hex string in clipboard: {clipboard_data}\n")
        idaapi.msg("Invalid hex string in clipboard\n")
        return

    if confirm_write(valid_hex, ea):
        write_bytes_to_address(ea, valid_hex)
        idaapi.msg("Hex bytes pasted to address %x\n" % ea)
    else:
        idaapi.msg("Hex paste cancelled\n")


class IDACtxEntry(idaapi.action_handler_t):
    """
    A basic Context Menu class to utilize IDA's action handlers.
    """

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

    def activate(self, ctx):
        """
        Execute the embedded action_function when this context menu is invoked.
        """
        self.action_function()
        return 1

    def update(self, ctx):
        """
        Ensure the context menu is always available in IDA.
        """
        return idaapi.AST_ENABLE_ALWAYS


repository: Finch4/Malware-Analysis-Reports
file: IAT_Auto_Label.py
description: My Malware Analysis Reports
wanted_name: IAT Auto Label
comment: Label dynamically resolved IAT entries
created_at: 2021-01-03 15:32:05+00:00
pushed_at: 2022-05-24 13:30:18+00:00
url:
https://github.com/Finch4/Malware-Analysis-Reports/blob/4f3baae07575e799db97ec22cb271d89c0fb0879/IAT_Auto_Label.py

content:
# Original Author: AGDCServices - https://github.com/AGDCservices/Ghidra-Scripts - https://www.youtube.com/channel/UCnpn999NpDMMPxZXW8sgZLA - Video Example [Ghidra]: https://youtu.be/wsdPmW0dt0I?t=820
# Converted to IDA: Finch - https://twitter.com/Finch39487976 - Video Example: https://youtu.be/BR65Y2kGxDY
# IAT Dumper: https://github.com/AGDCservices/Misc-Malware-Analysis-Tools/blob/main/Dump_Labeled_Iat_Memory.exe
# Watch the video of AGDCServices to understand how to use it


import idaapi
import ida_kernwin
import ida_ida
import idc
import idautils
import ida_nalt
import random
class myplugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Label dynamically resolved IAT entries"
    help="https://github.com/Finch4/Malware-Analysis-Reports"
    wanted_name = "IAT Auto Label"
    wanted_hotkey = "Alt-F8"
    def init(self):
        idaapi.msg("init() called!\n")
        return idaapi.PLUGIN_OK
    

    def parse_iat(self, file):
                labeledIatList = open(file, "r").read().splitlines()
                labeledIatDict = dict()
                for i in labeledIatList:
                    cur_rva, cur_iat_label = i.split('\t')
                    labeledIatDict[hex(ida_nalt.get_imagebase() + int(cur_rva, 16))] = cur_iat_label
                return labeledIatDict
            
            
    def apply_iat(self, iat):
            min_ea = ida_ida.inf_get_min_ea()
            max_ea = ida_ida.inf_get_max_ea()
            image_base = ida_nalt.get_imagebase()
            idaapi.msg(f"{hex(min_ea)} - {hex(max_ea)} - {hex(image_base)} \n")
            for i in range(min_ea, max_ea):
                mnemonic = idc.print_insn_mnem(i)
                if mnemonic == "call" and str(idc.print_operand(i, 0)).startswith("dword"):
                    try:
                        idc.set_name(idc.get_operand_value(i, 0), iat[hex(idc.get_operand_value(i, 0))]+str(random.randint(0, 1000))) # I added random.randint() because some functions already in the IAT are dynamically resolved again, so renaming them causes a conflict, if you want remove it. (Sample causing the problem: https://bazaar.abuse.ch/sample/49fd52a3f3d1d46dc065217e588d1d29fba4d978cd8fdb2887fd603320540f71/)
                        open("resolved_iat.txt","a").write(f"{hex(idc.get_operand_value(i, 0))}\n")
                    except:
                        open("unresolved_iat.txt","a").write(f"{hex(idc.get_operand_value(i, 0))}\n")


    def run(self, arg):        
            file = ida_kernwin.ask_file(0, "*.txt", "Load IAT")
            idaapi.msg(f"{file} Selected!\n")
            if file == None:
               idaapi.msg("Please select a file!")
            labeled_iat_dump_file_name = self.parse_iat(file)
            self.apply_iat(labeled_iat_dump_file_name)
            
        
        

    def term(self):
        idaapi.msg("term() called!\n")
        
def PLUGIN_ENTRY():
    return myplugin_t()


repository: repnz/ida-plugins
file: reg_xref.py
description: A collection of my IDA plugins
url: https://github.com/repnz/ida-plugins/blob/e63c1e721e961eec56b8fdcbe4c192fb8412e60a/reg_xref.py
wanted_name: RegisterXrefs
comment: Register Xrefs
created_at: 2020-08-01 12:54:03+00:00
pushed_at: 2020-08-07 15:42:22+00:00
content:
"""
Register Cross References by Ori Damari (github.com/repnz)
"""
import sark
import idaapi


def log(msg):
    idaapi.msg("[RegXref] {}\n".format(msg))


log("Plugin is loaded")


class RegisterXrefsPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "Register Xrefs"
    help = "List references to a register inside a function"
    wanted_name = "RegisterXrefs"
    wanted_hotkey = "Shift+Z"

    @staticmethod
    def init():
        return idaapi.PLUGIN_OK

    @staticmethod
    def term():
        pass

    @staticmethod
    def run(_):
        run()


def PLUGIN_ENTRY():
    return RegisterXrefsPlugin()


class RegisterReferencesView(idaapi.Choose):
    def __init__(self, name, reg):
        title = 'RegisterXrefs: Function {}: xrefs to {}'.format(name, reg)
        idaapi.Choose.__init__(self, title, [["Address", 16], ["Instruction", 50]])
        self.items = []

    def add_xref(self, ea):
        text = sark.Line(ea).disasm
        hex_address = hex(ea)

        if hex_address.endswith("L"):
            hex_address = hex_address[:-1]

        self.items.append([hex_address, text])

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

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

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

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

    @staticmethod
    def make_item(reg):
        return [reg, reg]

    def OnClose(self):
        return


def has_register_reference(ins, reg_id):
    for op in ins.operands:
        for reg in op.regs:
            if get_register_identifier(reg) == reg_id:
                return True
    return False


def get_highlight():
    highlight_result = idaapi.get_highlight(idaapi.get_current_viewer())
    return highlight_result[0] if highlight_result and highlight_result[1] else ''


def get_register_identifier(register_name):
    """
    In the IDA API, Every register has an ID.
    The issue is, sometimes the same register can be accessed in different sizes: For example,
    RAX, AX, AL. In IDA, the ID of the 'AX' register is different than the ID of 'AX' - this causes the cross
    reference list to be miss some references. To solve this issue, We list all the registers that have "incorrect" ID
    and translate them to other registers to get the correct ID.
    """
    register_translations = {
        'al': 'ax',
        'cl': 'cx',
        'dl': 'dx',
        'bl': 'bx',
        'ah': 'ax',
        'ch': 'cx',
        'dh': 'dx',
        'bh': 'bx',
        'spl': 'sp',
        'bpl': 'bp',
        'sil': 'si',
        'dil': 'di'
    }

    if register_name in register_translations:
        register_name = register_translations[register_name]

    return sark.get_register_id(register_name)


def run():
    try:
        current_function = sark.Function()
    except sark.exceptions.SarkNoFunction:
        log("Cannot xref registers outside of functions.")
        return

    register_name = get_highlight()

    try:
        register_id = get_register_identifier(register_name)
    except sark.exceptions.SarkInvalidRegisterName:
        log("Highlight a register to xref")
        return

    choose = RegisterReferencesView(current_function.name, register_name)

    for line in current_function.lines:
        if has_register_reference(line.insn, register_id):
            choose.add_xref(line.ea)

    choose.show()


repository: deepin-community/binwalk
file: src/scripts/binida.py
wanted_name: Binwalk IDA Plugin
comment: Scan the current IDB for file signatures
created_at: 2024-09-06 02:17:26+00:00
pushed_at: 2024-09-09 01:38:16+00:00
url:
https://github.com/deepin-community/binwalk/blob/dac903de9d0105ef1dfb4aa650e36a8af7ae27be/src/scripts/binida.py

content:
if idaapi.IDA_SDK_VERSION <= 695:
    import idc
    import idaapi
    import binwalk
if idaapi.IDA_SDK_VERSION >= 700:
    import ida_idc
    import ida_idaapi
    import binwalk

    from idaapi import *
else:
    pass

# use 'try' here for compatibility with older API
# use Actions API for handlers
try:
    class OpHandler(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

        def activate(self, ctx):
            arg = None
            a = binwalk_t()
            a.opcode_scan(arg)
            return 1

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

# use 'try' here for compatibility with older API
# use Actions API for handlers
try:
    class SigHandler(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

        def activate(self, ctx):
            arg = None
            b = binwalk_t()
            b.signature_scan(arg)
            return 1

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

class binwalk_t(idaapi.plugin_t):
    flags = 0
    comment = "Scan the current IDB for file signatures"
    help = ""
    wanted_name = "Binwalk IDA Plugin"
    wanted_hotkey = ""

    def init(self):
        if idaapi.IDA_SDK_VERSION <= 695:
            self.menu_context_1 = idaapi.add_menu_item(
                "Search/", "binwalk opcodes", "", 0, self.opcode_scan, (None,))
            self.menu_context_2 = idaapi.add_menu_item(
                "Search/", "binwalk signatures", "", 0, self.signature_scan, (None,))

        if idaapi.IDA_SDK_VERSION >= 700:
            # populate action menus
            action_desc = idaapi.action_desc_t(
                'my:opaction',  # action name. This acts like an ID and must be unique
                'Binwalk opcodes',  # text for this action
                OpHandler(),  # the action handler
                '',  # optional shortcut key
                'Binwalk opcodes',  # optional action tooltip for menus/toolbar
                )

            # Register the action
            idaapi.register_action(action_desc)
            idaapi.attach_action_to_menu(
                'Search/',
                'my:opaction',
                idaapi.SETMENU_APP)

            # populate action menus
            action_desc = idaapi.action_desc_t(
                'my:sigaction',
                'Binwalk signatures',
                SigHandler(),
                '',
                'Binwalk signatures',
                )

            # Register the action
            idaapi.register_action(action_desc)
            idaapi.attach_action_to_menu(
                'Search/',
                'my:sigaction',
                idaapi.SETMENU_APP)

        else:
            pass

        return idaapi.PLUGIN_KEEP

    def term(self):
        if idaapi.IDA_SDK_VERSION <= 695:
            idaapi.del_menu_item(self.menu_context_1)
            idaapi.del_menu_item(self.menu_context_2)
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.detach_action_from_menu(
                'Search/',
                'my:opaction')
            idaapi.detach_action_from_menu(
                'Search/',
                'my:sigaction')
        else:
            pass

        return None

    def run(self, arg):
        return None

    def signature_scan(self, arg):
        binwalk.scan(idc.GetIdbPath(), signature=True)

    def opcode_scan(self, arg):
        binwalk.scan(idc.GetIdbPath(), opcode=True)


def PLUGIN_ENTRY():
    return binwalk_t()


repository: matthewrhannon/IDAPython-Automation-Scripts
file: FiX_Int3_Functions.py
created_at: 2024-09-24 12:53:57+00:00
pushed_at: 2024-09-24 13:32:30+00:00
url:
https://github.com/matthewrhannon/IDAPython-Automation-Scripts/blob/cad9d3a5bb5a92ac65736b674a7166348a7ebcc6/FiX_Int3_Functions.py

content:
import concurrent.futures

import idc
import ida_range
import ida_ida
import ida_ua

try:
    from ida_idaapi import idaapi
except Exception:
    print("WARNING: NOTICE - from ida_idaapi import idaapi threw exception!")
    import idaapi
    pass
    
import ida_search
import ida_funcs
import ida_bytes
import ida_kernwin
import ida_ida
import ida_nalt
import idautils
import numpy as np

thread_num = -1
func_arrays = []

def intro():

    info = idaapi.get_inf_structure()
    filename = idaapi.get_input_file_path()
    entrypoint = info.start_ip
    imagebase = ida_nalt.get_imagebase()
    is_64bits = info.is_64bit()
    is_dll = info.is_dll()
    proc_name = ida_ida.inf_get_procname()
    
    print("info=" + str(info))
    print("filename=" + str(filename))    
    print("entrypoint=" + hex(entrypoint))
    print("imagebase=" + hex(imagebase))
    print("is_64bits=" + str(is_64bits))
    print("is_dll=" + str(is_dll))
    print("proc_name=" + str(proc_name))
    
    return True
    
def getInstructionBytes(offset):
        ins = idautils.DecodeInstruction(offset)
        ins_bytes = ida_bytes.get_bytes(offset, ins.size)
        return ins_bytes 

def remake_single_funcs():

    FUNCATTR_STARTY = 0x1401183D0
    FUNCATTR_ENDY = 0x14011841A

    #current_ea = 0
    count = 0
    dont_do_again = False
    for function_ea in idautils.Functions():
        dont_do_again = True
        
        if not dont_do_again:
            
            start_ea = idc.get_func_attr(function_ea, FUNCATTR_STARTY) #0x0000000141D99590 
            end_ea = idc.get_func_attr(function_ea, FUNCATTR_ENDY) #0x0000000141D9CD2C 
            
            func_length = end_ea - start_ea
           
            #current_ea = start_ea
            
            print("start_ea=" + hex(start_ea) + " --> end_ea=" + hex(end_ea) + " size=" + str(func_length))
            
            #start_ea = idc.to_ea(start_ea)
            ida_kernwin.jumpto(start_ea)
            
            #Save Function Comments
            function_disassembly_comments_Repeat = 1
            function_disassembly_comments = ''
            function_disassembly_comments = idc.get_func_cmt(start_ea, 1)
            if function_disassembly_comments == '' or str(''):
                function_disassembly_comments = idc.get_func_cmt(start_ea, 0)
                if function_disassembly_comments == '' or str(''):
                    function_disassembly_comments = ''
                    function_disassembly_comments_Repeat = 0
                   
            #Save Disassembly Comments
            disassembly_comments_Repeat = 1
            disassembly_comments = ''
            disassembly_comments = idc.get_cmt(start_ea, 1)
            if disassembly_comments == '' or str(''):
                disassembly_comments = idc.get_cmt(start_ea, 0)
                if disassembly_comments == '' or str(''):
                    disassembly_comments = ''
                    disassembly_comments_Repeat = 0

            #TODO: hexrays_comments
            #TODO: epilogue and 'postlog'
            
            #sanitycheck
            if start_ea % 0x04 != 0:
                print("[SANITY FAILURE] start_ea=" + hex(start_ea)+ " and end_ea=" + hex(end_ea) + "... continuing ... and destroying instructions inbetween...")
                num_bytes_to_un_encode = func_length
                
                ida_funcs.del_func(start_ea)
                ida_bytes.del_items(start_ea, 0, num_bytes_to_un_encode)
                
                continue
                
            #ida_funcs.del_func(function_ea)
            #ida_ua.create_insn(insn, start_ea)

            ida_funcs.del_func(function_ea)
            ida_ua.create_insn(start_ea)
            ida_funcs.add_func(start_ea)
            idc.set_func_cmt(start_ea, function_disassembly_comments, function_disassembly_comments_Repeat)
            idc.set_cmt(start_ea, disassembly_comments, disassembly_comments_Repeat)
            
            if not ida_bytes.has_xref(end_ea):
                if ida_bytes.is_func(end_ea):
                    #ida_kernwin.jumpto(end_ea)
                    ida_bytes.del_items(end_ea, 0, 1)
                    #insn = ida_ua.insn_t()
                    ida_funcs.del_func(end_ea)
            elif ida_bytes.has_xref(end_ea):
                ida_ua.create_insn(end_ea)
                ida_funcs.add_func(end_ea)
            else:
                print("!!! WTF !!! end_ea=" + hex(end_ea))
                
            #ida_funcs.update_func(1)
            #current_ea = idc.next_head(current_ea)
            
            count = count + 1
            if count % 500 == 0:
                print("count % 500 == 0! count=" + str(count)) #+ "/" + str(len(idautils.Functions)))
                #return count
            
        return count
        
def remake_all_funcs():

    #current_ea = 0
    count = 0
    for function_ea in idautils.Functions():
        start_ea = idc.get_func_attr(function_ea, idc.FUNCATTR_START) #0x0000000141D99590 
        end_ea = idc.get_func_attr(function_ea, idc.FUNCATTR_END) #0x0000000141D9CD2C 
        
        func_length = end_ea - start_ea
       
        #current_ea = start_ea
        
        print("start_ea=" + hex(start_ea) + " --> end_ea=" + hex(end_ea) + " size=" + str(func_length))
        
        #start_ea = idc.to_ea(start_ea)
        ida_kernwin.jumpto(start_ea)
        
        #Save Function Comments
        function_disassembly_comments_Repeat = 1
        function_disassembly_comments = ''
        function_disassembly_comments = idc.get_func_cmt(start_ea, 1)
        if function_disassembly_comments == '' or str(''):
            function_disassembly_comments = idc.get_func_cmt(start_ea, 0)
            if function_disassembly_comments == '' or str(''):
                function_disassembly_comments = ''
                function_disassembly_comments_Repeat = 0
               
        #Save Disassembly Comments
        disassembly_comments_Repeat = 1
        disassembly_comments = ''
        disassembly_comments = idc.get_cmt(start_ea, 1)
        if disassembly_comments == '' or str(''):
            disassembly_comments = idc.get_cmt(start_ea, 0)
            if disassembly_comments == '' or str(''):
                disassembly_comments = ''
                disassembly_comments_Repeat = 0
        
        
        #TODO: hexrays_comments
        #TODO: epilogue and 'postlog'
        
        #sanitycheck
        if start_ea % 0x04 != 0:
            print("[SANITY FAILURE] start_ea=" + hex(start_ea)+ " and end_ea=" + hex(end_ea) + "... continuing ... and destroying instructions inbetween...")
            num_bytes_to_un_encode = func_length
            
            ida_funcs.del_func(start_ea)
            ida_bytes.del_items(start_ea, 0, num_bytes_to_un_encode)
            
            continue
            
        #ida_funcs.del_func(function_ea)
        #ida_ua.create_insn(insn, start_ea)
        

        ida_funcs.del_func(function_ea)
        ida_ua.create_insn(start_ea)
        ida_funcs.add_func(start_ea)
        idc.set_func_cmt(start_ea, function_disassembly_comments, function_disassembly_comments_Repeat)
        idc.set_cmt(start_ea, disassembly_comments, disassembly_comments_Repeat)
        
        if not ida_bytes.has_xref(end_ea):
            if ida_bytes.is_func(end_ea):
                #ida_kernwin.jumpto(end_ea)
                ida_bytes.del_items(end_ea, 0, 1)
                #insn = ida_ua.insn_t()
                ida_funcs.del_func(end_ea)
        elif ida_bytes.has_xref(end_ea):
            ida_ua.create_insn(end_ea)
            ida_funcs.add_func(end_ea)
        else:
            print("!!! WTF !!! end_ea=" + hex(end_ea))
            
        #ida_funcs.update_func(1)
        #current_ea = idc.next_head(current_ea)
        
        
        count = count + 1
        if count % 500 == 0:
            print("count % 500 == 0! count=" + str(count)) #+ "/" + str(len(idautils.Functions)))
            #return count
        
    return count

countera = 0
def remake_some_funcs():

    global thread_num
    thread_num = thread_num + 1
    print("thread_num=" + str(thread_num))
    
    #current_ea = 0
    count = 0
    #for function_ea in idautils.Functions():
    for function_ea in func_arrays[thread_num]:
        start_ea = idc.get_func_attr(function_ea, idc.FUNCATTR_START) #0x0000000141D99590 
        end_ea = idc.get_func_attr(function_ea, idc.FUNCATTR_END) #0x0000000141D9CD2C 
        
        func_length = end_ea - start_ea
       
        #current_ea = start_ea
        
        print("start_ea=" + hex(start_ea) + " --> end_ea=" + hex(end_ea) + " size=" + str(func_length))
        
        #start_ea = idc.to_ea(start_ea)
        ida_kernwin.jumpto(start_ea)
        '''
        #Save Function Comments
        function_disassembly_comments_Repeat = 1
        function_disassembly_comments = ''
        function_disassembly_comments = idc.get_func_cmt(start_ea, 1)
        if function_disassembly_comments == '' or str(''):
            function_disassembly_comments = idc.get_func_cmt(start_ea, 0)
            if function_disassembly_comments == '' or str(''):
                function_disassembly_comments = ''
                function_disassembly_comments_Repeat = 0
               
        #Save Disassembly Comments
        disassembly_comments_Repeat = 1
        disassembly_comments = ''
        disassembly_comments = idc.get_cmt(start_ea, 1)
        if disassembly_comments == '' or str(''):
            disassembly_comments = idc.get_cmt(start_ea, 0)
            if disassembly_comments == '' or str(''):
                disassembly_comments = ''
                disassembly_comments_Repeat = 0
        '''
        
        #TODO: hexrays_comments
        #TODO: epilogue and 'postlog'
        
        #sanitycheck
        if start_ea % 0x04 != 0:
            print("[SANITY FAILURE] start_ea=" + hex(start_ea)+ " and end_ea=" + hex(end_ea) + "... continuing ... and destroying instructions inbetween...")
            num_bytes_to_un_encode = func_length
            
            ida_funcs.del_func(start_ea)
            ida_bytes.del_items(start_ea, 0, num_bytes_to_un_encode)
            
            count = count + 1        
            if count % 500 == 0:
                print("count % 500 == 0! count=" + str(count)) #+ "/" + str(len(idautils.Functions)))
                
            continue
            
        else:
            continue
        #ida_funcs.del_func(function_ea)
        #ida_ua.create_insn(insn, start_ea)
'''
        ida_funcs.del_func(function_ea)
        ida_ua.create_insn(start_ea)
        ida_funcs.add_func(start_ea)
        idc.set_func_cmt(start_ea, function_disassembly_comments, function_disassembly_comments_Repeat)
        idc.set_cmt(start_ea, disassembly_comments, disassembly_comments_Repeat)
        
        if not ida_bytes.has_xref(end_ea):
            if ida_bytes.is_func(end_ea):
                #ida_kernwin.jumpto(end_ea)
                ida_bytes.del_items(end_ea, 0, 1)
                #insn = ida_ua.insn_t()
                ida_funcs.del_func(end_ea)
        elif ida_bytes.has_xref(end_ea):
            ida_ua.create_insn(end_ea)
            ida_funcs.add_func(end_ea)
        else:
            print("!!! WTF !!! end_ea=" + hex(end_ea))
            
        #ida_funcs.update_func(1)
        #current_ea = idc.next_head(current_ea)
'''

# Yield successive n-sized
# chunks from l.
def divide_chunks(l, n):
    
    # looping till length l
    global countera
    for i in range(0, countera, n): 
        yield l[i:i + n]


def PLUGIN_ENTRY():
    intro()


    counter1 = 0
    
    funcs = idautils.Functions()
    for func_count in funcs:
        counter1 = counter1 + 1
    num_functions = counter1 #int(len(funcs), 10)
    print("num_functions=" + str(num_functions))
    
    num_threads = 8
    #thread_size = int(num_functions / num_threads, 10)
    thread_size = num_functions // num_threads
    print("thread_size=" + str(thread_size))
    n = thread_size
    #thread_size_remainder = num_functions % num_threads
    func_arrays = list(divide_chunks(funcs, thread_size))
    #func_arrays = [funcs[i * n:(i + 1) * n] for i in range(int((counter1 + n - 1) / n ))]
    #for i in np.arange(c/10):
    
    #print(str(func_arrays))
    
    print("Main thread continuing to run...")
    
    def worker():
        remake_some_funcs()
     
    pool = concurrent.futures.ThreadPoolExecutor(max_workers=num_threads)
    
    counter = 0
    print("Launching...")
    while counter < num_threads:
        pool.submit(worker)
        counter = counter + 1
        
    pool.shutdown(wait=True)
     
    print("Done!")
    
    
    
'''    
 def PLUGIN_ENTRY():
    intro()   
    print("Number functions remade = " + str(remake_all_funcs()))
'''


repository: Accenture/protobuf-finder
file: protobuf-finder.py
description: IDA Pro plugin for reconstructing original .proto files from binary.
wanted_name: Protobuf
comment: Protobuf string finder.
created_at: 2022-06-24 12:16:32+00:00
pushed_at: 2024-08-19 10:36:13+00:00
url:
https://github.com/Accenture/protobuf-finder/blob/3431fd4f9554000787ecf2662f1c4cc4a77f2141/protobuf-finder.py

content:
"""
Protobuf Finder
 > IDA plugin for reconstructing original .proto files from binary
"""

import re

# require google's protobuf library
from google.protobuf.descriptor_pb2 import FileDescriptorProto, FieldDescriptorProto
from google.protobuf.message import DecodeError

# IDA API imports
import idaapi
import ida_bytes
import ida_ida
import ida_kernwin


# Max protobuf size - used to load code from IDA database; 
MAX_PROTOBUF_SIZE = 0x100000

# Icon
icon = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x000\x00\x00\x00%\x08\x06\x00\x00\x00\x04\x19j\xaf\x00\x00\x00\x01sRGB\x00\xae\xce\x1c\xe9\x00\x00\x00\x04gAMA\x00\x00\xb1\x8f\x0b\xfca\x05\x00\x00\x00\tpHYs\x00\x00\x0e\xc3\x00\x00\x0e\xc3\x01\xc7o\xa8d\x00\x00\x06\x0cIDATXG\xed\x97{PTu\x14\xc7\xcf\xdd]vY\x9e\xbb\x08\x08\xb8\x12h>\x92t@\x13\x07\xb2\x04\x93\xac\x1c\tL\x19\x9f\x199\x16\x16\x1a\x98\xa3i9\x84J\xa33\x068\xcd\xa0\xd9\x03m0Lj4\x07\xa3\xc9\x07\xa4#\xf8\xd8\x89-5\x1b\x9f\xa4;$\x12\xb2\xbca\x1f\xf7\xf6\xfb\xdd=\xbb\xee\xb2w/\xeb\x8c\x7f\xf4\xc7~\x98\xe3\xfd}\xcf]\x99s\xf8\x9d\xf3\xfb\x9d\x05/^\xbcx\xd1\x10+&v\x97\x18\xe7`\xf5\xc4\x0e\x11[@\xec\x7f\t\r\x9c\x06\xe8\x18\xb4;\xa3\x9f\xa3\x9f\x7f\xec0\xf8\xf4\x98\xbb\xebr\x17W^\xb9\x9a\xb8\xbe\xe6\xe4{\xe8\xf2\x98\xc9k2v\x8dz%\xa9\x89\x05\x0b\xaf\x1f\xf4wv\x87\xfa\xaa\x03\xa8\xee6\xf7\x0f<\x19\x1ce({.\xa7\x92\x7f\xe9!\x8f\x94\x80~c^^C]]\xc9\x82\x06\x1dz\xac\xbc<<\x04\xb2\xa3#!\xd2W\x0e\xe1\xea\x90\xf6\xf0\xf9\x0b\xa7\xa8__\xf9\x0cyE\xcb\xc7\xa9\x84\x026\xcd\x02\xe9\xe8PT\xce0\xe4gz\xd8\x84\xcc\xea\xf4\xcdG\xd05$\x1e\'\xa0/)\x8a\xef\xff\xa5\xa6v\xfa\xa9\x0b\xaa{\x03F\xde\x17\xa1\x90CUb\x1c\xff\xa4H\xfc\x03\xa0S\xa1\xccL8T\xed\x18@\x121{\tI\xd4~\x10\xb83\x9d.\x05\x912\x12CV\xec\xd4\x84=\xa9\xf9M\xe8\x12E\x82OQn\x97\x14\xa8\x06\xce5|\xf5\x93\xbe\xc5\x1e<\xc51x\x8a1X\xb5sP\xf0\x94\x06b4\t\x1e\xb6\xbd\x17LZ\xda\xef\xc2X8Vu\xb6\xe5F-\xca!\xf1(\x01\xd0\xea\n\xd9\xb6\xd6\xc95\xf7\xdb\xd0\x01\xb0i\xec\x13N\xc1K\xa2c\xb4\x11\xd9\xb9E(\x07\xa3\'\x96e]\x02\x98\xab\xafv\xe1R\x90;=m1s\x8fm9\x80R\x94!\x13\xd8\xf1\xf5\xad\\CG\xcf\x1a\x96\x9c%5-\x0f\xd0\x0b\xa4\xe6#p\x05`\xf2\x0b0\x04\xcdxa\x89:5\xd5\x80.!\xe8N\xf0\x18\xf5\xed]a\x8a\xc0n\x94\x82\x9ci\xb9\xba\xf8\xc5\xa3\x1fe\xa0t\x8bh\x02{k\xbb\xe3O\xb5\x86o\xdd3~\x13\x18X\x16\xbd\xce0r\x85\xd1\x1c\xa4\xca\x0e[\xfe\xd65t\xb9\x83\xee\x82\x15\x8e\x8b\xf2\xed\xe3\xe6\x92\':\\\xe1\xc8;\xed\xbf\xb7\xcasjw\xc7\xa0K\x10\xb7\t\x94\x1cnW\x9d\xbc\xc4\xec5\x99Au]\x9d\x08gF\xa6C\xb8C\xc9\xd8z\xc1<,\xb4lR\xc5\x0f\x1e\x9f\x1a6.\xaf\xfc\xb2n\x822r\x17JAh?\x9ch\xd6\x1dE)\x88\xdb\x04\xb4\xff\xf8\x14\xb4vqSQ\xc2\xcf\xe3\xf3!\xc8?\x18\x15@\xa3\xa1\x1b\x98(M\xfd\xf0\xb5\x9b\x0b\xd15\x14\xf9\xf8\xa4T\xd1\x7f\x1a\x96\x94\xe6\x8dT\x868\x9f\xc9\x83h\xed\xef\x98\xb8\xecD\xf1\xf7(]\x10L\xa0\xa0\xaa7\xaf\xad\x13\xf2P\xf2\xd0\x1eP\xa7mC\x05\xb0\xe6\xd2u\xa8\x92\x07mV\'$\x88\xd5\xbd\r\xdb\xb8a\x83O\x802\x11F\xbc\x14\xa5T\xa3\x12\xa6\xfa\xef\x8b\xaf\xa5W\x17\xceA\xe9\x84\xcb=\xb0\xa1\xb2/Fw\xc7\xd2HK\x07]N\\\xd8\x1e\x0b\xdd]\xad\xa8\xec\xa7\x8b\xbdA\x05\xa0\xc1\xd3\xf7\xb6Q\x82\xfe\x9f\x91\xd6\xa5\x95\x98\xcf\x96\xa6\xb4+\xcd\xb5b\xb7\x92\x8c\xdc\x0f\xeb\x9e\x9a\x9b\xb01i\x91\xd3\xfd\xe0\xba\x03\x0c\xf7\xaa\xbb\xe0)\xe3\x16\x1f\xc4\x15\x0f\rJhh\xa3~\xdb\x05F\x0f}[\xf0\x14\xfbqj\xa3iuE\xdd\xbc@\xf68JA\xcc\xa4\x1f\xaa\xee\x9eOAi\xc7%\x81\xf3\xe7Z\xf6+e\xd6YE\x88\xe0\xe8i\x90\xb8\xea$*;4x\xc7\xc1\x8e\x06M\x13\x1b<\x89\xd2\xec]v\x8b\xfb##\xe3\x8b\xf0\xba\xb4\x99\xb2\x16\xf4\xb8\xa2`d\x86\x85c\xa6\x9fEi\xc7%\x81\xba\xd2X\x83Zn\xcaD)H\x80f\x1a\xcc\xfc\xf0\xafN\x7f\xff\xa0c\xe8\xf2\x94\x85\xc4\x1c\x9b\x19\xb8\x9b\x1b\xa2\xb9\xae\xc6r\xce\xd8\x01\x9f\x87\\\x84\x11L\x1f\xbey\x88\\"\x03\xa9\x91]\xb1>a\xfeut\xd9\x11l\xe2\x8a\xfc\xd0#\xc3\x14\xfdn;\x9f\xc7O\x134{[s Y\xd1\x92\xa0M9\xf8/Kk\xbd\x84\x18\xadw{\xd3\x12h3\xf3Ip\\\xbb\n\xee\xed+c\xcc\xf7TRrB\xab\xa4&\xd8\xedw\xc1\xba\x87\x0e\xc8\x8c\\i\xf3\xdb\x07\xaaQ:\xe1\xb6m\xc8E"M+\xea\xd4\xb2\x8c4\x1e]\x82L\xd20\x85\xa5\xcb\xfd?F)\x06-)\xfbLDH\xe6\xb4S\xd2\xa1\xe7\xf2\x07\xa8y\x8c\xe4~.\xbe?\x16\xb6\xb3q\xbcVI\x94:\x06\xb8\xd4\xa6\xec}\x82\xa7\x9d\xe0\x0eP\x18\x86\xb1<;\x9a\xcdTH\xdd\xf7\x03\xe5J3W\x90[\xde\xe5\xd2\\\x02$\x13s\xdc\xa5\xfa\x06\xed\xcd\xf5\xb8\xb6#\x0f\x00x?\xe2\x1a<\xcf\xdc\x07_\xc6\xc70V\x15\x99\xed.x\x8a\xdb\x04(\x85\x8b\xd4M*\x85x?X\xc8\x84q\xbb\x959\xdcx\xbbO\xf4\xcaGh\xb9\xd9G\x8a\xe4\x1c\x83D\xdf\xea:N\xf8\xf8\x01\xec\xd7\xfc\x06\xe3\xc0\x94s<\xf3\x13\xd1\x8bN4\x01J%\xe9\x870\xdf~\xd1Q\xa1\xcf\x04\xaa-G,\xe5(\xc5\xd0ge\xcd\x99\x1d\xa6\x96\x0e\xa0\x86\xa4U&\x10JB\xed\'+=\xfd\xe6\xc1\xefP\xbae\xc8\x04(\x07\xd7\x0e[*\x03\x8b\xe8_\xa2\xa3\x17R\x96\xef\xeeu:a\x84\xa8\xcc\xbf1\xef\xc7"\xa9\x02%\x1f<M\xc2\tY\xb0\x0e\xfc5\x1e\x8d(\x1e%@\xfa\xa1\'i\x14\x9b\xe9\xc3\x88\xf7C\xb3\x81-&7\xb9\xdb~\xe0\xfe\\\x91!a\r[\x93\xe2\x18\xa8/\xf3A\xaf5\x89\xe4w0\tR\xf7\x10<+\x9bI\xd0y2\xa2x\x96\x00\x85\xf6C\x80\xb4\x7f\xc8~\xb8\xa2\xb7\x1c\xfe]\xcf9\xde\xbc<\x9c6%\x14:N\x97\x83\xd9\x1aWR\x9c\x04\x0e\x15>L\xe2\xd3w\xa5$xr\x96\xcaG\xacf\x9e\xfeVt\xb7\x1dy\xa4/\xf5\x94\x8c-\xfa7dJ\x19?j\xd8\xb2\xb7}S\xb0\xe9\xd8\x08\x85n\xc72u\x1dJ\x1e\xee\xec\x98x\x90\xca\x1cv\x87~\x9a\x85\xa5\x05\xcd3*\n5\xbf\xf2\xbf\xc5\'\xd8\x00\x93\x1b\xbe!;.\xfc\xe5\xc3\x8b\x17/^\xbc<^\x00\xfe\x03t\xd3\x12d\xcd*\xb7\xe1\x00\x00\x00\x00IEND\xaeB`\x82'
icon_id = idaapi.load_custom_icon(data=icon, format="png")

class util:
    def syntax_highlight_proto(line):
        colored_line = line
        # Strings
        str = re.search("\".*\"",colored_line)
        if str:
            colored_line = colored_line[:str.start()] + idaapi.COLSTR(colored_line[str.start():str.end()], idaapi.SCOLOR_DSTR) + colored_line[str.end():]
        # Variables
        var = re.search("\w+\s*(?==)",colored_line)
        if var:
            colored_line = colored_line[:var.start()] + idaapi.COLSTR(colored_line[var.start():var.end()], idaapi.SCOLOR_CODNAME) + colored_line[var.end():]
        # Objects
        obj = re.search("\w+\s*(?={)",colored_line)
        if obj:
            colored_line = colored_line[:obj.start()] + idaapi.COLSTR(colored_line[obj.start():obj.end()], idaapi.SCOLOR_IMPNAME) + colored_line[obj.end():]
        # Numbers
        num = re.search("\W\d+\W",colored_line)
        if num:
            colored_line = colored_line[:num.start()] + idaapi.COLSTR(colored_line[num.start():num.end()], idaapi.SCOLOR_DNUM) + colored_line[num.end():]

        return colored_line


# IDA plugin Action Handler
class ProtobufFinder(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    # Called when the Search -> Protobuf Finder is clicked
    def activate(self, ctx):
        print("[Protobuf] Search started.")
        runProtod() # run search
        print("[Protobuf] Search done.")
        pass

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


class protobuf_fetch_t(idaapi.plugin_t):
    comment = "Protobuf string finder."
    help = "A plugin to find a parse embedded proto files in binaries with basic syntax highlight."
    wanted_name = "Protobuf"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        self.printBanner()

        protobuf_desc = idaapi.action_desc_t(
            'protobuf:search',      # The action name. This acts like an ID and must be unique
            'Protobuf Finder',      # The action text.
             ProtobufFinder(),      # The action handler.
            '',                     # Optional: the action shortcut
            '',                      # Optional: the action tooltip (available in menus/toolbar)
            icon_id)                      # Optional: the action icon (shows when in menus/toolbars)
        idaapi.register_action(protobuf_desc)

        # Insert the action in the menu
        if not idaapi.attach_action_to_menu("Search", "protobuf:search", idaapi.SETMENU_APP):
            print("[Protobuf] Failed attaching to menu.")

        return idaapi.PLUGIN_KEEP
      
    def run(self):
        pass

    def term(self):
        pass

    def printBanner(self):
        print("---------------------------------------------------------------------------------------------")
        print("Protobuf string finder loaded.")
        print("Run via Search -> Protobuf Finder")
        print("---------------------------------------------------------------------------------------------")


def runProtod():
    extractor = ProtobufExtractor()
    extractor.extract()


def PLUGIN_ENTRY():
    return protobuf_fetch_t()


# =======================================================================================
# rewritten Protod code to Python3, using IDA API
# =======================================================================================

###########
# helpers
###########

def is_valid_filename(filename):
    '''
    Check if given filename may be valid
    '''
    charset = 'abcdefghijklmnopqrstuvwxyz0123456789-_/$,.[]()'
    for char in filename.lower():
        if chr(char) not in charset:
            return False
    return True


def decode_varint128(stream):
    '''
    Decode Varint128 from buffer
    '''
    bits = ''
    count = 0
    for stream_byte in stream:
        count += 1
        raw_byte = stream_byte
        bits += (bin((raw_byte&0x7F))[2:]).rjust(7,'0')
        if (raw_byte&0x80) != 0x80:
            break
    return (int(bits, 2), count)


def render_type(field_type, package):
    '''
    Return the string representing a given type inside a given package
    '''
    i = 0
    nodes = field_type.split('.')
    nodes_ = package.split('.')
    for i in range(len(nodes)):
        if i < len(nodes_):
            if nodes[i] != nodes_[i]:
                return '.'.join(nodes[i:])
        else:
            return '.'.join(nodes[i:])
    return '.'.join(nodes[i:])


#############################
# Protobuf fields walker
#############################

class ProtobufFieldsWalker:
    '''
    Homemade Protobuf fields walker

    This class allows Protod to walk the fields
    and determine the probable size of the protobuf
    serialized file.
    '''

    def __init__(self, stream):
        self._stream = stream
        self._size =  -1

    def get_size(self):
        return self._size

    def walk(self):
        end = False
        offset = 0
        while (not end) and (offset<len(self._stream)):
            # read tag
            tag = self._stream[offset]
            offset += 1
            wire_type = tag&0x7
            if wire_type == 0:
                value, size = decode_varint128(self._stream[offset:])
                offset += size
            elif wire_type == 1:
                offset += 8
            elif wire_type == 2:
                value, size = decode_varint128(self._stream[offset:])
                offset += size + value
            elif wire_type == 5:
                offset += 4
            elif wire_type == 3:
                continue
            elif wire_type == 4:
                continue
            else:
                end = True
        self._size = offset-1


#############################
# Serialized metadata parsing
#############################

class FileDescriptorDisassembler:
    '''
    Core disassembling class

    This class parses the provided serialized data and
    produces one or many .proto files.
    '''

    def __init__(self, file_desc):
        self.desc = file_desc

    def getLabel(self, l):
        return [None, 'optional', 'required', 'repeated'][l]

    def getTypeStr(self, t):
        types = [
            None,
            'double',
            'float',
            'int64',
            'uint64',
            'int32',
            'fixed64',
            'fixed32',
            'bool',
            'string',
            'group',
            'message',
            'bytes',
            'uint32',
            'enum',
            'sfixed32',
            'sfixed64',
            'sint32',
            'sint64'
        ]
        return types[t]

    def renderEnum(self, enum, depth=0, package='', nested=False):
        buffer = '\n'
        buffer += '%senum %s {\n' % (' '*depth, enum.name)
        for value in enum.value:
            buffer += '%s%s = %d;\n' % (' '*(depth+1), value.name, value.number)
        buffer += '%s}' % (' '*depth)
        buffer += '\n'
        return buffer

    def renderField(self, field, depth=0, package='', nested=False):
        buffer = ''
        try:
            if field.HasField('type'):
                # message case
                if field.type == FieldDescriptorProto.TYPE_MESSAGE or field.type == FieldDescriptorProto.TYPE_ENUM:
                    field.type_name = render_type(field.type_name[1:], package)
                    buffer += '%s%s %s %s = %d;\n' % (' '*depth, self.getLabel(field.label), field.type_name, field.name, field.number)
                else:
                    if field.HasField('default_value'):
                        if self.getTypeStr(field.type) == 'string':
                            field.default_value = '"%s"'% field.default_value
                        buffer += '%s%s %s %s = %d [default = %s];\n' % (' '*depth, self.getLabel(field.label), self.getTypeStr(field.type), field.name, field.number, field.default_value)
                    else:
                        buffer += '%s%s %s %s = %d;\n' % (' '*depth, self.getLabel(field.label), self.getTypeStr(field.type), field.name, field.number)
        except ValueError:
            buffer += '%smessage %s {\n' % (' '*depth, field.name)
            _package = package+'.'+field.name

            if len(field.nested_type)>0:
                for nested in field.nested_type:
                    buffer += self.renderField(nested, depth+1, _package, nested=True)
            if len(field.enum_type)>0:
                for enum in field.enum_type:
                    buffer += self.renderEnum(enum, depth+1, _package)
            if len(field.field)>0:
                for field in field.field:
                    buffer += self.renderField(field, depth+1, _package)
            buffer += '%s}' % (' '*depth)
            buffer += '\n\n'
        return buffer


    def render(self, filename=None):
        print('[Protobuf][+] Processing %s' % self.desc.name)
        buffer = ''
        buffer += 'package %s;\n\n' % self.desc.package

        # add dependencies
        if len(self.desc.dependency)>0:
            for dependency in self.desc.dependency:
                buffer += 'import "%s";\n' % dependency
            buffer += '\n'

        if len(self.desc.enum_type)>0:
            for enum in self.desc.enum_type:
                buffer += self.renderEnum(enum, package=self.desc.package)
        if len(self.desc.message_type)>0:
            for message in self.desc.message_type:
                buffer += self.renderField(message, package=self.desc.package)
        
        # print result into IDA window
        textWindow = ida_kernwin.simplecustviewer_t()
        if textWindow.Create("Protobuf file: " + self.desc.name):
            for line in buffer.split('\n'):
                textWindow.AddLine(util.syntax_highlight_proto(line))
            textWindow.Show()
        else:
            print ("[Protobuf] Failed to open window - it is probably already open.")


#############################
# Main code
#############################

class ProtobufExtractor:
    def __init__(self):
        pass

    def extract(self):
        protos = []
        searchStartAddr = 0
        
        while True:
            # search binary for ".proto" string
            r = ida_bytes.bin_search(searchStartAddr,ida_ida.MAXADDR,bytes([0x2E, 0x70, 0x72, 0x6F, 0x74,0x6F]),bytes([0xFF,0xFF,0xFF,0xFF,0xFF,0xFF]),1,1)
            if r == idaapi.BADADDR:
                print("[Protobuf][dbg] Search results into BADADDR (not found). Break!")
                break

            searchStartAddr = r+1
            print(f"[Protobuf][dbg] String \".proto\" found at {hex(r)}.")
            
            # Find beginning of the whole proto string
            for j in range(64):
                try:
                    protostring = ida_bytes.get_bytes(r-j-1, MAX_PROTOBUF_SIZE, False)
                    
                    if decode_varint128(protostring[1:])[0]==(j+5) and is_valid_filename(protostring[1+1:j+6+1]):
                        print(f"[Protobuf][dbg] Protostring identified: {protostring[:40]} [...]")
                       
                        # Walk the fields and get a probable size
                        walker = ProtobufFieldsWalker(protostring)
                        walker.walk()
                        probable_size = walker.get_size()
                       
                        print(f"[Protobuf][dbg] Protostring section size {hex(probable_size)}.")
                       
                        """
                        Probable size approach is not perfect,
                        we add a delta of 1024 bytes to be sure
                        not to miss something =)
                        """
                        for k in range(probable_size+1024, 0, -1):
                            try:
                                fds = FileDescriptorProto()
                                fds.ParseFromString(protostring[:k])
                                protos.append(protostring[:k])
                                print('[Protobuf][dbg][i] Found protofile %s (%d bytes)' % (protostring[:j+5], k))
                                break
                            except DecodeError:
                                pass
                            except UnicodeDecodeError:
                                pass
                        break
                except IndexError:
                    pass
        if protos:
            # Load successively each binary proto file and rebuild it from scratch
            seen = []
            for content in protos:
                try:
                    # Load the prototype
                    fds  = FileDescriptorProto()
                    fds.ParseFromString(content)
                    res = FileDescriptorDisassembler(fds)
                    if len(res.desc.name)>0:
                        if res.desc.name not in seen:
                            res.render() # disassemble and show .proto file in IDA window
                            seen.append(res.desc.name)
                except DecodeError:
                    pass
        else:
            ida_kernwin.info("No embedded proto files were discovered.")


repository: ChiChou/IDA-ObjCExplorer
file: ObjCExplore.py
description: Objective C classdump for IDA Pro
comment: classdump
created_at: 2020-08-09 17:08:41+00:00
pushed_at: 2020-09-01 08:33:16+00:00
url:
https://github.com/ChiChou/IDA-ObjCExplorer/blob/f5b584975622bc0751a8a2f007f0c932830283e1/ObjCExplore.py

content:
import struct
from collections import namedtuple

import idc
import idaapi
import ida_segment
import ida_bytes
import ida_nalt
import ida_idaapi

import ida_kernwin as kw
from idaapi import PluginForm
from PyQt5 import QtWidgets, QtGui


PLUGIN_NAME = 'ObjCExplorer'


def cstr(ea):
    try:
        return ida_bytes.get_strlit_contents(ea, -1, ida_nalt.STRTYPE_C).decode()
    except Exception as e:
        print('Unable to decode string at %s' % hex(ea))
        raise e


class Objc2Class(object):
    """
    struct __objc2_class
    {
        __objc2_class *isa;
        __objc2_class *superclass;
        void *cache;
        void *vtable;
        __objc2_class_ro *info;
    };
    """
    fmt = '<QQQQQ'
    length = struct.calcsize(fmt)

    def __init__(self, data, offset=0):
        (self.isa,
         self.superclass,
         self.cache,
         self.vtable,
         self.info) = struct.unpack_from(self.fmt, data, offset)


class Objc2ClassRo(object):
    """
    struct __objc2_class_ro
    {
        uint32_t flags;
        uint32_t ivar_base_start;
        uint32_t ivar_base_size;
        uint32_t reserved;
        void *ivar_lyt;
        char *name;
        __objc2_meth_list *base_meths;
        __objc2_prot_list *base_prots;
        __objc2_ivar_list *ivars;
        void *weak_ivar_lyt;
        __objc2_prop_list *base_props;
    };
    """
    fmt = '<IIIIQQQQQQQ'
    length = struct.calcsize(fmt)

    def __init__(self, data, offset=0):
        (self.flags,
         self.ivar_base_start,
         self.ivar_base_size,
         self.reserved,
         self.ivar_lyt,
         self.name,
         self.base_meths,
         self.base_prots,
         self.ivars,
         self.weak_ivar_lyt,
         self.base_props) = struct.unpack_from(self.fmt, data, offset)


class Objc2Method(object):
    fmt = '<QQQ'
    length = struct.calcsize(fmt)

    def __init__(self, data, offset=0):
        (self.name, self.types, self.imp) = struct.unpack_from(
            self.fmt, data, offset)

Post14Method = namedtuple('Method', ['name', 'types', 'imp'])

def method_list(ea):
    if not ea:
        return

    count = ida_bytes.get_dword(ea + 4)
    name = idc.get_segm_name(ea)
    first = ea + 8

    def post14format(addr):
        for _ in range(3):
            data = ida_bytes.get_bytes(addr, 4)
            offset, = struct.unpack('<i', data)
            yield addr + offset
            addr += 4
        
    is14 = name and (name.endswith(':__objc_const_ax') or name.endswith(':__objc_methlist'))
    for i in range(count):
        if is14:
            # iOS 14
            yield Post14Method(*post14format(first + i * 12))

        else:
            ea_method_t = first + i * Objc2Method.length
            data = ida_bytes.get_bytes(ea_method_t, Objc2Method.length)
            yield Objc2Method(data)


class Base(object):
    def __init__(self, name, ea):
        self.name = name
        self.ea = ea

    def __repr__(self):
        return '<%s "%s">' % (self.__class__.__name__, self.name)


class Clazz(Base):
    def __init__(self, name, ea):
        super().__init__(name, ea)
        self.methods = {}


class Protocol(Base):
    def __init__(self, name, ea):
        super().__init__(name, ea)
        self.methods = []


class ClassDump(object):
    def __init__(self, output=None, verbose=False):
        self.classes = []
        self.protocols = []
        self.class_lookup = {}
        self.protocol_lookup = {}
        self.lookup = {}
        self.output = output
        self.verbose = verbose

    def print(self, *args):
        if self.output is not None:
            print(*args, file=self.output)
        elif self.verbose:
            print(*args)

    def parse(self):
        if ida_segment.get_segm_by_name('DYLD_CACHE_HEADER'):
            seg = ida_segment.get_first_seg()

            def handle(seg):
                name = ida_segment.get_segm_name(seg)
                try:
                    mod, segname = name.split(':')
                except ValueError:
                    return

                if segname == '__objc_protolist':
                    self.handle_proto_seg(seg)
                elif segname == '__objc_classlist':
                    self.handle_class_seg(seg)

            while seg:
                handle(seg)
                seg = ida_segment.get_next_seg(seg.start_ea)

            return

        protocols = ida_segment.get_segm_by_name('__objc_protolist')
        if protocols:
            self.handle_proto_seg(protocols)

        classes = ida_segment.get_segm_by_name('__objc_classlist')
        if classes:
            self.handle_class_seg(classes)

    def handle_proto_seg(self, protocols):
        for ea in range(protocols.start_ea, protocols.end_ea, 8):
            self.handle_protocol(ea)

            if not (self.output or self.verbose) and len(self.protocols) > 4096:
                print('Threshold exceeded')
                break

    def handle_class_seg(self, classes):
        for ea in range(classes.start_ea, classes.end_ea, 8):
            self.handle_class(ea)

            if not (self.output or self.verbose) and len(self.protocols) > 4096:
                print('Threshold exceeded')
                break

    def handle_protocol(self, ea):
        protocol_ea = ida_bytes.get_qword(ea)
        protocol_name = cstr(ida_bytes.get_qword(protocol_ea + 8))
        method_list_ea = ida_bytes.get_qword(protocol_ea + 3 * 8)
        p = Protocol(protocol_name, ea=protocol_ea)
        self.print('@protocol', protocol_name)
        # todo: support class methods
        for method in method_list(method_list_ea):
            key = '- ' + cstr(method.name)
            p.methods.append(key)
            self.print(key)
        self.print('@end')
        self.print()
        self.protocols.append(p)
        self.protocol_lookup[p.name] = p
        self.lookup[ea] = p

    def handle_class(self, ea):
        clazz_ea = ida_bytes.get_qword(ea)
        clazz = Objc2Class(ida_bytes.get_bytes(clazz_ea, Objc2Class.length))
        # if clazz.info & 7 != 0:
        # swift

        meta_class = Objc2Class(
            ida_bytes.get_bytes(clazz.isa, Objc2Class.length))
        meta_class.info = (meta_class.info >> 3) << 3
        meta_info = Objc2ClassRo(ida_bytes.get_bytes(
            meta_class.info, Objc2ClassRo.length))

        clazz.info = (clazz.info >> 3) << 3
        clazz_info = Objc2ClassRo(ida_bytes.get_bytes(
            clazz.info, Objc2ClassRo.length))

        c = Clazz(cstr(clazz_info.name), ea=clazz_ea)

        self.print('@interface', cstr(clazz_info.name))
        for method in method_list(meta_info.base_meths):
            key = '+ ' + cstr(method.name)
            c.methods[key] = method.imp
            self.print(key)

        for method in method_list(clazz_info.base_meths):
            key = '- ' + cstr(method.name)
            c.methods[key] = method.imp
            self.print(key)

        self.print('@end')
        self.print()

        self.classes.append(c)
        self.class_lookup[c.name] = c
        self.lookup[ea] = c


class ClassView(PluginForm):
    # todo: classref
    cols = ['Name', 'Address']

    def __init__(self):
        super(ClassView, self).__init__()
        self.data = None
        self.tree = None

    def dblclick(self, item):
        '''Handle double click event.'''
        try:
            idaapi.jumpto(int(item.text(1), 16))
        except:
            pass

    def load_data(self):
        classdump = ClassDump()
        classdump.parse()

        self.data = classdump

        class_root = QtWidgets.QTreeWidgetItem(self.tree)
        class_root.setText(0, 'Classes')
        class_root.setExpanded(True)

        x = lambda ea: '0x%X' % ea

        for clazz in self.data.classes:
            item = QtWidgets.QTreeWidgetItem(class_root)
            item.setText(0, clazz.name)
            item.setText(1, x(clazz.ea))

            for method, imp in clazz.methods.items():
                child = QtWidgets.QTreeWidgetItem(item)
                child.setText(0, method)
                child.setText(1, x(imp))

        protocol_root = QtWidgets.QTreeWidgetItem(self.tree)
        protocol_root.setText(0, 'Protocols')
        
        # todo: refactor
        for proto in self.data.protocols:
            item = QtWidgets.QTreeWidgetItem(protocol_root)
            item.setText(0, proto.name)
            item.setText(1, x(proto.ea))

            for method in proto.methods:
                child = QtWidgets.QTreeWidgetItem(item)
                child.setText(0, method)
                # todo:
                child.setText(1, x(proto.ea))


    def OnCreate(self, form):
        '''Called when the plugin form is created'''

        self.parent = self.FormToPyQtWidget(form)

        self.tree = QtWidgets.QTreeWidget()
        self.tree.setColumnCount(len(self.cols))
        self.tree.setHeaderLabels(self.cols)
        self.tree.itemDoubleClicked.connect(self.dblclick)

        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.tree)
        layout.setSpacing(0)
        layout.setContentsMargins(0, 0, 0, 0)

        self.load_data()

        self.tree.setColumnWidth(0, 512)
        self.tree.setColumnWidth(1, 512)
        self.parent.setLayout(layout)

    def OnClose(self, form):
        '''Called when the plugin form is closed.'''
        del self

    def Show(self):
        '''Creates the form is not created or focuses it if it was.'''
        return PluginForm.Show(self, 'ClassDump')


class ObjCExplorer(ida_idaapi.plugin_t):
    """Class that is required for the code to be recognized as
    a plugin by IDA."""
    flags = 0
    comment = "classdump"
    help = comment
    wanted_name = PLUGIN_NAME
    wanted_hotkey = "Ctrl-Shift-E"

    def is_compatible(self):
        return idaapi.IDA_SDK_VERSION >= 700

    def init(self):
        return (ida_idaapi.PLUGIN_OK if
                self.is_compatible() else ida_idaapi.PLUGIN_SKIP)

    def run(self, arg):
        kw.show_wait_box('Building class information')
        try:
            view = ClassView()
            view.Show()
        finally:
            kw.hide_wait_box()

    def term(self):
        pass


# todo: add this to menu
def export():
    path = kw.ask_file(1, '*.txt', 'Save to')
    if not path:
        print('Cancelled')
        return

    with open(path, 'a') as fp:
        classdump = ClassDump(output=fp)
        classdump.parse()


def PLUGIN_ENTRY():
    return ObjCExplorer()


repository: Russinovich/IDA7.5_SDK
file: python/examples/hexrays/vds17.py
description: SDK and IDAPython for IDA Pro 7.5
wanted_name: Structure offsets (IDAPython)
comment: Sample plugin17 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/vds17.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 shows how to use "Select offsets" widget (select_udt_by_offset() call).
#      This plugin repeats the Alt-Y functionality.
#      Usage: place cursor on the union field and press Shift-T
#
#      This is a rewrite in Python of the vds17 example that comes with hexrays sdk.
#

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: clovme/WTools
file: software/Reverse Cracking/IDA Professional 9.0/python/examples/hexrays/vds5.py
wanted_name: Hex-Rays show C graph (IDAPython)
comment: Sample plugin5 for Hex-Rays decompiler
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/hexrays/vds5.py

content:
"""
summary: show ctree graph

description:
  Registers an action that can be used to show the graph of the ctree.
  The current item will be highlighted in the graph.

  The command shortcut is `Ctrl+Shift+G`, and is also added
  to the context menu.

  To display the graph, we produce a .gdl file, and
  request that ida displays that using `ida_gdl.display_gdl`.
"""

import ida_pro
import ida_hexrays
import ida_kernwin
import ida_gdl
import ida_lines
import ida_idaapi

ACTION_NAME = "vds5.py:displaygraph"
ACTION_SHORTCUT = "Ctrl+Shift+G"

CL_WHITE            = ((255)+  (255<<8)+  (255<<16)) #   0
CL_BLUE             = ((0  )+  (0  <<8)+  (255<<16)) #   1
CL_RED              = ((255)+  (0  <<8)+  (0  <<16)) #   2
CL_GREEN            = ((0  )+  (255<<8)+  (0  <<16)) #   3
CL_YELLOW           = ((255)+  (255<<8)+  (0  <<16)) #   4
CL_MAGENTA          = ((255)+  (0  <<8)+  (255<<16)) #   5
CL_CYAN             = ((0  )+  (255<<8)+  (255<<16)) #   6
CL_DARKGREY         = ((85 )+  (85 <<8)+  (85 <<16)) #   7
CL_DARKBLUE         = ((0  )+  (0  <<8)+  (128<<16)) #   8
CL_DARKRED          = ((128)+  (0  <<8)+  (0  <<16)) #   9
CL_DARKGREEN        = ((0  )+  (128<<8)+  (0  <<16)) #  10
CL_DARKYELLOW       = ((128)+  (128<<8)+  (0  <<16)) #  11
CL_DARKMAGENTA      = ((128)+  (0  <<8)+  (128<<16)) #  12
CL_DARKCYAN         = ((0  )+  (128<<8)+  (128<<16)) #  13
CL_GOLD             = ((255)+  (215<<8)+  (0  <<16)) #  14
CL_LIGHTGREY        = ((170)+  (170<<8)+  (170<<16)) #  15
CL_LIGHTBLUE        = ((128)+  (128<<8)+  (255<<16)) #  16
CL_LIGHTRED         = ((255)+  (128<<8)+  (128<<16)) #  17
CL_LIGHTGREEN       = ((128)+  (255<<8)+  (128<<16)) #  18
CL_LIGHTYELLOW      = ((255)+  (255<<8)+  (128<<16)) #  19
CL_LIGHTMAGENTA     = ((255)+  (128<<8)+  (255<<16)) #  20
CL_LIGHTCYAN        = ((128)+  (255<<8)+  (255<<16)) #  21
CL_LILAC            = ((238)+  (130<<8)+  (238<<16)) #  22
CL_TURQUOISE        = ((64 )+  (224<<8)+  (208<<16)) #  23
CL_AQUAMARINE       = ((127)+  (255<<8)+  (212<<16)) #  24
CL_KHAKI            = ((240)+  (230<<8)+  (140<<16)) #  25
CL_PURPLE           = ((160)+  (32 <<8)+  (240<<16)) #  26
CL_YELLOWGREEN      = ((154)+  (205<<8)+  (50 <<16)) #  27
CL_PINK             = ((255)+  (192<<8)+  (203<<16)) #  28
CL_ORANGE           = ((255)+  (165<<8)+  (0  <<16)) #  29
CL_ORCHID           = ((218)+  (112<<8)+  (214<<16)) #  30
CL_BLACK            = ((0  )+  (0  <<8)+  (0  <<16)) #  31

COLORS_LUT = {
    CL_WHITE             : "white",
    CL_BLUE              : "blue",
    CL_RED               : "red",
    CL_GREEN             : "green",
    CL_YELLOW            : "yellow",
    CL_MAGENTA           : "magenta",
    CL_CYAN              : "cyan",
    CL_DARKGREY          : "darkgrey",
    CL_DARKBLUE          : "darkblue",
    CL_DARKRED           : "darkred",
    CL_DARKGREEN         : "darkgreen",
    CL_DARKYELLOW        : "darkyellow",
    CL_DARKMAGENTA       : "darkmagenta",
    CL_DARKCYAN          : "darkcyan",
    CL_GOLD              : "gold",
    CL_LIGHTGREY         : "lightgrey",
    CL_LIGHTBLUE         : "lightblue",
    CL_LIGHTRED          : "lightred",
    CL_LIGHTGREEN        : "lightgreen",
    CL_LIGHTYELLOW       : "lightyellow",
    CL_LIGHTMAGENTA      : "lightmagenta",
    CL_LIGHTCYAN         : "lightcyan",
    CL_LILAC             : "lilac",
    CL_TURQUOISE         : "turquoise",
    CL_AQUAMARINE        : "aquamarine",
    CL_KHAKI             : "khaki",
    CL_PURPLE            : "purple",
    CL_YELLOWGREEN       : "yellowgreen",
    CL_PINK              : "pink",
    CL_ORANGE            : "orange",
    CL_ORCHID            : "orchid",
    CL_BLACK             : "black",
}

def get_color_name(c):
    return COLORS_LUT[c] if c in COLORS_LUT.keys() else "?"

class cfunc_graph_t: # alas we can't inherit gdl_graph_t
    def __init__(self, highlight):
        self.items = [] # list of citem_t
        self.highlight = highlight
        self.succs = [] # list of lists of next nodes
        self.preds = [] # list of lists of previous nodes

    def nsucc(self, n):
        return len(self.succs[n]) if self.size() else 0

    def npred(self, n):
        return len(self.preds[n]) if self.size() else 0

    def succ(self, n, i):
        return self.succs[n][i]

    def pred(self, n, i):
        return self.preds[n][i]

    def size(self):
        return len(self.preds)

    def add_node(self):
        n = self.size()

        def resize(array, new_size):
            if new_size > len(array):
                while len(array) < new_size:
                    array.append([])
            else:
                array = array[:new_size]
            return array

        self.preds = resize(self.preds, n+1)
        self.succs = resize(self.succs, n+1)
        return n

    def add_edge(self, x, y):
        self.preds[y].append(x)
        self.succs[x].append(y)

    def get_expr_name(self, expr):
        name = expr.print1(None)
        name = ida_lines.tag_remove(name)
        name = ida_pro.str2user(name)
        return name

    def get_node_label(self, n):
        item = self.items[n]
        op = item.op
        insn = item.cinsn
        expr = item.cexpr
        parts = [ida_hexrays.get_ctype_name(op)]
        if op == ida_hexrays.cot_ptr:
            parts.append(".%d" % expr.ptrsize)
        elif op == ida_hexrays.cot_memptr:
            parts.append(".%d (m=%d)" % (expr.ptrsize, expr.m))
        elif op == ida_hexrays.cot_memref:
            parts.append(" (m=%d)" % (expr.m,))
        elif op in [
                ida_hexrays.cot_obj,
                ida_hexrays.cot_var]:
            name = self.get_expr_name(expr)
            parts.append(".%d %s" % (expr.refwidth, name))
        elif op in [
                ida_hexrays.cot_num,
                ida_hexrays.cot_helper,
                ida_hexrays.cot_str]:
            name = self.get_expr_name(expr)
            parts.append(" %s" % (name,))
        elif op == ida_hexrays.cit_goto:
            parts.append(" LABEL_%d" % insn.cgoto.label_num)
        elif op == ida_hexrays.cit_asm:
            parts.append("<asm statements; unsupported ATM>")
            # parts.append(" %a.%d" % ())
        parts.append("\\n")
        parts.append("ea: %08X" % item.ea)
        if item.is_expr() and not expr.type.empty():
            parts.append("\\n")
            tstr = expr.type._print()
            parts.append(tstr if tstr else "?")
        return "".join(parts)

    def get_node_color(self, n):
        item = self.items[n]
        if self.highlight is not None and item.obj_id == self.highlight.obj_id:
            return CL_GREEN
        return None

    def gen_gdl(self, fname):
        with open(fname, "wb") as out:
            def write_out(s):
                out.write(s.encode("UTF-8"))
            write_out("graph: {\n")

            write_out("// *** nodes\n")
            for n in range(len(self.items)):
                item = self.items[n]
                node_label = self.get_node_label(n)
                node_props = [""]
                if n == 0:
                    node_props.append("vertical_order: 0")
                color = self.get_node_color(n)
                if color is not None:
                    node_props.append("color: %s" % get_color_name(color))
                write_out("""node: { title: "%d" label: "%d: %s" %s}\n""" % (
                    n,
                    n,
                    node_label,
                    " ".join(node_props)))

            write_out("// *** edges\n")
            for n in range(len(self.items)):
                item = self.items[n]
                write_out("// edges %d -> ?\n" % n)
                for i in range(self.nsucc(n)):
                    t = self.succ(n, i)
                    label = ""
                    if item.is_expr():
                        target = self.items[t]
                        if item.x is not None and item.x == target:
                            label = "x"
                        elif item.y is not None and item.y == target:
                            label = "y"
                        elif item.z is not None and item.z == target:
                            label = "z"
                        if label:
                            label = """ label: "%s" """ % label
                    write_out("""edge: { sourcename: "%s" targetname: "%s"%s}\n""" % (
                        str(n), str(t), label))

            write_out("}\n")

    def dump(self):
        print("%d items:" % len(self.items))
        for i in self.items:
            print("\t%s (%08x)" % (i, i.ea))

        print("succs:")
        for s in self.succs:
            print("\t%s" % s)

        print("preds:")
        for p in self.preds:
            print("\t%s" % p)


class graph_builder_t(ida_hexrays.ctree_parentee_t):

    def __init__(self, cg):
        ida_hexrays.ctree_parentee_t.__init__(self)
        self.cg = cg
        self.reverse = [] # (citem_t, node#) tuples

    def add_node(self, i):
        for k, _ in self.reverse:
            if i.obj_id == k.obj_id:
                ida_kernwin.warning("bad ctree - duplicate nodes! (i.ea=%x)" % i.ea)
                self.cg.dump()
                return -1

        n = self.cg.add_node()
        if n <= len(self.cg.items):
            self.cg.items.append(i)
        self.cg.items[n] = i
        self.reverse.append((i, n))
        return n

    def process(self, i):
        n = self.add_node(i)
        if n < 0:
            return n
        if len(self.parents) > 1:
            lp = self.parents.back().obj_id
            for k, v in self.reverse:
                if k.obj_id == lp:
                    p = v
                    break
            self.cg.add_edge(p, n)
        return 0

    def visit_insn(self, i):
        return self.process(i)

    def visit_expr(self, e):
        return self.process(e)


class display_graph_ah_t(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def activate(self, ctx):
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        vu.get_current_item(ida_hexrays.USE_KEYBOARD)
        highlight = vu.item.e if vu.item.is_citem() else None

        cg = cfunc_graph_t(highlight)
        gb = graph_builder_t(cg)
        gb.apply_to(vu.cfunc.body, None)

        import tempfile
        fname = tempfile.mktemp(suffix=".gdl")
        cg.gen_gdl(fname)
        ida_gdl.display_gdl(fname)
        return 1

    def update(self, ctx):
        return ida_kernwin.AST_ENABLE_FOR_WIDGET if \
            ctx.widget_type == ida_kernwin.BWN_PSEUDOCODE else \
            ida_kernwin.AST_DISABLE_FOR_WIDGET


class vds5_hooks_t(ida_hexrays.Hexrays_Hooks):
    def populating_popup(self, widget, handle, vu):
        ida_kernwin.attach_action_to_popup(vu.ct, None, ACTION_NAME)
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays show C graph (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin5 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            ida_kernwin.register_action(
                ida_kernwin.action_desc_t(
                    ACTION_NAME,
                    "Hex-Rays show C graph (IDAPython)",
                    display_graph_ah_t(),
                    ACTION_SHORTCUT))
            self.vds5_hooks = vds5_hooks_t()
            self.vds5_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds5_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: Titvt/GolangEnhancer
file: golang_enhancer.py
description: A plugin for IDA to enhance the ability of decompile Golang programs.
wanted_name: Golang Enhancer
comment: Golang Enhancer
created_at: 2021-12-20 15:30:07+00:00
pushed_at: 2021-12-20 16:53:03+00:00
url:
https://github.com/Titvt/GolangEnhancer/blob/32e86ea78fbe2ee9a6b2190e302792cbc3e87aa7/golang_enhancer.py

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

import string

import idc
from idaapi import *

desc = '''
Golang Enhancer V1.0
GitHub: https://github.com/Titvt/GolangEnhancer/
'''

logo = r'''
 _____           _                              _____           _                                           
|  __ \         | |                            |  ___|         | |                                          
| |  \/   ___   | |   __ _   _ __     __ _     | |__    _ __   | |__     __ _   _ __     ___    ___   _ __  
| | __   / _ \  | |  / _` | | '_ \   / _` |    |  __|  | '_ \  | '_ \   / _` | | '_ \   / __|  / _ \ | '__| 
| |_\ \ | (_) | | | | (_| | | | | | | (_| |    | |___  | | | | | | | | | (_| | | | | | | (__  |  __/ | |    
 \____/  \___/  |_|  \__,_| |_| |_|  \__, |    \____/  |_| |_| |_| |_|  \__,_| |_| |_|  \___|  \___| |_|    
                                      __/ |                                                                 
                                     |___/                                                                        
'''

str_execute_failed = '''
执行失败：该程序不是一个带符号的Golang程序或该Golang版本不支持（通常是因为太低）。
Execute failed: This program is not a Golang program with symbols or the Golang version is not supported (usually because too low).
'''

str_golang_version = '''
Golang版本：{0}
Golang Version: {0}
'''

str_result = '''
执行成功：{0}个函数已恢复，{1}个函数未恢复。
Execute succeeded: {0} functions have been recovered, {1} functions have not been recovered.
'''

argv = {
    0: '_QWORD __usercall {}<rax>();',
    1: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>);',
    2: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>);',
    3: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>, _QWORD arg3<rcx>);',
    4: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>, _QWORD arg3<rcx>, _QWORD arg4<rdi>);',
    5: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>, _QWORD arg3<rcx>, _QWORD arg4<rdi>, _QWORD arg5<rsi>);',
    6: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>, _QWORD arg3<rcx>, _QWORD arg4<rdi>, _QWORD arg5<rsi>, _QWORD arg6<r8>);',
    7: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>, _QWORD arg3<rcx>, _QWORD arg4<rdi>, _QWORD arg5<rsi>, _QWORD arg6<r8>, _QWORD arg7<r9>);',
    8: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>, _QWORD arg3<rcx>, _QWORD arg4<rdi>, _QWORD arg5<rsi>, _QWORD arg6<r8>, _QWORD arg7<r9>, _QWORD arg8<r10>);',
    9: '_QWORD __usercall {}<rax>(_QWORD arg1<rax>, _QWORD arg2<rbx>, _QWORD arg3<rcx>, _QWORD arg4<rdi>, _QWORD arg5<rsi>, _QWORD arg6<r8>, _QWORD arg7<r9>, _QWORD arg8<r10>, _QWORD arg9<r11>);',
}


def get_valid_name(name):
    ret = ''

    for i in name:
        if i in string.ascii_letters + string.digits:
            ret += i
        else:
            ret += '_'

    return ret


class GolangEnhancer(plugin_t):
    flags = PLUGIN_MOD + PLUGIN_DRAW + PLUGIN_UNL
    wanted_name = 'Golang Enhancer'
    wanted_hotkey = 'Ctrl+Alt+G'
    comment = 'Golang Enhancer'
    help = desc

    def init(self):
        print(self.help)
        return PLUGIN_OK

    def term(self):
        pass

    def run(self, arg):
        print(logo)

        go_ver = get_name_ea(BADADDR, 'runtime.buildVersion.str')

        if go_ver == BADADDR:
            print(str_execute_failed)
            return

        go_ver = get_strlit_contents(go_ver, -1, STRTYPE_C).decode()[2:]
        print(str_golang_version.format(go_ver))

        fn_def = 0
        fn_und = 0

        for i in range(get_func_qty()):
            func = getn_func(i)
            name = get_valid_name(get_func_name(func.start_ea))
            argc = 0

            insn = insn_t()
            decode_prev_insn(insn, func.end_ea)
            ea = insn.ea

            if idc.print_insn_mnem(ea) == 'jmp' and idc.print_operand(ea, 0) == name:
                fn_def += 1

                while argc < 9:
                    decode_prev_insn(insn, ea)
                    ea = insn.ea

                    if idc.print_insn_mnem(ea) == 'call' and idc.print_operand(ea, 0).startswith('runtime_morestack'):
                        break

                    if idc.print_insn_mnem(ea) == 'mov':
                        argc += 1
            else:
                fn_und += 1
                argc = 0

            idc.SetType(func.start_ea, argv[argc].format(name))

        print(str_result.format(fn_def, fn_und))


def PLUGIN_ENTRY():
    return GolangEnhancer()


repository: mahmoudimus/ida-themr
file: src/ida_themr_plugin.py
description: A pure python IDA Pro theme maker (and plugin!) for porting Visual Studio Code themes to IDA Pro.
wanted_name: IDA Theming Toolkit
comment: IDA Theming Toolkit
created_at: 2025-04-30 23:07:10+00:00
pushed_at: 2025-05-08 18:34:51+00:00
url:
https://github.com/mahmoudimus/ida-themr/blob/933f6476403612aa5b81a1f1fbc35c3cd22c01e6/src/ida_themr_plugin.py

content:
import configparser
import json
import os
import pathlib
from datetime import datetime
from textwrap import dedent

# For IDA integration
import idaapi
from PyQt5.QtCore import QEvent, Qt, QTimer
from PyQt5.QtGui import (
    QColor,
    QFont,
    QIcon,
    QKeySequence,
    QPalette,
    QPixmap,
    QTextCursor,
)
from PyQt5.QtWidgets import (
    QApplication,
    QCheckBox,
    QComboBox,
    QFileDialog,
    QHBoxLayout,
    QHeaderView,
    QLabel,
    QLineEdit,
    QMainWindow,
    QMenu,
    QMessageBox,
    QPushButton,
    QShortcut,
    QSizePolicy,
    QSplitter,
    QStyleFactory,
    QTabWidget,
    QTextEdit,
    QTreeWidget,
    QTreeWidgetItem,
    QVBoxLayout,
    QWidget,
)

# Configuration paths
USER_IDADIR = pathlib.Path(idaapi.get_user_idadir())
USER_CFGDIR = USER_IDADIR / "cfg"
QTAPP_STYLE_CFG = USER_CFGDIR / "qtappstyle.cfg"
PLUGIN_NAME = "IDAThemr"


def load_style():
    """Read the previously selected style and dark mode setting from the config file."""
    config = configparser.ConfigParser()
    try:
        config.read(QTAPP_STYLE_CFG)
    except Exception:
        pass

    style = config.get("default", "qt_style", fallback="Fusion")
    dark_mode = config.getboolean("default", "dark_mode", fallback=False)
    return style, dark_mode


def save_style(style: str, dark_mode: bool = False):
    """Save the chosen style and dark mode setting to the config file."""
    config = configparser.ConfigParser()
    config["default"] = {
        "qt_style": style,
        "dark_mode": str(dark_mode).lower(),
    }
    USER_CFGDIR.mkdir(parents=True, exist_ok=True)
    with open(QTAPP_STYLE_CFG, "w") as configfile:
        config.write(configfile)


def toggle_dark_theme(on: bool, property_name: str = "os-dark-theme"):
    """Toggle the OS theme for the IDA GUI."""
    app = QApplication.instance()
    # Flip the dynamic property on every widget
    for w in app.allWidgets():
        w.setProperty(property_name, on)
        w.style().unpolish(w)
        w.style().polish(w)
        w.update()


class StyleSheetWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)

        # Force widget font across the board, and monospace in the editor
        self.setStyleSheet(
            """
            QWidget {
                font-family: -apple-system, system-ui, BlinkMacSystemFont,
                             "Segoe UI", Roboto, "Helvetica Neue",
                             Arial, sans-serif !important;
            }
            QTextEdit {
                font-family: "Maple Mono NF Light", "Cascadia Code", "Consolas" !important;
                font-size: 10em !important;
            }
        """
        )

        self.tape = []
        self.tape_pos = -1
        self.current_file = None

        # Create main layout
        layout = QVBoxLayout(self)

        # Add file path display
        self.file_path_label = QLabel("Current stylesheet: <none>", self)
        layout.addWidget(self.file_path_label)

        # Create splitter for raw file and editor
        self.splitter = QSplitter(Qt.Vertical)

        # Raw file view
        self.raw_file_view = QTextEdit(self)
        self.raw_file_view.setReadOnly(True)
        self.raw_file_view.setAcceptRichText(False)
        self.raw_file_view.setPlaceholderText(
            "Raw stylesheet file content will appear here"
        )

        # Editor section with search bar
        self.editor_widget = QWidget(self)
        editor_layout = QVBoxLayout(self.editor_widget)

        self.search_bar = QLineEdit(self)
        self.search_bar.setPlaceholderText("Search in stylesheet (Ctrl+F)")
        self.search_bar.textChanged.connect(self.onSearchTextChanged)

        self.style_text_edit = QTextEdit(self)
        self.style_text_edit.setAcceptRichText(False)
        self.style_text_edit.textChanged.connect(self.onStyleTextChanged)

        editor_layout.addWidget(self.search_bar)
        editor_layout.addWidget(self.style_text_edit)

        # Add to splitter
        self.splitter.addWidget(self.raw_file_view)
        self.splitter.addWidget(self.editor_widget)
        self.splitter.setSizes([200, 400])  # Initial sizes

        layout.addWidget(self.splitter)

        # Add buttons
        button_layout = QHBoxLayout()

        self.load_button = QPushButton("Load File", self)
        self.load_button.clicked.connect(self.onLoadFile)

        self.apply_button = QPushButton("Apply", self)
        self.apply_button.clicked.connect(self.onApplyButton)
        self.apply_button.setEnabled(False)

        self.save_button = QPushButton("Save", self)
        self.save_button.clicked.connect(self.onSaveFile)

        button_layout.addWidget(self.load_button)
        button_layout.addWidget(self.apply_button)
        button_layout.addWidget(self.save_button)

        layout.addLayout(button_layout)

        # Shortcuts
        QShortcut(QKeySequence(Qt.Key_F3), self).activated.connect(self.onNextSearchHit)
        QShortcut(QKeySequence(Qt.CTRL + Qt.Key_F), self).activated.connect(
            self.onFocusSearchBar
        )
        QShortcut(QKeySequence(Qt.CTRL + Qt.Key_S), self).activated.connect(
            self.applyStyleSheet
        )
        QShortcut(QKeySequence(Qt.CTRL + Qt.ALT + Qt.Key_Z), self).activated.connect(
            self.onUndo
        )
        QShortcut(QKeySequence(Qt.CTRL + Qt.ALT + Qt.Key_Y), self).activated.connect(
            self.onRedo
        )
        QShortcut(QKeySequence(Qt.Key_F1), self).activated.connect(self.onHelp)

        # Load current stylesheet
        self.loadStyleSheet()

    def onLoadFile(self):
        """Load a stylesheet from a file."""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Open Stylesheet", "", "Stylesheet Files (*.qss *.css);;All Files (*)"
        )

        if file_path:
            try:
                with open(file_path, "r") as f:
                    content = f.read()

                self.current_file = file_path
                self.file_path_label.setText(f"Current stylesheet: {file_path}")
                self.raw_file_view.setPlainText(content)
                self.style_text_edit.setPlainText(content)

                # Add to tape
                self.tape = [content]
                self.tape_pos = 0
                self.apply_button.setEnabled(False)
            except Exception as e:
                QMessageBox.critical(self, "Error", f"Failed to load file: {str(e)}")

    def onSaveFile(self):
        """Save the current stylesheet to a file."""
        if self.current_file:
            default_path = self.current_file
        else:
            default_path = ""

        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "Save Stylesheet",
            default_path,
            "Stylesheet Files (*.qss *.css);;All Files (*)",
        )

        if file_path:
            try:
                with open(file_path, "w") as f:
                    f.write(self.style_text_edit.toPlainText())

                self.current_file = file_path
                self.file_path_label.setText(f"Current stylesheet: {file_path}")
                self.raw_file_view.setPlainText(self.style_text_edit.toPlainText())
                QMessageBox.information(
                    self, "Success", "Stylesheet saved successfully."
                )
            except Exception as e:
                QMessageBox.critical(self, "Error", f"Failed to save file: {str(e)}")

    def onUndo(self, checked=False):
        if self.tape_pos <= 0:
            return
        self.tape_pos -= 1
        self.style_text_edit.setPlainText(self.tape[self.tape_pos])
        self.applyStyleSheet(stateless=True)

    def onRedo(self, checked=False):
        if self.tape_pos >= len(self.tape) - 1:
            return
        self.tape_pos += 1
        self.style_text_edit.setPlainText(self.tape[self.tape_pos])
        self.applyStyleSheet(stateless=True)

    def onHelp(self):
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle("Help")
        msg_box.setText("Available shortcuts:")
        msg_box.setInformativeText(
            dedent(
                """\
            F1: show help dialog
            Ctrl+S: apply current changes
            Ctrl+F: go to search bar
            F3: go to next search hit
            Ctrl+Alt+Z: revert to last applied style sheet
            Ctrl+Alt+Y: redo last reverted style sheet
        """
            )
        )
        msg_box.exec_()

    def onSearchTextChanged(self, text):
        if not self.style_text_edit.find(text):
            self.search_bar.setStyleSheet("color: red;")
            self.style_text_edit.moveCursor(QTextCursor.Start)
        else:
            self.search_bar.setStyleSheet("color: green;")

    def onNextSearchHit(self):
        search = self.search_bar.text()
        if not self.style_text_edit.find(search):
            self.style_text_edit.moveCursor(QTextCursor.Start)
            self.style_text_edit.find(search)

    def onFocusSearchBar(self):
        self.search_bar.setFocus()

    def onStyleTextChanged(self):
        self.apply_button.setEnabled(True)

    def onApplyButton(self, checked=False):
        self.applyStyleSheet()
        # Update the raw view with the current content
        self.raw_file_view.setPlainText(self.style_text_edit.toPlainText())

    def loadStyleSheet(self):
        ss = QApplication.instance().styleSheet()
        self.tape = [ss]
        self.tape_pos = 0
        self.style_text_edit.setPlainText(ss)
        self.raw_file_view.setPlainText(ss)
        self.apply_button.setEnabled(False)
        self.file_path_label.setText("Current stylesheet: <from application>")

    def applyStyleSheet(self, stateless=False):
        ss = self.style_text_edit.toPlainText()
        QApplication.instance().setStyleSheet(ss)
        if not stateless:
            # Discard any "future" states
            self.tape = self.tape[: self.tape_pos + 1]
            self.tape.append(ss)
            self.tape_pos += 1
        self.apply_button.setEnabled(False)


class ObjectInspector(QWidget):
    """
    Qt widget inspector.
    Automatically inspects widgets on hover without stealing focus.
    Presents inspection data in a tree widget with collapsible categories.
    Includes deep metadata: style engine, cursor, autofill, focus, dynamic properties, palette, actions.
    """

    def __init__(self, parent=None):
        super().__init__(parent)

        # Store reference to main window for accessing the status bar
        self.main_window = self.get_main_window()

        # Always on top as a tool window, but does not grab focus
        self.setWindowFlags(self.windowFlags() | Qt.Tool)
        self.setAttribute(Qt.WA_ShowWithoutActivating)

        # Flag to track if updates are suspended
        self.updates_suspended = False
        # Flag to track if updates are manually suspended (for context menu)
        self.menu_suspended = False

        # Flag for expand on updates
        self.expand_on_updates = False

        # For status message handling
        self.status_timer = QTimer(self)
        self.status_timer.setSingleShot(True)
        self.status_timer.timeout.connect(self.restore_default_status)

        # Tracking categories for tree widget
        self.categories = {}

        # Store context menu data
        self.context_menu_data = None

        # Install event filter for hover events
        app = QApplication.instance()
        app.installEventFilter(self)

        # Shortcut F7 to select parent
        sch = QShortcut(QKeySequence(Qt.Key_F7), self)
        sch.setContext(Qt.ApplicationShortcut)
        sch.activated.connect(self.select_parent)

        self.selected_widget = None

        self._create_ui()
        self._set_tree_style()

        # Enable key tracking
        self.setFocusPolicy(Qt.StrongFocus)

        # Show default status
        self.set_status_message("Hold Ctrl OR Shift to suspend updates")

    def get_main_window(self):
        """Get reference to the main window."""
        parent = self.parent()
        while parent is not None:
            if isinstance(parent, QMainWindow):
                return parent
            parent = parent.parent()
        return None

    def _create_ui(self):
        # Top buttons
        self.btn_select_parent = QPushButton("Select parent", self)
        self.btn_select_parent.released.connect(self.select_parent)

        self.btn_copy = QPushButton("Copy All", self)
        self.btn_copy.released.connect(self.copy_to_clipboard)

        self.btn_export = QPushButton("Export...", self)
        self.btn_export.released.connect(self.export_data)

        self.btn_expand_all = QPushButton("Expand All", self)
        self.btn_expand_all.released.connect(self.expand_all_categories)

        self.btn_collapse_all = QPushButton("Collapse All", self)
        self.btn_collapse_all.released.connect(self.collapse_all_categories)

        self.chk_expand_on_updates = QCheckBox("Expand on Updates", self)
        self.chk_expand_on_updates.stateChanged.connect(self.toggle_expand_on_updates)

        # Help label
        self.lbl_help = QLabel(
            "Hover over any widget to inspect it; press F7 to select its parent. Right-click for context menu.",
            self,
        )

        spacer = QWidget(self)
        spacer.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred)

        # Top controls layout
        top_layout = QHBoxLayout()
        top_layout.setContentsMargins(2, 2, 2, 2)
        top_layout.setSpacing(4)
        top_layout.addWidget(self.btn_select_parent)
        top_layout.addWidget(self.btn_copy)
        top_layout.addWidget(self.btn_export)
        top_layout.addWidget(self.btn_expand_all)
        top_layout.addWidget(self.btn_collapse_all)
        top_layout.addWidget(self.chk_expand_on_updates)
        top_layout.addWidget(spacer)

        # Second row layout
        second_row_layout = QHBoxLayout()
        second_row_layout.setContentsMargins(2, 2, 2, 2)
        second_row_layout.setSpacing(4)
        second_row_layout.addWidget(self.lbl_help)
        second_row_layout.addStretch()

        # Inspection tree
        self.tree_inspection = QTreeWidget(self)
        self.tree_inspection.setColumnCount(2)
        self.tree_inspection.setHeaderLabels(["Property", "Value"])
        self.tree_inspection.header().setSectionResizeMode(
            0, QHeaderView.ResizeToContents
        )
        self.tree_inspection.header().setSectionResizeMode(1, QHeaderView.Stretch)
        self.tree_inspection.setEditTriggers(QTreeWidget.NoEditTriggers)
        self.tree_inspection.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.tree_inspection.setContextMenuPolicy(Qt.CustomContextMenu)
        self.tree_inspection.customContextMenuRequested.connect(self.show_context_menu)

        # Main layout
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(4, 4, 4, 4)
        main_layout.setSpacing(6)
        main_layout.addLayout(top_layout)
        main_layout.addLayout(second_row_layout)
        main_layout.addWidget(self.tree_inspection)
        self.setLayout(main_layout)

    def set_status_message(self, message, timeout=0):
        """
        Set a status message with optional timeout.
        If timeout > 0, the message will be shown for that many milliseconds
        and then revert to the default message.
        """
        # Cancel any pending timer
        self.status_timer.stop()

        # Set the new message on the main window's status bar if available
        if self.main_window and hasattr(self.main_window, "statusBar"):
            self.main_window.statusBar().showMessage(message)

        # If timeout specified, schedule restoration of default message
        if timeout > 0:
            self.status_timer.start(timeout)

    def restore_default_status(self):
        """Restore the default status message based on the current state."""
        if self.updates_suspended:
            self.set_status_message("Updates suspended")
        else:
            self.set_status_message("Hold Ctrl OR Shift to suspend updates")

    def toggle_expand_on_updates(self, state):
        """Toggle whether categories should be expanded on inspection updates."""
        self.expand_on_updates = state == Qt.Checked

    def expand_all_categories(self):
        """Expand all categories in the tree."""
        for category_idx in range(self.tree_inspection.topLevelItemCount()):
            category = self.tree_inspection.topLevelItem(category_idx)
            category.setExpanded(True)

    def collapse_all_categories(self):
        """Collapse all categories in the tree."""
        for category_idx in range(self.tree_inspection.topLevelItemCount()):
            category = self.tree_inspection.topLevelItem(category_idx)
            category.setExpanded(False)

    def show_context_menu(self, position):
        """Show a context menu for the tree item at the given position."""
        item = self.tree_inspection.itemAt(position)
        if not item:
            return

        # Suspend updates while the context menu is active
        self.menu_suspended = True
        self.updates_suspended = True
        self.set_status_message("Updates suspended for menu operation")

        menu = QMenu(self)
        # Connect to aboutToHide to resume updates if menu is cancelled
        menu.aboutToHide.connect(self.resume_updates_after_menu)

        # Store item data before showing menu
        is_category = item.parent() is None

        if is_category:
            # For category, store all property data
            category_name = item.text(0)
            properties = []

            for prop_idx in range(item.childCount()):
                prop_item = item.child(prop_idx)
                properties.append((prop_item.text(0), prop_item.text(1)))

            self.context_menu_data = {
                "type": "category",
                "name": category_name,
                "properties": properties,
                "item": item,  # Keep reference for expand/collapse actions
            }

            # Category item
            copy_action = menu.addAction("Copy Category")
            copy_action.triggered.connect(self.copy_context_category)

            expand_action = menu.addAction("Expand")
            expand_action.triggered.connect(self.expand_context_category)

            collapse_action = menu.addAction("Collapse")
            collapse_action.triggered.connect(self.collapse_context_category)
        else:
            # For property, just store the property data
            prop_name = item.text(0)
            prop_value = item.text(1)

            self.context_menu_data = {
                "type": "property",
                "name": prop_name,
                "value": prop_value,
            }

            # Property item
            copy_action = menu.addAction("Copy Property")
            copy_action.triggered.connect(self.copy_context_property)

        menu.exec_(self.tree_inspection.viewport().mapToGlobal(position))

    def resume_updates_after_menu(self):
        """Resume updates after the context menu is hidden."""
        # Only resume if suspension was due to menu (not Ctrl/Shift keys)
        if self.menu_suspended:
            self.menu_suspended = False
            # Check if Ctrl/Shift are still pressed
            modifiers = QApplication.keyboardModifiers()
            if not (modifiers & Qt.ControlModifier or modifiers & Qt.ShiftModifier):
                self.updates_suspended = False
                self.restore_default_status()

    def copy_context_category(self):
        """Copy all properties in the context menu category to the clipboard."""
        if not self.context_menu_data or self.context_menu_data["type"] != "category":
            return

        category_name = self.context_menu_data["name"]
        properties = self.context_menu_data["properties"]

        lines = [f"{category_name}:"]
        for prop_name, prop_value in properties:
            lines.append(f"  {prop_name}: {prop_value}")

        QApplication.instance().clipboard().setText("\n".join(lines))
        self.set_status_message(f"Category '{category_name}' copied to clipboard", 3000)

        # Resume updates after copy
        self.resume_updates_after_menu()

    def copy_context_property(self):
        """Copy a single property from the context menu to the clipboard."""
        if not self.context_menu_data or self.context_menu_data["type"] != "property":
            return

        prop_name = self.context_menu_data["name"]
        prop_value = self.context_menu_data["value"]

        text = f"{prop_name}: {prop_value}"
        QApplication.instance().clipboard().setText(text)
        self.set_status_message(f"Property '{prop_name}' copied to clipboard", 3000)

        # Resume updates after copy
        self.resume_updates_after_menu()

    def expand_context_category(self):
        """Expand the category from the context menu."""
        if not self.context_menu_data or self.context_menu_data["type"] != "category":
            return

        item = self.context_menu_data["item"]
        if item and self.tree_inspection.indexOfTopLevelItem(item) >= 0:
            item.setExpanded(True)

        # Resume updates after expand
        self.resume_updates_after_menu()

    def collapse_context_category(self):
        """Collapse the category from the context menu."""
        if not self.context_menu_data or self.context_menu_data["type"] != "category":
            return

        item = self.context_menu_data["item"]
        if item and self.tree_inspection.indexOfTopLevelItem(item) >= 0:
            item.setExpanded(False)

        # Resume updates after collapse
        self.resume_updates_after_menu()

    def export_data(self):
        """Export the inspection data to a file."""
        if not self.selected_widget:
            self.set_status_message("No widget selected to export", 3000)
            return

        # Suspend updates during export
        was_suspended = self.updates_suspended
        self.updates_suspended = True
        self.set_status_message("Updates suspended for export operation")

        # Prepare data for export
        data = self._get_data_as_dict()

        # Ask user for file location and format
        file_path, selected_filter = QFileDialog.getSaveFileName(
            self,
            "Save Inspection Data",
            f"widget_inspection_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            "JSON Files (*.json);;INI Files (*.ini);;All Files (*)",
        )

        if not file_path:
            # Resume updates if export was cancelled
            if not was_suspended:
                self.updates_suspended = False
                self.restore_default_status()
            return  # User canceled

        try:
            if file_path.endswith(".json") or "JSON" in selected_filter:
                self._export_as_json(file_path, data)
            elif file_path.endswith(".ini") or "INI" in selected_filter:
                self._export_as_ini(file_path, data)
            else:
                # Default to JSON if no extension specified
                if not "." in file_path:
                    file_path += ".json"
                self._export_as_json(file_path, data)

            self.set_status_message(f"Exported data to {file_path}", 3000)
        except Exception as e:
            self.set_status_message(f"Error exporting data: {str(e)}", 5000)
        finally:
            # Resume updates if they weren't suspended before
            if not was_suspended:
                self.updates_suspended = False
                # Status will be restored by the timer set in set_status_message

    def _get_data_as_dict(self):
        """Get the inspection data as a nested dictionary."""
        data = {}

        # Iterate through all categories
        for category_idx in range(self.tree_inspection.topLevelItemCount()):
            category = self.tree_inspection.topLevelItem(category_idx)
            category_name = category.text(0)
            category_data = {}

            # Iterate through properties in this category
            for prop_idx in range(category.childCount()):
                prop_item = category.child(prop_idx)
                prop_name = prop_item.text(0)
                prop_value = prop_item.text(1)
                category_data[prop_name] = prop_value

            data[category_name] = category_data

        return data

    def _export_as_json(self, file_path, data):
        """Export data as JSON."""
        with open(file_path, "w") as f:
            json.dump(data, f, indent=2)

    def _export_as_ini(self, file_path, data):
        """Export data as INI file."""
        config = configparser.ConfigParser()

        for category, props in data.items():
            # Replace any characters not allowed in INI section names
            safe_category = category.replace("[", "(").replace("]", ")")
            config[safe_category] = props

        with open(file_path, "w") as f:
            config.write(f)

    def _set_tree_style(self):
        font = QFont("Monospace")
        font.setStyleHint(QFont.TypeWriter)
        self.tree_inspection.setFont(font)

    def _update_suspended_state(self):
        """Check and update the suspended state based on modifier keys."""
        # Don't change state if menu is active
        if self.menu_suspended:
            return True

        modifiers = QApplication.keyboardModifiers()
        was_suspended = self.updates_suspended
        self.updates_suspended = bool(
            modifiers & Qt.ControlModifier or modifiers & Qt.ShiftModifier
        )

        if self.updates_suspended != was_suspended:
            if self.updates_suspended:
                self.set_status_message("Updates suspended")
            else:
                self.set_status_message("Hold Ctrl OR Shift to suspend updates")

        return self.updates_suspended

    def eventFilter(self, obj, event):
        # Check if modifiers changed on any event
        if event.type() in (QEvent.KeyPress, QEvent.KeyRelease):
            self._update_suspended_state()

        # Handle Enter events for widget inspection
        elif event.type() == QEvent.Enter and isinstance(obj, QWidget):
            if obj is self or self.isAncestorOf(obj):
                return super().eventFilter(obj, event)

            # Check keyboard modifiers
            if not self._update_suspended_state() and obj is not self:
                self._inspect_widget(obj)

        return super().eventFilter(obj, event)

    def select_parent(self):
        if not self.selected_widget:
            return
        parent = self.selected_widget.parent()
        if parent and parent.inherits("QWidget"):
            self._inspect_widget(parent)

    def _get_specific_css_selector(self, widget):
        """Generate the most specific CSS selector for a widget."""
        widget_type = widget.metaObject().className()
        object_name = widget.objectName()

        if object_name:
            return f'{widget_type}[objectName="{object_name}"]'
        else:
            return widget_type

    def _get_css_selector_with_parent(self, widget):
        """Generate a CSS selector that includes the immediate parent."""
        if not widget.parent():
            return self._get_specific_css_selector(widget)

        parent_selector = self._get_specific_css_selector(widget.parent())
        widget_selector = self._get_specific_css_selector(widget)
        return f"{parent_selector} > {widget_selector}"

    def _get_full_hierarchy_css_selector(self, widget):
        """Generate a CSS selector with the full widget hierarchy from root to widget."""
        selectors = []
        current = widget

        while current:
            selectors.insert(0, self._get_specific_css_selector(current))
            current = current.parent()

        return " > ".join(selectors)

    def _add_css_selector_hierarchy(self, widget, css_selectors):
        """Add CSS selectors for the widget and all its parents to the css_selectors dict."""
        # Add the most specific selector for this widget
        css_selectors["CSS Selector (Specific)"] = self._get_specific_css_selector(
            widget
        )

        # Build a list of all ancestors
        ancestors = []
        current = widget
        while current.parent():
            ancestors.insert(0, current.parent())
            current = current.parent()

        # Add a selector for each level of the hierarchy
        if ancestors:
            # First, add the immediate parent relationship
            css_selectors["CSS Selector (With Parent)"] = (
                self._get_css_selector_with_parent(widget)
            )

            # Then add the full hierarchy selector
            css_selectors["CSS Selector (Full Hierarchy)"] = (
                self._get_full_hierarchy_css_selector(widget)
            )

            # Add intermediate levels if there are more than one parent
            if len(ancestors) > 1:
                # Start building from the root
                path_selectors = []
                for i, ancestor in enumerate(ancestors):
                    path_selectors.append(self._get_specific_css_selector(ancestor))
                    # Skip the immediate parent (already covered) and full hierarchy (already covered)
                    if 0 < i < len(ancestors) - 1:
                        path_selectors_copy = path_selectors.copy()
                        path_selectors_copy.append(
                            self._get_specific_css_selector(widget)
                        )
                        css_selectors[f"CSS Selector (With {i+1} Parents)"] = (
                            " > ".join(path_selectors_copy)
                        )

    def _add_category(self, name, collapsed=True):
        """Add a category to the tree widget."""
        category = QTreeWidgetItem(self.tree_inspection)
        category.setText(0, name)
        category.setExpanded(not collapsed and not self.expand_on_updates)
        self.categories[name] = category
        return category

    def _add_property(self, category, prop, val, icon=None):
        """Add a property to a category in the tree widget."""
        item = QTreeWidgetItem(category)
        item.setText(0, prop)
        item.setText(1, val)
        if icon:
            item.setIcon(1, icon)
        return item

    def _inspect_widget(self, widget):
        # Disconnect previous destroyed signal
        if self.selected_widget is not None:
            try:
                self.selected_widget.destroyed.disconnect(self.on_widget_destroyed)
            except TypeError:
                pass

        self.selected_widget = widget

        # Clear existing tree and categories
        self.tree_inspection.clear()
        self.categories = {}

        # Clear context menu data
        self.context_menu_data = None

        if not widget:
            self._add_property(self.tree_inspection, "<no widget>", "")
            return

        # Create categories with appropriate collapsed state (might be overridden by expand_on_updates)
        self._add_category("Screen Information", collapsed=True)
        self._add_category("Basic Information", collapsed=False)
        self._add_category("CSS Selectors", collapsed=True)
        self._add_category("Parent Information", collapsed=False)
        self._add_category("Widget State", collapsed=False)
        self._add_category("Style & Appearance", collapsed=False)
        self._add_category("Layout", collapsed=False)
        self._add_category("Properties", collapsed=False)
        self._add_category("Palette Colors", collapsed=True)
        self._add_category("Interactivity", collapsed=False)
        self._add_category("Miscellaneous", collapsed=False)

        # Screen info
        screen = QApplication.instance().primaryScreen()
        res = screen.size()
        phy = screen.physicalSize()
        dpi = screen.logicalDotsPerInch()
        self._add_property(
            self.categories["Screen Information"],
            "Screen Resolution",
            f"{res.width()}×{res.height()} px",
        )
        self._add_property(
            self.categories["Screen Information"],
            "Physical Size",
            f"{phy.width()}×{phy.height()} mm",
        )
        self._add_property(self.categories["Screen Information"], "DPI", f"{dpi:.1f}")

        # Basic info
        self._add_property(
            self.categories["Basic Information"],
            "Type",
            widget.metaObject().className(),
        )
        self._add_property(
            self.categories["Basic Information"],
            "Name",
            widget.objectName() or "<none>",
        )
        self._add_property(
            self.categories["Basic Information"],
            "Children Count",
            str(len(widget.children())),
        )

        # CSS selectors
        css_selectors = {}
        self._add_css_selector_hierarchy(widget, css_selectors)
        for name, value in css_selectors.items():
            self._add_property(self.categories["CSS Selectors"], name, value)

        # Parent info
        parent = widget.parent()
        self._add_property(
            self.categories["Parent Information"],
            "Parent Type",
            parent.metaObject().className() if parent else "<none>",
        )
        self._add_property(
            self.categories["Parent Information"],
            "Parent Name",
            parent.objectName() if parent and parent.objectName() else "<none>",
        )

        # Widget state
        self._add_property(
            self.categories["Widget State"], "Enabled", str(widget.isEnabled())
        )
        self._add_property(
            self.categories["Widget State"], "Visible", str(widget.isVisible())
        )
        self._add_property(
            self.categories["Widget State"], "Under Mouse", str(widget.underMouse())
        )
        self._add_property(
            self.categories["Widget State"], "Has Focus", str(widget.hasFocus())
        )

        # Style & Appearance
        self._add_property(
            self.categories["Style & Appearance"],
            "Style Engine",
            widget.style().metaObject().className(),
        )
        self._add_property(
            self.categories["Style & Appearance"],
            "Cursor Shape",
            str(widget.cursor().shape()),
        )
        self._add_property(
            self.categories["Style & Appearance"],
            "AutoFillBackground",
            str(widget.autoFillBackground()),
        )
        self._add_property(
            self.categories["Style & Appearance"],
            "Focus Policy",
            str(widget.focusPolicy()),
        )

        # Font
        font = widget.font()
        self._add_property(
            self.categories["Style & Appearance"],
            "Font",
            f"family={font.family()}, size={font.pointSize()}, weight={font.weight()}, italic={font.italic()}",
        )

        # Layout
        geom = widget.geometry()
        self._add_property(
            self.categories["Layout"],
            "Geometry",
            f"x={geom.x()}, y={geom.y()}, width={geom.width()}, height={geom.height()}",
        )

        # Layout margins
        layout = widget.layout()
        if layout:
            m = layout.contentsMargins()
            self._add_property(
                self.categories["Layout"],
                "Layout Margins",
                f"left={m.left()}, top={m.top()}, right={m.right()}, bottom={m.bottom()}",
            )
        else:
            self._add_property(
                self.categories["Layout"], "Layout Margins", "<no layout>"
            )

        # Style sheet
        self._add_property(
            self.categories["Layout"], "StyleSheet", widget.styleSheet() or "<none>"
        )

        # Dynamic properties
        dp = [name.data().decode() for name in widget.dynamicPropertyNames()]
        self._add_property(
            self.categories["Properties"],
            "Dynamic Properties",
            ", ".join(dp) or "<none>",
        )

        # Palette colors - dynamic version
        pal = widget.palette()
        roles = {}
        for i in range(QPalette.NColorRoles):
            try:
                col = pal.color(i)
                if not col.isValid():
                    continue

                # find the matching attribute name on QPalette
                name = next(
                    (
                        name
                        for name, val in QPalette.__dict__.items()
                        if isinstance(val, int) and val == i
                    ),
                    f"Role_{i}",  # Fallback name if not found
                )
                # store the integer role (not the color string)
                roles[name] = i
            except Exception:
                continue  # Skip any role that causes an error

        # now pull each color back out and append
        for name in sorted(roles):
            role = roles[name]
            color = pal.color(role).name()
            pix = QPixmap(16, 16)
            pix.fill(QColor(color))
            icon = QIcon(pix)
            self._add_property(self.categories["Palette Colors"], name, color, icon)

        # Shortcuts
        shortcuts = widget.findChildren(QShortcut)
        sc_list = [sc.key().toString() for sc in shortcuts]
        self._add_property(
            self.categories["Interactivity"],
            "Shortcuts",
            ", ".join(sc_list) or "<none>",
        )

        # Actions
        actions = widget.actions()
        act_list = [f"{act.text()} ({act.shortcut().toString()})" for act in actions]
        self._add_property(
            self.categories["Interactivity"], "Actions", "; ".join(act_list) or "<none>"
        )

        # Paint engine
        engine = widget.paintEngine()
        self._add_property(
            self.categories["Miscellaneous"],
            "Paint Engine",
            engine.__class__.__name__ if engine else "<none>",
        )

        # If expand_on_updates is enabled, expand all categories
        if self.expand_on_updates:
            self.expand_all_categories()

        # Connect destroyed signal
        widget.destroyed.connect(self.on_widget_destroyed)

    def copy_to_clipboard(self):
        if not self.selected_widget:
            return

        # Suspend updates during copy
        was_suspended = self.updates_suspended
        self.updates_suspended = True
        self.set_status_message("Updates suspended for copy operation")

        lines = []

        # Iterate through all categories
        for category_idx in range(self.tree_inspection.topLevelItemCount()):
            category = self.tree_inspection.topLevelItem(category_idx)
            category_name = category.text(0)
            lines.append(f"{category_name}:")

            # Iterate through properties in this category
            for prop_idx in range(category.childCount()):
                prop_item = category.child(prop_idx)
                prop_name = prop_item.text(0)
                prop_value = prop_item.text(1)
                lines.append(f"  {prop_name}: {prop_value}")

        QApplication.instance().clipboard().setText("\n".join(lines))
        self.set_status_message("All data copied to clipboard", 3000)

        # Resume updates if they weren't suspended before
        if not was_suspended:
            self.updates_suspended = False
            # Status will be restored by timer set in set_status_message

    def on_widget_destroyed(self, _):
        self.selected_widget = None


class QtStyleSelector(QWidget):
    """Widget for selecting Qt style and toggling dark mode."""

    def __init__(self, parent=None):
        super(QtStyleSelector, self).__init__(parent)
        self.setup_ui()
        self.load_saved_settings()

    def setup_ui(self):
        layout = QVBoxLayout(self)

        # Style selector group
        style_group = QWidget(self)
        style_layout = QHBoxLayout(style_group)

        style_label = QLabel("Qt Style:", self)
        self.style_combo = QComboBox(self)
        self.style_combo.addItems(QStyleFactory.keys())

        style_layout.addWidget(style_label)
        style_layout.addWidget(self.style_combo)

        # Dark mode toggle
        self.dark_mode_checkbox = QCheckBox("Dark Mode", self)

        # Buttons
        button_layout = QHBoxLayout()
        self.apply_button = QPushButton("Apply", self)
        self.apply_button.clicked.connect(self.apply_settings)

        self.save_button = QPushButton("Save Settings", self)
        self.save_button.clicked.connect(self.save_settings)

        button_layout.addWidget(self.apply_button)
        button_layout.addWidget(self.save_button)

        # Add to main layout
        layout.addWidget(style_group)
        layout.addWidget(self.dark_mode_checkbox)
        layout.addLayout(button_layout)
        layout.addStretch()

    def load_saved_settings(self):
        try:
            style, dark_mode = load_style()

            # Set style combo box
            index = self.style_combo.findText(style)
            if index >= 0:
                self.style_combo.setCurrentIndex(index)

            # Set dark mode checkbox
            self.dark_mode_checkbox.setChecked(dark_mode)

        except Exception as e:
            QMessageBox.warning(self, "Warning", f"Failed to load settings: {str(e)}")

    def apply_settings(self):
        style = self.style_combo.currentText()
        dark_mode = self.dark_mode_checkbox.isChecked()

        # Apply style
        QApplication.instance().setStyle(QStyleFactory.create(style))

        # Apply dark mode
        toggle_dark_theme(dark_mode)

        QMessageBox.information(self, "Success", "Settings applied successfully.")

    def save_settings(self):
        try:
            style = self.style_combo.currentText()
            dark_mode = self.dark_mode_checkbox.isChecked()

            save_style(style, dark_mode)
            QMessageBox.information(self, "Success", "Settings saved successfully.")

        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save settings: {str(e)}")


class IDAThemingToolkit(QMainWindow):
    """
    Main application window for the IDA Theming Toolkit.
    Provides access to widget inspector, style sheet editor, and theme settings.
    """

    def __init__(self, parent=None):
        super(IDAThemingToolkit, self).__init__(parent)
        self.setWindowTitle("IDA Theming Toolkit")
        self.resize(1200, 800)

        # Create central widget with tabs
        self.central_widget = QTabWidget(self)
        self.setCentralWidget(self.central_widget)

        # Create status bar first so the object inspector can access it
        self.statusBar().showMessage("Ready")

        # Create tabs
        self.object_inspector = ObjectInspector(
            self
        )  # Pass self as parent so it can access the statusBar
        self.stylesheet_editor = StyleSheetWidget(self)
        self.style_selector = QtStyleSelector(self)

        # Add tabs
        self.central_widget.addTab(self.object_inspector, "Widget Inspector")
        self.central_widget.addTab(self.stylesheet_editor, "Stylesheet Editor")
        self.central_widget.addTab(self.style_selector, "Qt Style & Theme")

        # Set up tab change handler to update status bar
        self.central_widget.currentChanged.connect(self.on_tab_changed)

        # Load initial settings
        self.load_initial_settings()

    def on_tab_changed(self, index):
        """Update the status bar when tabs are changed."""
        if index == 0:  # Widget Inspector tab
            self.statusBar().showMessage("Hold Ctrl OR Shift to suspend updates")
        elif index == 1:  # Stylesheet Editor tab
            self.statusBar().showMessage("Ready - Stylesheet Editor")
        elif index == 2:  # Qt Style & Theme tab
            self.statusBar().showMessage("Ready - Style & Theme Settings")

    @classmethod
    def load_initial_settings(cls):
        try:
            style, dark_mode = load_style()

            # Apply saved settings if they exist
            if style:
                QApplication.instance().setStyle(QStyleFactory.create(style))

            if dark_mode:
                toggle_dark_theme(dark_mode)

            QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
            QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
            QApplication.setAttribute(
                Qt.HighDpiScaleFactorRoundingPolicy.PassThrough, True
            )
            idaapi.msg(
                f"{PLUGIN_NAME}: Applied style: {style}, Dark mode: {dark_mode}{os.linesep}"
            )
        except Exception as e:
            if cls is IDAThemingToolkit:
                idaapi.msg("Failed to load settings: {str(e)}\n")
            else:
                cls.statusBar().showMessage(f"Failed to load settings: {str(e)}")


class IDARootPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = "IDA Theming Toolkit"
    help = "Tools for customizing IDA's appearance"
    wanted_name = "IDA Theming Toolkit"
    wanted_hotkey = "Alt-T"

    def init(self):
        IDAThemingToolkit.load_initial_settings()
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        # Launch the toolkit
        toolkit = IDAThemingToolkit()
        toolkit.show()

    def term(self):
        pass


# Entry point for IDA Plugin
def PLUGIN_ENTRY():
    return IDARootPlugin()


if __name__ == "__main__":
    toolkit = IDAThemingToolkit()
    toolkit.show()


repository: tmr232/ida-plugins
file: lca/lca.py
description: IDA Pro Plugins
url: https://github.com/tmr232/ida-plugins/blob/c333ebecc9b3d641cdfdc9ab0d66f73e0b551fc3/lca/lca.py
wanted_name: Lowest Common Ancestors
comment: Lowest Common Ancestors
created_at: 2015-05-27 12:35:29+00:00
pushed_at: 2015-05-31 13:05:09+00:00
content:
from awesome.context import ignored
import sark
import sark.graph
import networkx as nx
import sark.ui
import idc
import idaapi


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

PADDING = 1
PAD_WIDTH = 3


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


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

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


    return RemoveTargetHandler


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

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


    return DisableSourceHandler


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

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


    return EnableSourceHandler


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

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

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


    return AddAddressHandler


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

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

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

    return AddFunctionHandler


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

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

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

        self._lca_graph = nx.DiGraph()

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

        self._current_node_id = 0

        self._disabled_sources = set()

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

        self._node_ids = {}


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

    def OnGetText(self, node_id):
        return pad(idc.Name(self[node_id]))

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

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

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

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

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

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

        self._targets.add(target)

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

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

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

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

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

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

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

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

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

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


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


    def OnRefresh(self):
        self.Clear()

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

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

        self._node_ids = node_ids

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

        self.color_nodes()

        return True

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

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

    def OnDeactivate(self):
        self._unregister_handlers()

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

        return True

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

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

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

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

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

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

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

        return False


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

        def _activate(self, ctx):
            lca_plugin.show_graph()

    return LCAViewerStarter


def idaview_add_target_handler(lca_plugin):
    class IDAViewAddTargetHandler(sark.ui.ActionHandler):
        TEXT = "Add LCA Target"

        def _activate(self, ctx):
            if lca_plugin._lca_viewer:
                with ignored(KeyError):
                    lca_plugin._lca_viewer.add_target(ctx.cur_ea)
                    lca_plugin._lca_viewer.rebuild_graph()
                    idaapi.msg("[LCA] Target Added: {}\n".format(idc.Name(ctx.cur_ea)))

    return IDAViewAddTargetHandler


def idaview_hooks(idaview_handler):
    class Hooks(idaapi.UI_Hooks):
        def finish_populating_tform_popup(self, form, popup):
            if idaapi.get_tform_type(form) == idaapi.BWN_DISASM:
                idaapi.attach_action_to_popup(form, popup, idaview_handler.get_name(), "")

    return Hooks


class LCA(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "Lowest Common Ancestors"
    help = "Lowest Common Ancestors"
    wanted_name = "Lowest Common Ancestors"
    wanted_hotkey = ""

    def init(self):
        self._lca_viewer = None

        self._lca_starter = lca_viewer_starter(self)
        self._lca_starter.register()
        idaapi.attach_action_to_menu("View/Graph Overview", self._lca_starter.get_name(), idaapi.SETMENU_APP)

        self._idaview_handler  = idaview_add_target_handler(self)
        self._idaview_handler.register()
        self._hooks = idaview_hooks(self._idaview_handler)()
        self._hooks.hook()

        return idaapi.PLUGIN_KEEP

    def term(self):
        self._lca_starter.unregister()
        self._idaview_handler.unregister()
        self._hooks.unhook()

    def show_graph(self):
        if self._lca_viewer is None:
            self._lca_viewer = LCAGraph("LCA Graph")
        self._lca_viewer.Show()


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


def PLUGIN_ENTRY():
    return LCA()


repository: synacktiv/vmx_intrinsics
file: vmx_intrinsics.py
description: VMX intrinsics plugin for Hex-Rays decompiler
wanted_name: VMXIntrinsic
comment: VMX intrinsics plugin for Hex-Rays decompiler
created_at: 2019-10-28 13:02:46+00:00
pushed_at: 2019-10-28 14:38:27+00:00
url:
https://github.com/synacktiv/vmx_intrinsics/blob/add1d90cff9fd47459c66a92e5cf8431383ea02b/vmx_intrinsics.py

content:
#
# This script is mostly based on the previous work done by https://dougallj.wordpress.com/2018/06/04/writing-a-hex-rays-plugin-vmx-intrinsics/
# Now it's a python plugin and some renaming was necessary as MICROCODE API has changed (cf ida_hexrays.py)
# Use of mop_a/mop_addr_t for taking the address of registers for vmread/vmclear/vmptrld/vmptrst/vmxon
#

import idaapi
import ida_hexrays
import ida_allins
import ida_idp

# Identify 'tt' register
# import ida_hexrays
# for i in range(0, 0xFF):
#     print("0x{0:02X} => {1}".format(i, ida_hexrays.get_mreg_name(i, 1)))
IDA_TT = 0xC0

def is_gte_ida74():
    major, minor = map(int, idaapi.get_kernel_version().split("."))
    return (major == 7 and minor >= 4)

class MicroInstruction(ida_hexrays.minsn_t):

    def __init__(self, opcode, ea):
        ida_hexrays.minsn_t.__init__(self, ea)
        self.opcode = opcode
        self.l.zero()
        self.r.zero()
        self.d.zero()

class CallBuilder():

    def __init__(self, cdg, name, return_type=idaapi.tinfo_t(idaapi.BT_VOID)):
        self.emitted = False
        self.cdg = cdg
        self.callinfo = ida_hexrays.mcallinfo_t()
        self.callinfo.callee = idaapi.BADADDR
        self.callinfo.solid_args = 0x00
        self.callinfo.call_spd = 0x00
        self.callinfo.stkargs_top = 0x00
        self.callinfo.cc = idaapi.CM_CC_FASTCALL
        self.callinfo.return_type = return_type
        self.callinfo.flags = idaapi.FCI_SPLOK | idaapi.FCI_FINAL | idaapi.FCI_PROP
        self.callinfo.role = idaapi.ROLE_UNK

        glbhigh_off = cdg.mba.get_stack_region().off + cdg.mba.get_stack_region().size
        # what memory is visible to the call : GLBLOW - GLBHIGH
        self.callinfo.visible_memory.add(ida_hexrays.ivl_t(0x00, 0x100000))
        self.callinfo.visible_memory.add(ida_hexrays.ivl_t(glbhigh_off, 0xFFFFFFFFFFFFFFFF - glbhigh_off))
        # spoiled locations : GLBLOW - GLBHIGH
        self.callinfo.spoiled.mem.add(ida_hexrays.ivl_t(0x00, 0x100000))
        self.callinfo.spoiled.mem.add(ida_hexrays.ivl_t(glbhigh_off, 0xFFFFFFFFFFFFFFFF - glbhigh_off))

        self.callins = MicroInstruction(ida_hexrays.m_call, self.cdg.insn.ea)
        self.callins.l.make_helper(name)
        self.callins.d.t = ida_hexrays.mop_f
        self.callins.d.size = 0x00
        self.callins.d.f = self.callinfo

        if (return_type.is_void()):
            self.ins = self.callins
        else:
            self.callins.d.size = return_type.get_size()
            self.ins = MicroInstruction(ida_hexrays.m_mov, self.cdg.insn.ea)
            self.ins.l.t = ida_hexrays.mop_d
            self.ins.l.d = self.callins
            self.ins.l.size = self.callins.d.size
            self.ins.d.t = ida_hexrays.mop_r
            self.ins.d.r = 0x00
            self.ins.d.size = self.callins.d.size

    def add_register_argument(self, t, operand):
        ca = ida_hexrays.mcallarg_t()
        ca.t = idaapi.mop_r
        ca.r = operand
        ca.type = t
        ca.size = t.get_size()
        self.callinfo.args.push_back(ca)
        self.callinfo.solid_args += 1

    def add_register_address_argument(self, t, operand):
        addr_t = ida_hexrays.mop_addr_t()
        addr_t.t = idaapi.mop_r
        addr_t.r = operand
        addr_t.type = t
        addr_t.size = t.get_size()
        addr_t.insize = t.get_size()
        addr_t.outsize = t.get_size()

        ca = ida_hexrays.mcallarg_t()
        ca.t = idaapi.mop_a
        ca.a = addr_t
        t.create_ptr(t)
        ca.type = t
        ca.size = t.get_size()
        self.callinfo.args.push_back(ca)
        self.callinfo.solid_args += 1

    def set_return_register(self, reg):
        self.ins.d.r = reg

    def emit(self):
        if self.emitted == False:
            self.cdg.mb.insert_into_block(self.ins , self.cdg.mb.tail)
            self.emitted = True

    def emit_und_reg(self, reg, size):
        ins = MicroInstruction(ida_hexrays.m_und, self.cdg.insn.ea)
        ins.d.t = idaapi.mop_r
        ins.d.r = reg
        ins.d.size = size
        self.cdg.mb.insert_into_block(ins, self.cdg.mb.tail)

    def emit_reg_equals_number(self, result_reg, reg, number, size):
        ins = MicroInstruction(ida_hexrays.m_setz, self.cdg.insn.ea)
        ins.l.t = idaapi.mop_r
        ins.l.r = reg
        ins.l.size = size
        ins.r.make_number(number, size)
        ins.d.t = idaapi.mop_r
        ins.d.r = result_reg
        ins.d.size = 1;
        self.cdg.mb.insert_into_block(ins, self.cdg.mb.tail)

class VMXFilter(ida_hexrays.microcode_filter_t):

    def __init__(self, itype):
        ida_hexrays.microcode_filter_t.__init__(self)
        self.itype = itype
        self.installed = False
        self.toggle_install()

    def match(self, cdg):
        return cdg.insn.itype == self.itype

    def install(self):
        ida_hexrays.install_microcode_filter(self, True);
        self.installed = True

    def uninstall(self):
        ida_hexrays.install_microcode_filter(self, False);
        self.installed = False

    def toggle_install(self):
        if self.installed:
            self.uninstall()
        else:
            self.install()

class VMXVmwrite(VMXFilter):

    def __init__(self):
        VMXFilter.__init__(self, ida_allins.NN_vmwrite)
        self.name = "__vmx_vmwrite"

    def apply(self, cdg):
        builder = CallBuilder(cdg, self.name, idaapi.tinfo_t(idaapi.BT_INT8 | idaapi.BTMT_UNSIGNED))
        builder.add_register_argument(idaapi.tinfo_t(idaapi.BT_INT64 | idaapi.BTMT_UNSIGNED), cdg.load_operand(0))
        builder.add_register_argument(idaapi.tinfo_t(idaapi.BT_INT64 | idaapi.BTMT_UNSIGNED), cdg.load_operand(1))
        builder.emit()
        builder.set_return_register(IDA_TT)
        builder.emit_reg_equals_number(ida_hexrays.mr_zf, IDA_TT, 1, 1)
        builder.emit_reg_equals_number(ida_hexrays.mr_cf, IDA_TT, 2, 1)
        return idaapi.MERR_OK

class VMXVmread(VMXFilter):

    def __init__(self):
        VMXFilter.__init__(self, ida_allins.NN_vmread)
        self.name = "__vmx_vmread"

    def apply(self, cdg):
        builder = CallBuilder(cdg, self.name, idaapi.tinfo_t(idaapi.BT_INT8 | idaapi.BTMT_UNSIGNED))
        builder.add_register_argument(idaapi.tinfo_t(idaapi.BT_INT64 | idaapi.BTMT_UNSIGNED), cdg.load_operand(1))
        builder.add_register_address_argument(idaapi.tinfo_t(idaapi.BT_INT64 | idaapi.BTMT_UNSIGNED), cdg.load_operand(0))
        builder.emit()
        builder.set_return_register(IDA_TT)
        builder.emit_und_reg(ida_hexrays.mr_zf, 1)
        builder.emit_und_reg(ida_hexrays.mr_cf, 1)
        return idaapi.MERR_OK

class VMXVoidReturn(VMXFilter):

    def __init__(self, itype, name):
        VMXFilter.__init__(self, itype)
        self.name = name

    def apply(self, cdg):
        builder = CallBuilder(cdg, self.name, idaapi.tinfo_t(idaapi.BT_INT8 | idaapi.BTMT_UNSIGNED))
        builder.emit()
        builder.set_return_register(IDA_TT)
        builder.emit_reg_equals_number(ida_hexrays.mr_zf, IDA_TT, 1, 1)
        builder.emit_reg_equals_number(ida_hexrays.mr_cf, IDA_TT, 2, 1)
        return idaapi.MERR_OK

class VMXVoid(VMXFilter):

    def __init__(self, itype, name):
        VMXFilter.__init__(self, itype)
        self.name = name

    def apply(self, cdg):
        builder = CallBuilder(cdg, self.name)
        builder.emit()
        builder.emit_und_reg(ida_hexrays.mr_zf, 1)
        builder.emit_und_reg(ida_hexrays.mr_cf, 1)
        return idaapi.MERR_OK

class VMXVMCSReturn(VMXFilter):

    def __init__(self, itype, name):
        VMXFilter.__init__(self, itype)
        self.name = name

    def apply(self, cdg):
        builder = CallBuilder(cdg, self.name, idaapi.tinfo_t(idaapi.BT_INT8 | idaapi.BTMT_UNSIGNED))
        builder.add_register_address_argument(idaapi.tinfo_t(idaapi.BT_INT64 | idaapi.BTMT_UNSIGNED), cdg.load_operand(0))
        builder.emit()
        builder.set_return_register(IDA_TT)
        builder.emit_reg_equals_number(ida_hexrays.mr_zf, IDA_TT, 1, 1)
        builder.emit_reg_equals_number(ida_hexrays.mr_cf, IDA_TT, 2, 1)
        return idaapi.MERR_OK

class VMXIntrinsicIDAPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = "VMX intrinsics plugin for Hex-Rays decompiler"
    help = ""
    wanted_name = "VMXIntrinsic"
    wanted_hotkey = ""

    def __init__(self):
        self.filters = []

    def init(self):
        if is_gte_ida74() == False:
            print("[-] {0} : for ida >= 7.4 only".format(self.wanted_name))
            return idaapi.PLUGIN_SKIP
        if idaapi.ph.id != idaapi.PLFM_386 and not idaapi.get_inf_structure().is_64bit():
            print("[-] {0} : for x64 only".format(self.wanted_name))
            return idaapi.PLUGIN_SKIP
        if not ida_hexrays.init_hexrays_plugin():
            print("[-] {0} : no decompiler available, skipping".format(self.wanted_name))
            return idaapi.PLUGIN_SKIP
        self.add_filter(VMXVmwrite())
        self.add_filter(VMXVmread())
        self.add_filter(VMXVoidReturn(ida_allins.NN_vmlaunch, "__vmx_vmlaunch"))
        self.add_filter(VMXVoidReturn(ida_allins.NN_vmresume, "__vmx_vmresume"))
        self.add_filter(VMXVoid(ida_allins.NN_vmxoff, "__vmx_off"))
        self.add_filter(VMXVMCSReturn(ida_allins.NN_vmclear, "__vmx_vmclear"))
        self.add_filter(VMXVMCSReturn(ida_allins.NN_vmptrld, "__vmx_vmptrld"))
        self.add_filter(VMXVMCSReturn(ida_allins.NN_vmptrst, "__vmx_vmptrst"))
        self.add_filter(VMXVMCSReturn(ida_allins.NN_vmxon, "__vmx_on"))

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        idaapi.warning("[-] {0} cannot be run as a script in IDA.".format(self.wanted_name))

    def add_filter(self, f):
        self.filters.append(f)

    def term(self):
        for f in self.filters:
            f.toggle_install()
        self.filters = []

def PLUGIN_ENTRY():
    return VMXIntrinsicIDAPlugin()


repository: zmrbak/IDAPython
file: 配套代码/ida10/python/examples/hexrays/vds3.py
description: 在线课程《IDAPython脚本化软件逆向分析入门》课程源码
wanted_name: Hex-Rays if-inverter (IDAPython)
comment: Sample plugin3 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/vds3.py

content:
"""
summary: invert if/else blocks

description:
  Registers an action that can be used to invert the `if`
  and `else` blocks of a `ida_hexrays.cif_t`.

  For example, a statement like

      if ( cond )
      {
        statements1;
      }
      else
      {
        statements2;
      }

  will be displayed as

      if ( !cond )
      {
        statements2;
      }
      else
      {
        statements1;
      }

  The modifications are persistent: the user can quit & restart
  IDA, and the changes will be present.

author: EiNSTeiN_ (einstein@g3nius.org)
"""

from __future__ import print_function

import idautils

import ida_kernwin
import ida_hexrays
import ida_netnode
import ida_idaapi
import ida_idp

import traceback

NETNODE_NAME = '$ hexrays-inverted-if'

inverter_actname = "vds3:invert"

class invert_action_handler_t(ida_kernwin.action_handler_t):
    def __init__(self, inverter):
        ida_kernwin.action_handler_t.__init__(self)
        self.inverter = inverter

    def activate(self, ctx):
        vdui = ida_hexrays.get_widget_vdui(ctx.widget)
        self.inverter.invert_if_event(vdui)
        return 1

    def update(self, ctx):
        return ida_kernwin.AST_ENABLE_FOR_WIDGET if \
            ctx.widget_type == ida_kernwin.BWN_PSEUDOCODE else \
            ida_kernwin.AST_DISABLE_FOR_WIDGET


class hexrays_callback_info(object):

    def __init__(self):
        self.vu = None

        self.node = ida_netnode.netnode()
        if not self.node.create(NETNODE_NAME):
            # node exists
            self.load()
        else:
            self.stored = []

        return

    def load(self):

        self.stored = []

        try:
            data = self.node.getblob(0, 'I')
            if data:
                self.stored = eval(data.decode("UTF-8"))
                print('Invert-if: Loaded %s' % (repr(self.stored), ))
        except:
            print('Failed to load invert-if locations')
            traceback.print_exc()
            return

        return

    def save(self):

        try:
            self.node.setblob(repr(self.stored).encode("UTF-8"), 0, 'I')
        except:
            print('Failed to save invert-if locations')
            traceback.print_exc()
            return

        return

    def invert_if(self, insn):

        if insn.opname != 'if':
            return False

        cif = insn.details

        if not cif.ithen or not cif.ielse:
            return False

        ida_hexrays.qswap(cif.ithen, cif.ielse)
        # Make a copy of 'cif.expr': 'lnot' might destroy its toplevel
        # cexpr_t and return a pointer to its direct child (but we'll want to
        # 'swap' it later, the 'cif.expr' cexpr_t object must remain valid.)
        cond = ida_hexrays.cexpr_t(cif.expr)
        notcond = ida_hexrays.lnot(cond)

        cif.expr.swap(notcond)

        return True

    def add_location(self, ea):
        if ea in self.stored:
            self.stored.remove(ea)
        else:
            self.stored.append(ea)
        self.save()
        return

    def find_if_statement(self, vu):

        vu.get_current_item(ida_hexrays.USE_KEYBOARD)
        item = vu.item

        if item.is_citem() and item.it.op == ida_hexrays.cit_if and item.it.to_specific_type.cif.ielse is not None:
            return item.it.to_specific_type

        if vu.tail.citype == ida_hexrays.VDI_TAIL and vu.tail.loc.itp == ida_hexrays.ITP_ELSE:
            # for tail marks, we know only the corresponding ea,
            # not the pointer to if-statement
            # find it by walking the whole ctree
            class if_finder_t(ida_hexrays.ctree_visitor_t):
                def __init__(self, ea):
                    ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST | ida_hexrays.CV_INSNS)

                    self.ea = ea
                    self.found = None
                    return

                def visit_insn(self, i):
                    if i.op == ida_hexrays.cit_if and i.ea == self.ea:
                        self.found = i
                        return 1 # stop enumeration
                    return 0

            iff = if_finder_t(vu.tail.loc.ea)
            if iff.apply_to(vu.cfunc.body, None):
                return iff.found

        return

    def invert_if_event(self, vu):

        i = self.find_if_statement(vu)
        if not i:
            return False

        if self.invert_if(i):
            vu.refresh_ctext()
            self.add_location(i.ea)

        return True

    def restore(self, cfunc):

        class visitor(ida_hexrays.ctree_visitor_t):

            def __init__(self, inverter, cfunc):
                ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST | ida_hexrays.CV_INSNS)
                self.inverter = inverter
                self.cfunc = cfunc
                return

            def visit_insn(self, i):
                try:
                    if i.op == ida_hexrays.cit_if and i.ea in self.inverter.stored:
                        self.inverter.invert_if(i)
                except:
                    traceback.print_exc()
                return 0 # continue enumeration

        visitor(self, cfunc).apply_to(cfunc.body, None)

        return


class vds3_hooks_t(ida_hexrays.Hexrays_Hooks):
    def __init__(self, i):
        ida_hexrays.Hexrays_Hooks.__init__(self)
        self.i = i

    def populating_popup(self, widget, phandle, vu):
        ida_kernwin.attach_action_to_popup(vu.ct, None, inverter_actname)
        return 0

    def maturity(self, cfunc, maturity):
        if maturity == ida_hexrays.CMAT_FINAL:
            self.i.restore(cfunc)
        return 0

class idp_hooks_t(ida_idp.IDP_Hooks):
    def __init__(self, i):
        ida_idp.IDP_Hooks.__init__(self)
        self.i = i

    # 'node' refers to index of the named node, this index became invalid after
    # privrange moving, so we recreate the node here to update nodeidx
    def ev_privrange_changed(self, old_privrange, delta):
        i.node.create(NETNODE_NAME)


# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays if-inverter (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin3 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            i = hexrays_callback_info()
            ida_kernwin.register_action(
                ida_kernwin.action_desc_t(
                    inverter_actname,
                    "Invert then/else",
                    invert_action_handler_t(i),
                    "I"))
            self.vds3_hooks = vds3_hooks_t(i)
            self.vds3_hooks.hook()
            # we need this hook to react to privrange moving event
            self.idp_hooks = idp_hooks_t(i)
            self.idp_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds3_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: kiwidoggie-productions/IDA-ConsoleHacking-Scripts
file: PS2/SnDebugPlugin.py
description: Assortment of IDA scripts for an assortment of game consoles, PS2, PS4, etc
wanted_name: PS2 SnDebug Symbols Loader
comment: PlayStation 2 .sndebug symbols loader
created_at: 2019-12-21 14:39:05+00:00
pushed_at: 2021-02-26 18:50:47+00:00
url:
https://github.com/kiwidoggie-productions/IDA-ConsoleHacking-Scripts/blob/ea3446bf291468bd71fd1634d3b28c5452509c03/PS2/SnDebugPlugin.py

content:
import idaapi

#------------------------------------------------------------------------------
# IDA Plugin
# Initial creation: Oct 27, 2019 (kd)
# Plugin port: Dec 20, 2019 (Cleanup for submission, kd)
# Fix bug: Jan 2, 2020 (kd)
# Install by placing in <IDA Root>/plugins
#------------------------------------------------------------------------------
VERSION = "v0.1"
AUTHORS = [ 'kd', '' ]

# https://unit42.paloaltonetworks.com/using-idapython-to-make-your-life-easier-part-1/
def ReadAsciiString(ea):
    stringData = ""
    while True:
        if idaapi.get_byte(ea) != 0:
            stringData += chr(idaapi.get_byte(ea))
        else:
            break

        ea += 1
    
    return stringData

#
# Align(input address, alignment)
# Aligns the input address to the specified alignment, and returns the new value
#
def Align(inputEa, alignment):
    if (inputEa % alignment) == 0:
        return inputEa
    
    num = alignment - (inputEa % alignment)
    return inputEa + num

#
# Entry structure holding the name, and target (in executable) address
#
class SnDebugEntry:
    nameAddress = idaapi.BADADDR
    targetAddress = idaapi.BADADDR
    unknown08 = 0

    def __init__(self, parseEa):
        currentPos = 0

        self.nameAddress = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.targetAddress = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown08 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4

#
# String entry structure, holding the offset of the string, and the actual string data
#
class SnStringEntry:
    offset = 0
    stringData = ""

    def __init__(self, pOffset, pString):
        self.offset = pOffset
        self.stringData = pString

#
# The .sndebug header, this has a bunch of information I'm not fully aware of at this time
# but as we find out, will fill in the unknowns
#
class SnDebugHeader:
    def __init__(self, parseEa):
        # Keep current position for tracking
        currentPos = 0

        # Read out the magic 'SNR1' or 'SNR2'
        self.magic = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4

        idaapi.msg("info: magic: (%x)\n" % self.magic)

        # Unknown data
        self.unknown04 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown08 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        
        # Start of symbol table
        self.symTabStart = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4

        # RE4 == 7409, SOCOM = ?
        # Parse out the symbol count
        self.symbolCount = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4

        # More unknowns
        self.unknown14 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown18 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown1C = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        
        # Offset to end of table, from this point on is zeros
        self.endOfTable = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        
        # More unknowns
        self.unknown24 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown28 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown2C = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown30 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown34 = idaapi.get_32bit(parseEa + currentPos)
        currentPos = currentPos + 4
        self.unknown38 = idaapi.get_32bit(parseEa + currentPos)

#
# IDA Plugin implementation
#
class ps2_sndebugplugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "PlayStation 2 .sndebug symbols loader"
    help = "This will parse PlayStation 2 debug symbols from the .sndebug code section"
    wanted_name = "PS2 SnDebug Symbols Loader"
    wanted_hotkey = ""

    symbolList = [ ]
    entries = [ ]

    #
    # IDA calls this on startup
    #
    def init(self):
        idaapi.msg("PlayStation 2 .sndebug symbol loader initialized\n")
        return idaapi.PLUGIN_OK
    
    #
    # This is called when the user invokes this script
    #
    def run(self, arg):
        #idaapi.msg("info: run called\n")

        info = idaapi.get_inf_structure()
        if info.is_32bit() == False:
            idaapi.msg("err: not 32-bit target\n")
            return
        
        if info.is_be():
            idaapi.msg("err: not little-endian target\n")
            return
        
        idaapi.msg("info: processor (%s)\n" % info.procName)

        # Get the required segments
        snDebugSegment = idaapi.get_segm_by_name(".sndata")
        idaapi.msg("info: .sndata segment: (0x%x)\n" % snDebugSegment.start_ea)

        textSegment = idaapi.get_segm_by_name(".text")
        idaapi.msg("info: .text segment: (0x%x)\n" % textSegment.start_ea)

        # Get the start and end ea's
        snDebugStartEa = snDebugSegment.start_ea
        snDebugEndEa = snDebugSegment.end_ea

        # Check to see if our start effective address is valid
        if snDebugStartEa is idaapi.BADADDR or snDebugEndEa is idaapi.BADADDR:
            return
        
        # read out the header
        self.header = SnDebugHeader(snDebugStartEa)
        if self.header.magic != 0x31524e53 and self.header.magic != 0x32524e53:
            idaapi.msg("err: invalid header\n")
            return
        
        # Print something useful
        symbolCount = self.header.symbolCount
        endOffset = self.header.endOfTable

        idaapi.msg("info: preparing to parse (%d) symbols... [endOffset: (%x)]\n" % (symbolCount, endOffset))
        
        # Set current position to first string
        firstEntry = idaapi.get_32bit(self.header.symTabStart)
        currentPos = (firstEntry - snDebugStartEa) 

        for _ in range(symbolCount):
            symOffset = snDebugStartEa + currentPos
            symStringData = ReadAsciiString(snDebugStartEa + currentPos)
            currentPos += len(symStringData) + 1
            #idaapi.msg("info: adding (%d) (%x) (%s)\n" % (symIndex, symOffset, symStringData))
            # Add it to our list
            self.symbolList.append(SnStringEntry(symOffset, symStringData))
        
        # Set current position to start of symbol table
        currentPos = self.header.symTabStart

        for _ in range(symbolCount):
            self.entries.append(SnDebugEntry(currentPos))
            currentPos += 12 # skip the size of the entry
        
        for index, entry in enumerate(self.entries):
            symbol = None
            for _, sym in enumerate(self.symbolList):
                # Check if this symbol offset is equal to the symbol entry offset
                #idaapi.msg("(%x) == (%x)\n" % (sym.offset, entry.nameAddress))
                if sym.offset == entry.nameAddress:
                    #idaapi.msg("info: got sym (%x)\n" % sym.offset)
                    symbol = sym
                    break
            #
            # After we loop, check to see if we have a symbol
            if symbol is None:
                idaapi.msg("err: could not find symbol for (%d)\n" % index)
                continue
            #
            #func = idaapi.get_func(entry.targetAddress)
            #if func is None:
            #    idaapi.msg("no func (0x%x, %s)\n" % (entry.targetAddress, symbol.stringData))
            #
            #idaapi.msg("succ: set_name(0x%x, 0x%x, %s)\n" % (symbol.offset, entry.targetAddress, symbol.stringData))
            idaapi.set_name(entry.targetAddress, symbol.stringData, idaapi.SN_NOWARN)
        
        idaapi.msg("succ: label all symbols completed\n")
    
    # This is terminating the plugin
    def term(self):
        #idaapi.msg("term called\n")
        pass

#
# This is how IDA finds the entry points
#
def PLUGIN_ENTRY():
    return ps2_sndebugplugin_t()


repository: miunasu/REAI_IDA
file: REAI_plugin.py
url: https://github.com/miunasu/REAI_IDA/blob/7e9fa881dd2c382e98791207d7680e86d5bc09ed/REAI_plugin.py
wanted_name: REAI
comment: REAI_PLUGIN
created_at: 2025-04-17 08:20:33+00:00
pushed_at: 2025-04-18 09:41:46+00:00
description:
Use idapython create call topology, AI analysis function layer by layer. Get the target function conclusion.利用idapython创建函数调用拓扑，AI逐层分析函数，最终总结目标函数的功能。

content:
import idaapi
import ida_funcs
import idc
import requests
import json
import ida_hexrays
import urllib3
import threading
import queue

urllib3.disable_warnings()

# CodeCheck check flag
CodeCheck_flag = True

AI_return = queue.Queue()
exception_code_collection = []
processed_func = []
processing_func = []

# Function information base of dictionary structure
# func start ea：node object
function_info = {}

# LLM API config
API_KEY = "api_key"  
API_URL = "api_url"  
MODEL = "model name"  

# prompt language
chinese_prompt = "你是一个恶意代码分析师,现在分析一个C语言函数,以json格式把结果返回给我,要求包含两个对象,'des':对该函数功能的中文描述。'name':给该函数一个合适的英文名称。代码如下："
english_prompt = "You are a malware analyst. Now analyze a C language function and return the result to me in JSON format. The result should contain two objects: 'des' — a English description of the function's purpose, and 'name' — an appropriate English name for the function. The code is as follows: "
pre_prompt = english_prompt

def chat_with_AI(prompt):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }

    payload = {
        "model": MODEL,  
        "messages": [
            {"role": "user", "content": pre_prompt + prompt}
        ],
        "temperature": 1.0, # recommend 1.0
        "max_tokens": 1024
    }


    try:
        response = requests.post(
            API_URL,
            headers=headers,
            data=json.dumps(payload)
        )
        response.raise_for_status()  
        result = json.loads(response.json()['choices'][0]['message']['content'].split("```json")[-1].split("```")[0])
        return result if response.status_code == 200 else None

    except Exception as e:
        print(f"[REAI] error: {type(e).__name__} - {str(e)}")
        print(f'[REAI] text : {response.text}')
        return None

# rename function, pre 'AI_' to mark AI recognition
def rename_function(func_ea, new_name):
    current_name = ida_funcs.get_func_name(func_ea)
    if new_name != current_name and 'sub_' not in new_name:
        return idaapi.set_name(func_ea, 'AI_' + new_name, idaapi.SN_NOWARN)
    return False


# get caller address and call address
# eg: 0x2255a7(caller address) call to : 0x784c(next function)
def get_function_calls(func_ea):
    cfunc = ida_hexrays.decompile(func_ea)
    function_calls = {}

    class Visitor(ida_hexrays.ctree_visitor_t):
        def __init__(self):
            ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_PARENTS)

        def visit_expr(self, e):
            if e.op == ida_hexrays.cot_call:
                function_calls[e.ea] = e.x.obj_ea
            return 0

    visitor = Visitor()
    visitor.apply_to(cfunc.body, None)

    return function_calls
    

# add comment to pseudocode
def add_decompiled_comment(ea, comment):
    cfunc = ida_hexrays.decompile(ea)
    if not cfunc:
        print(f"[REAI] Failed to decompile function at {hex(ea)}")
        return False

    class Visitor(ida_hexrays.ctree_visitor_t):
        def __init__(self):
            ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_PARENTS)

        def visit_expr(self, e):
            if e.ea == ea:
                loc = ida_hexrays.treeloc_t()
                loc.ea = e.ea
                loc.itp = idaapi.ITP_SEMI
                cfunc.set_user_cmt(loc, comment)
                cfunc.save_user_cmts()
                cfunc.refresh_func_ctext()
                return 1  
            return 0

    visitor = Visitor()
    visitor.apply_to(cfunc.body, None)
    cfunc.refresh_func_ctext()
    return True

# create function for all call insn
def call_add_func(ea):
    try:
        func = ida_funcs.get_func(ea)
        func_start = func.start_ea
        func_end = func.end_ea
        c_func = ida_hexrays.decompile(ea)
        point = func_start        
        while point <= func_end:
            disasm = idc.GetDisasm(point)
            # 'call    near ptr unk_5D5DACB'
            if 'call' in disasm and 'unk_' in disasm:
                address = int(disasm.split(' ')[-1].strip()[4:], 16)
                flag = ida_funcs.add_func(address)
                if flag == True:
                    print(f'[REAI] decompile call insn at {hex(point)}')
                else:
                    print(f'[REAI] fail decompile call insn at {hex(point)}')

            point = idc.next_head(point)
        c_func.refresh_func_ctext()

    except Exception as e:
        raise Exception(f"[REAI] Exception during decompile at {hex(ea)}: {e}")


# exception code check
def exception_code_check(ea, caller_ea_arg = 0):
    global exception_code_collection
    global processed_func
    global processing_func

    func = ida_funcs.get_func(ea)
    func_name = ida_funcs.get_func_name(ea)
    if caller_ea_arg != 0 and (ea in processed_func or ea in processing_func or func == None or 'sub_' not in func_name):
        return True
    
    processing_func.append(ea)
    call_add_func(ea)

    if CodeCheck_flag == True:
        cfunc = idaapi.decompile(ea)
        if "JUMPOUT" in str(cfunc) or 'MEMORY[' in str(cfunc):
            exception_code_collection.append(hex(ea))

    sub_func = get_function_calls(ea)
    for caller_ea, call_ea in sub_func.items():        
        exception_code_check(call_ea, caller_ea)

    processed_func.append(ea)
    if caller_ea_arg == 0:
        processing_func = []
        if exception_code_collection == []:
            return True
        else:
            return False

class node:
    def __init__(self):
        # start ea of parent function
        self.parent = []
        # sub function info. ea:caller_ea
        self.child = {}
        # parent function chain
        self.parent_chain = set()
        # current function start ea
        self.ea = 0
        # caller function address, used to write comments
        self.caller_ea = []

    
# get call topology
def get_call_topology(ea, caller_ea, parent_ea, parent_chain, count):
    # ea = current function start address; caller_ea = address of the function that calls this function; parent_ea = start address of the parent function; parent_chain = parent function chain
    
    global function_info
    
    # func node create
    current_node = node()
    current_node.ea = ea
    current_node.caller_ea.append(caller_ea)
    current_node.parent_chain = parent_chain | {parent_ea}
    current_node.parent.append(parent_ea)

    sub_func = get_function_calls(ea)
    for caller_ea, call_ea in sub_func.items():
        func_name = ida_funcs.get_func_name(call_ea)
        func = ida_funcs.get_func(call_ea)
        if current_node.child.get(call_ea) == None and call_ea not in current_node.parent_chain and call_ea != ea and func != None and 'sub_' in func_name:
            current_node.child[call_ea] = caller_ea

    function_info[ea] = current_node
    count += 1
    func_name = ida_funcs.get_func_name(ea)
    
    for call_ea, caller_ea in current_node.child.items():
        if function_info.get(call_ea) == None:
            get_call_topology(call_ea, caller_ea, ea, current_node.parent_chain, count)
        else:
            function_info[call_ea].caller_ea.append(caller_ea)
            function_info[call_ea].parent.append(ea)
            function_info[call_ea].parent_chain = function_info[call_ea].parent_chain | current_node.parent_chain | {ea}


# print call topology
def pt():
    for ea, func_node in function_info.items():
        parent = [hex(i) for i in func_node.parent]
        parent_chain = [hex(i) for i in func_node.parent_chain]
        del parent_chain[0]
        caller_ea = [hex(i) for i in func_node.caller_ea]
        func_name = ida_funcs.get_func_name(func_node.ea)
        print('----------------------')
        print(f'func name : {func_name}')
        print(f'parent func: {parent}')
        print(f'parent chain: {parent_chain}')
        print(f'caller ea: {caller_ea}')

        for i, j in func_node.child.items():
            print(f'sub func call ea: {hex(i)} caller ea: {hex(j)}')


def select_address_check(ea):
    func = ida_funcs.get_func(ea) 
    c_func = ida_hexrays.decompile(ea)
    if func == None or c_func == None:
        raise Exception(f'[REAI] select wrong address, try again')
    return func.start_ea


def AI_work(ea, pseudo_code, func_name):
    if (api_result := chat_with_AI(pseudo_code)):
        new_func_name = api_result.get("name", func_name)
        description = api_result.get("des", "None")
        info = []
        info.append(ea)
        info.append(new_func_name)
        info.append(description)
        AI_return.put(info)
    else:
        print(f'[REAI] bad api result, func: {func_name} ea: {hex(ea)}')
        AI_return.put(['bad'])


# use call topology to analyze function
def AI_ananalyze(func_start):
    global function_info
    count = 0
    while(len(function_info[func_start].child) != 0):
        print('------------------------------')
        print(f'[REAI] round {count}')
        if count >= 30:
            # maybe infinite loop, check it
            pt()
            raise Exception('[REAI] out of 30 round, exit')

        round_list = []
        thread_list = []
        for ea, func_node in function_info.items():
            if len(func_node.child) == 0:
                round_list.append(func_node.ea)
                c_func = ida_hexrays.decompile(ea)
                c_func.refresh_func_ctext()
                func_name = ida_funcs.get_func_name(ea)
                sub_thread = threading.Thread(target=AI_work, args=(func_node.ea, str(c_func), func_name))
                thread_list.append(sub_thread)
                sub_thread.start()

        for i in thread_list:
            i.join()
            info = AI_return.get()
            if info[0] == 'bad':
                print('[REAI] detect API error, continue')
                continue
            func_ea = info[0]
            new_func_name = info[1]
            description = info[2]
            name_count = 0
            old_name = ida_funcs.get_func_name(func_ea)
            rename_flag = rename_function(func_ea, new_func_name)
            temp_name = new_func_name + '_'

            while rename_flag == False:
                if name_count >= 5:
                    print('[REAI] multiple rename, fail!')
                    break
                new_func_name = temp_name + str(name_count)
                rename_flag = rename_function(func_ea, new_func_name)
                name_count += 1
                
            if rename_flag == False:
                print(f"[REAI] fail !!! Renamed: {old_name} -> {'AI_' + new_func_name}")
                name = old_name
            else:
                print(f"[REAI] Renamed: {old_name} -> {'AI_' + new_func_name}")
                name = 'AI_' + new_func_name

            ida_hexrays.decompile(func_ea).refresh_func_ctext()

            for i in function_info[func_ea].caller_ea:
                add_decompiled_comment(i, description + f' by func: {name}')
        

        for i in round_list:
            del function_info[i]
        
        for ea, func_node in function_info.items():
            for i in round_list:
                if func_node.child.get(i) != None:
                    del func_node.child[i]
        
        print(f'[REAI] Finish round {count}, parse list: {", ".join([hex(addr) for addr in round_list])}')
        count += 1

    c_func = ida_hexrays.decompile(func_start)
    c_func.refresh_func_ctext()
    func_name = ida_funcs.get_func_name(func_start)
    
    if (api_result := chat_with_AI(str(c_func))):
        if rename_function(func_start, api_result.get("name", func_name)):
            idc.set_func_cmt(func_start,api_result.get("des", "None"), True)
            print(f"[REAI] Renamed: {func_name} -> {api_result['name']}")
    c_func.refresh_func_ctext()


class func_analyze(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
    
    def activate(self, ctx):
        global function_info
        global CodeCheck_flag
        global processed_func
        CodeCheck_flag = False
        ea = idc.get_screen_ea()
        func_start = select_address_check(ea)
        exception_code_check(ea)
        processed_func = []
        get_call_topology(func_start, 0, 0, set(), 0)
        print('[REAI] start AI parser')
        AI_ananalyze(func_start)
        function_info = {}
        print("[REAI] Processing completed!")

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

class exception_code_check_action(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
    
    def activate(self, ctx):
        global exception_code_collection
        global processed_func
        global CodeCheck_flag
        CodeCheck_flag = True
        ea = idc.get_screen_ea()
        select_address_check(ea)

        if exception_code_check(ea) != True:
            print("[REAI] There has exception code")
            print(exception_code_collection)
        else:
            print(f"[REAI] Exception code is not exist, all function num: {len(processed_func)}")     
        exception_code_collection = []
        processed_func = []
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


class call_topology_print(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
    
    def activate(self, ctx):
        global function_info
        global CodeCheck_flag
        global processed_func
        CodeCheck_flag = False
        ea = idc.get_screen_ea()
        func_start = select_address_check(ea)
        exception_code_check(ea)
        processed_func = []
        get_call_topology(func_start, 0, 0, set(), 0)   
        print('[REAI] Print call topology:')
        pt()
        function_info = {}

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


def register_actions():
    action1 = idaapi.action_desc_t(
        "REAI:func_analyze",
        "func_analyze",  # 
        func_analyze(),
        "", 
        "AI work for you!",
        0
    )

    action2 = idaapi.action_desc_t(
        "REAI:exception_code_check",
        "exception_code_check",  # 
        exception_code_check_action(),
        "", 
        "exception code check",
        0
    )

    action3 = idaapi.action_desc_t(
        "REAI:call_topology_print",
        "call_topology_print",  # 
        call_topology_print(),
        "", 
        "call topology print",
        0
    )

    idaapi.register_action(action1)
    idaapi.register_action(action2)
    idaapi.register_action(action3)

def unregister_actions():
    idaapi.unregister_action("REAI:func_analyze")
    idaapi.unregister_action("REAI:exception_code_check") 
    idaapi.unregister_action("REAI:call_topology_print") 

class REAI_PLUGIN(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "REAI_PLUGIN"
    help = "REAI_PLUGIN"
    wanted_name = "REAI"
    def init(self):
        print('[REAI] REAI is loading...')
        register_actions()
        self.menu = ContextMenuHooks()
        self.menu.hook()
        print('[REAI] REAI is ready')
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        print('[REAI] Right click on pseudocode view')
        return

    def term(self):
        if self.menu:
            self.menu.unhook()
        unregister_actions()


def PLUGIN_ENTRY():
    return REAI_PLUGIN()


class ContextMenuHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        # Add actions to the context menu of the Pseudocode view
        if idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE:
            idaapi.attach_action_to_popup(form, popup, "REAI:func_analyze", "REAI/")
            idaapi.attach_action_to_popup(form, popup, "REAI:exception_code_check", "REAI/")
            idaapi.attach_action_to_popup(form, popup, "REAI:call_topology_print", "REAI/")


repository: idapython/src
file: examples/ui/tabular_views/custom/lazy_loaded_chooser.py
description: IDAPython project for Hex-Ray's IDA Pro
wanted_name: Test a lazy-loader chooser in Python
comment: This is a comment.
created_at: 2015-03-14 20:09:27+00:00
pushed_at: 2025-03-13 15:18:11+00:00
url:
https://github.com/idapython/src/blob/af978e79e68af701c45bab401afce5145aa4b80b/examples/ui/tabular_views/custom/lazy_loaded_chooser.py

content:
"""
summary: a lazy-loaded, tree-like data view

description:
  Brings lazy-loading of folders to the tree-like tabular views.

  The important bit to enable this are:

    * ida_kernwin.Choose.OnLazyLoadDir

keywords: chooser, folders, actions

see_also: choose, choose_multi, chooser_with_folders

level: advanced
"""

import idaapi

class Data_t:
    def children(self, path):
        num_children = self._simple_hash(path) & 0x7
        num_folders  = num_children >> 1

        for n in range(num_folders):
            yield True, "Folder %d" % (n + 1)
        for n in range(num_children - num_folders):
            yield False, "Item %d" % (n + 1)

    def _simple_hash(self, path):
        s = 0
        for c in path:
            s += ord(c)
        return s

#-----------------------

# dirspec for the dirtree
class my_dirspec_t(idaapi.dirspec_t):
    def __init__(self):
        idaapi.dirspec_t.__init__(self)
        self.inodes = []
        self.name_index = {}

    def add_entry(self, dirpath, name):
        new_inode = len(self.inodes)
        self.inodes.append(name)
        self.name_index[dirpath, name] = new_inode
        return new_inode

    def get_name(self, inode, flags=0):
        if inode >= 0 and inode < len(self.inodes):
            return self.inodes[inode]

    def get_inode(self, dirpath, name):
        return self.name_index.get((dirpath, name), idaapi.direntry_t.BADIDX)

    def n_inodes(self):
        return len(self.inodes)

    def get_attrs(self, inode):
        pass

    def rename_inode(self, inode, newname):
        pass

    def unlink_inode(self, inode):
        pass

# chooser with dirtree, lazy-loaded
class my_tree_t(idaapi.Choose):
    inherited = idaapi.Choose

    def __init__(self):
        self.dirspec = my_dirspec_t()
        self.dirtree = idaapi.dirtree_t(self.dirspec)

        self.data = Data_t()

        flags = idaapi.CH_TM_FULL_TREE
        columns = [
            ["Name", 10 | idaapi.CHCOL_PLAIN]
        ]

        self.inherited.__init__(self, "My tree", columns, flags)

    def OnGetSize(self):
        return self.dirspec.n_inodes()

    def OnGetLine(self, index):
        inode = self.OnIndexToInode(index)
        return [self.dirspec.get_name(inode)]

    def OnGetDirTree(self):
        return self.dirspec, self.dirtree

    def OnIndexToInode(self, n):
        return n

    # TODO: add "dirtree" parameter when the pywraps are ready
    def OnLazyLoadDir(self, dir_path):
        dirtree = self.dirtree  # should use the parameter when ready

        for is_dir, entry in self.data.children(dir_path):
            if is_dir:
                dirtree.mkdir(entry)
            else:
                inode = self.dirspec.add_entry(dir_path, entry)
                dirtree.link(inode)
        return True

# plugin
class lazy_chooser_plugin_t(idaapi.plugin_t):
    flags = 0
    comment = "This is a comment."
    help = "This is a test."
    wanted_name = "Test a lazy-loader chooser in Python"
    wanted_hotkey = "Alt-Shift-F12"

    def init(self):
        self.tree = my_tree_t()
        print("Registered plugin test-lazy-chooser");
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        self.tree.Show()

    def term(self):
        print("Un-registered plugin test-lazy-chooser");

def PLUGIN_ENTRY():
    return lazy_chooser_plugin_t()


repository: computerline1z/IDAPatcher6.1
file: idapatcher-r1.1.3.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.3.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.3
#
# 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.3"

#--------------------------------------------------------------------------
# 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 # Error
        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 # Error
        if self.add_menu_item_helper("Edit/Patch program/", "Edit selection...", "", 0, self.show_edit_form, None):  
            return 1 # Error
        if self.add_menu_item_helper("Edit/Patch program/", "Fill selection...", "", 0, self.show_fill_form, None):  
            return 1 # Error 
        if self.add_menu_item_helper("Edit/", "Import data...", "Shift-I", 0, self.show_import_form, None):   
            return 1 # Error
        # 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 # OK

    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
        if start_ea != idaapi.BADADDR and end_ea != idaapi.BADADDR:
            if end_ea > start_ea:
                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()
        return

    # 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 = 0x90 #nop

        # 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 = {}

# Dummy function
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 add_menu_items() return 1, then we known that faild to add menu
                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)
        print("[*] Initialize Finished ! Ctrl + U to Refesh Patched Windows")         
        return idaapi.PLUGIN_KEEP

    # NOTE : run() get call and terminate term(), so we 
    #           better don't cause run
    def run(self, arg):
        pass
        #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):
        global idapatcher_manager
        global patch_notify_hook
        if 'patch_notify_hook' in globals():
            patch_notify_hook.unhook()
            del patch_notify_hook
        if 'idapatcher_manager' in globals():
            idapatcher_manager.del_menu_items()
            del idapatcher_manager
        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: clovme/WTools
file: software/Reverse Cracking/IDA Professional 9.0/python/examples/hexrays/vds6.py
wanted_name: Hex-Rays space remover (IDAPython)
comment: Sample plugin6 for Hex-Rays decompiler
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/hexrays/vds6.py

content:
"""
summary: superficially modify the decompilation output

description:
  modifies the decompilation output in a superficial manner,
  by removing some white spaces

  Note: this is rather crude, not quite "pythonic" code.
"""

import idautils
import idc
import ida_idaapi
import ida_hexrays
import ida_lines

do_dbg = False

def dbg(msg):
    if do_dbg:
        print(msg)
    pass

def is_cident_char(c):
    return c.isalnum() or c == '_'

def my_tag_skipcodes(l, storage):
    n = ida_lines.tag_skipcodes(l)
    if n > 0:
        dbg("Skipping %d chars ('%s')" % (n, l[0:n]))
        storage.append(l[0:n])
    return l[n:]

def remove_spaces(sl):

    dbg("*" * 80)
    l = sl.line

    global do_dbg
    # do_dbg = l.find("const char *") > -1

    out = []

    def push(c):
        dbg("Appending '%s'" % c)
        out.append(c)

    # skip initial spaces, do not compress them
    def eat_spaces(l):
        while True:
            l = my_tag_skipcodes(l, out)
            if not l:
                break
            c = l[0]
            if not c.isspace():
                break
            push(c)
            l = l[1:]
        return l
    l = eat_spaces(l)

    # remove all spaces except in string and char constants
    delim = None # if not None, then we are skipping until 'delim'
    last = None # last seen character

    while True:
        dbg("-" * 60)
        dbg("l: '%s'" % l)
        dbg("d: '%s'" % delim)
        dbg("out: '%s'" % out)

        # go until comments
        l = my_tag_skipcodes(l, out)
        if l.startswith("//"):
            push(l)
            break
        nchars = ida_lines.tag_advance(l, 1)
        push(l[0:nchars])
        l = l[nchars:]
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        dbg("c: '%s', last: '%s', l: '%s'" % (c, last, l))
        if delim:
            # we're inside a literal.
            if c == delim:
                delim = None # literal ended
        elif c == '"' or c == "'":
            delim = c # string/char literal started
        elif c.isspace():
            end = l.lstrip()
            nptr = my_tag_skipcodes(end, out)
            dbg("end: '%s', nptr: '%s'" % (end, nptr))
            # do not concatenate idents
            if not is_cident_char(last) or not is_cident_char(nptr[0]):
                l = end
                c = l[0] if l else ''
                dbg("new l: '%s'" % l)
        last = l[0] if l else ''

    sl.line = "".join(out)


class vds6_hooks_t(ida_hexrays.Hexrays_Hooks):
    def func_printed(self, cfunc):
        for sl in cfunc.get_pseudocode():
            remove_spaces(sl);
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays space remover (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin6 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.vds6_hooks = vds6_hooks_t()
            self.vds6_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds6_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: plowsec/idangr
file: ida_angr.py
description: Use angr within IDA
url: https://github.com/plowsec/idangr/blob/a402c8c943cbbaa20335a574ecc576ad4a6f11eb/ida_angr.py
wanted_name: Angr Plugin
comment: Angr integration plugin
created_at: 2024-06-10 08:03:57+00:00
pushed_at: 2024-06-12 08:56:44+00:00
content:
import ida_kernwin
import idaapi
import ida_hexrays
import ida_funcs
import importlib
import sys
import builtins


from ida_angr_lib.log import logger
from ida_angr_lib.globals import simgr
from ida_angr_lib import ida_painter

sys.dont_write_bytecode = True
my_hooks = None

def clear_module_cache(module_name):
    if module_name in sys.modules:
        del sys.modules[module_name]

def paint_all_ea_in_function(ea):

    if not "painter" in builtins.__dict__ or not "aggregated_cov" in builtins.__dict__:
        logger.warning(f"Run an angr exploration first")
        return

    func = ida_funcs.get_func(ea)
    if func is not None:
        
        func_start = func.start_ea
        func_end = func.end_ea
        ea_in_function = [ea for ea in aggregated_cov if func_start <= ea <= func_end]
        logger.debug(f"{len(ea_in_function)} EAs to paint in current function...")
        for ea in ea_in_function:
            ida_painter.color_lines(ea)

builtins.__dict__["paint_all"] = paint_all_ea_in_function

class CustomRightClickHandler(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def activate(self, ctx):
        # Get the current decompiled function
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        if vu:
            ea = vu.cfunc.entry_ea
            logger.debug("EA of the currently decompiled function: %x" % ea)
            logger.debug("Reloading extension code...")
            clear_module_cache('ida_angr_lib.ida_angr_main')
            
            import ida_angr_lib.ida_angr_main as ida_angr_main
            importlib.reload(ida_angr_main)
            ida_angr_main.build_call_state(ea)
        return 1

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

class ExploreFromHereHandler(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def activate(self, ctx):
        # Get the current decompiled function
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        if vu:
            ea = vu.cfunc.entry_ea
            logger.debug("EA of the currently decompiled function: %x" % ea)
            logger.debug("Exploring from here...")
            clear_module_cache('ida_angr_lib.ida_angr_main')
            
            import ida_angr_lib.ida_angr_main as ida_angr_main
            importlib.reload(ida_angr_main)
            ida_angr_main.explore_from_here(ea)
        return 1

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

class ShowCoverageHandler(ida_kernwin.action_handler_t):

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

    def activate(self, ctx):
        
        # Get the current decompiled function
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        if vu:
            func_ea = vu.cfunc.entry_ea
            paint_all_ea_in_function(func_ea)
              
        return 1

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

# Define the actions
ACTION_NAME_BUILD_CALL_STATE = "idangr:right_click:build_call_state"
ACTION_LABEL_BUILD_CALL_STATE = "Build call state"
ACTION_SHORTCUT_BUILD_CALL_STATE = ""

ACTION_NAME_EXPLORE_FROM_HERE = "idangr:right_click:explore_from_here"
ACTION_LABEL_EXPLORE_FROM_HERE = "Explore from here"
ACTION_SHORTCUT_EXPLORE_FROM_HERE = ""

ACTION_NAME_SHOW_COVERAGE = "idangr:right_click:show_cov"
ACTION_LABEL_SHOW_COVERAGE = "Show angr coverage in current function"
ACTION_SHORTCUT_SHOW_COVERAGE = ""

def register_action():
    action_desc_build_call_state = ida_kernwin.action_desc_t(
        ACTION_NAME_BUILD_CALL_STATE,
        ACTION_LABEL_BUILD_CALL_STATE,
        CustomRightClickHandler(),
        ACTION_SHORTCUT_BUILD_CALL_STATE,
        "Build call state for current function",
        -1
    )

    action_desc_explore_from_here = ida_kernwin.action_desc_t(
        ACTION_NAME_EXPLORE_FROM_HERE,
        ACTION_LABEL_EXPLORE_FROM_HERE,
        ExploreFromHereHandler(),
        ACTION_SHORTCUT_EXPLORE_FROM_HERE,
        "Explore from the current function",
        -1
    )

    action_desc_show_coverage = ida_kernwin.action_desc_t(
        ACTION_NAME_SHOW_COVERAGE,
        ACTION_LABEL_SHOW_COVERAGE,
        ShowCoverageHandler(),
        ACTION_SHORTCUT_SHOW_COVERAGE,
        "Show angr coverage in current function",
        -1
    )

    # Register the actions
    res1 = ida_kernwin.register_action(action_desc_build_call_state)
    res2 = ida_kernwin.register_action(action_desc_explore_from_here)
    res3 = ida_kernwin.register_action(action_desc_show_coverage)
    logger.debug(f"Register action result (Build call state): {res1}")
    logger.debug(f"Register action result (Explore from here): {res2}")
    logger.debug(f"Register action result (Show coverage): {res3}")

    return res1 and res2


def unregister_action():
    res1 = ida_kernwin.unregister_action(ACTION_NAME_BUILD_CALL_STATE)
    res2 = ida_kernwin.unregister_action(ACTION_NAME_EXPLORE_FROM_HERE)
    res2 = ida_kernwin.unregister_action(ACTION_NAME_SHOW_COVERAGE)
    logger.debug(f"Unregister action result (Build call state): {res1}")
    logger.debug(f"Unregister action result (Explore from here): {res2}")
    logger.debug(f"Unregister action result (Show coverage): {res3}")


class my_hooks_t(ida_kernwin.UI_Hooks):
    def populating_widget_popup(self, widget, popup):
        widget_type = ida_kernwin.get_widget_type(widget)

        if widget_type == ida_kernwin.BWN_PSEUDOCODE:
            res1 = ida_kernwin.attach_action_to_popup(widget, popup, ACTION_NAME_BUILD_CALL_STATE)
            res2 = ida_kernwin.attach_action_to_popup(widget, popup, ACTION_NAME_EXPLORE_FROM_HERE)
            
            if not ida_kernwin.attach_action_to_popup(widget, popup, ACTION_NAME_SHOW_COVERAGE) or not res1 or not res2:
                logger.error(f"Failed attach_action_to_popup")

class AngrPlugin(idaapi.plugin_t):
    
    flags = idaapi.PLUGIN_KEEP
    comment = "Angr integration plugin"
    help = "Integrates angr with IDA Pro"
    wanted_name = "Angr Plugin"
    wanted_hotkey = "Ctrl-Shift-A"

    def init(self):
        logger.debug("Initializing Angr Plugin")
        if register_action():
            global my_hooks
            my_hooks = my_hooks_t()
            my_hooks.hook()
            logger.debug("UI Hooks set up")
            return idaapi.PLUGIN_OK
        else:
            return idaapi.PLUGIN_SKIP

    def term(self):
        logger.debug("Terminating Angr Plugin")

        # idk why but this is called immediately after init?
        
        """
        if my_hooks:
            my_hooks.unhook()
            logger.debug("UI Hooks removed")
        unregister_action()
        """

    def run(self, arg):
        logger.debug("Angr Plugin is running")

def PLUGIN_ENTRY():
    return AngrPlugin()


