# plugins

repository: the-m3chanic/RE-IDA-Plugins
file: ArrayX/ArrayX.py
wanted_name: ArrayX
comment: Array exporting made easy
created_at: 2024-10-06 15:31:21+00:00
pushed_at: 2024-10-06 15:41:39+00:00
url:
https://github.com/the-m3chanic/RE-IDA-Plugins/blob/3ccadcc23692c43f43facc1d69396a69e2459d0e/ArrayX/ArrayX.py

content:
import idaapi 
import idc 
import ida_kernwin

class ArrayX(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Array exporting made easy"
    help = "Allows for convenient exporting of arrays of various sizes"
    wanted_name = "ArrayX"
    wanted_hotkey = "Ctrl-Shift-E"
    EXPORT_FILE = "array_export.txt"
    EXPORT_TO_FILE = True


    def isValidHexAddr(addr: str) -> bool:
        """ Validates if the given address is a valid hex address """

        try: 
            flag1 = addr.startswith("0x")
        except AttributeError:
            flag1 = False
        
        try:
            flag2 = int(addr, 16)
        except (ValueError, TypeError):
            flag2 = 0
        
        return (flag1 and flag2)

    def convertToHexInt(self, addr: str) -> int:
        """ Converts a hex string to an integer """

        if not ArrayX.isValidHexAddr(addr):
            ida_kernwin.warning(
                "Not a valid hex address!"
            )
            ArrayX.term(self)
        else:
            return int(addr, 16)


    def getArrayAddress():
        """ Prompts the user for the base address of the array """

        baseAddr = ida_kernwin.ask_str(
            "0x0",
            0,
            "Enter the base address of the array you want to export",
        )

        idaapi.msg(f"Base address: {baseAddr}\n")
        baseAddr = ArrayX.convertToHexInt(baseAddr)
        return baseAddr
        
        
    def getArrayType(self):
        """ Prompts the user for the type of the array """

        arrayTypes = ["byte", "word", "dword", "qword", "string", "wideString"]
        arrayType = ida_kernwin.ask_str(
            "Byte",
            0,
            "Enter the type of the array you want to export\n[Byte, Word, Dword, Qword, String, wideString]",
        )

        try: 
            arrayType = arrayType.lower()
        except AttributeError:
            ida_kernwin.warning(
                "Not a valid array type!"
            )
            ArrayX.term(self)
        
        if arrayType not in arrayTypes:
            ida_kernwin.warning(
                "Not a valid array type!"
            )
            ArrayX.term(self)
        

        return arrayType.lower() 

    def getArraySize(self):
        """ Prompts the user for the size of the array """

        arraySize = ida_kernwin.ask_str(
            "0",
            0,
            "Enter the size of the array you want to export",
        )

        try:
            arraySize = int(arraySize)
        except ValueError:
            ida_kernwin.warning(
                "Not a valid array size!"
            )
            ArrayX.term(self)

        return arraySize


    def exportByteArr(baseAddr: int, arraySize: int) -> list:
        """ Exports a byte array """

        array = []
        addr = baseAddr

        for i in range(arraySize):
            x = (idaapi.get_bytes(addr, 1))
            addr += 1
            x = int.from_bytes(x, byteorder="little")
            array.append(x)
        
        return array
    
    def exportWordArr(baseAddr: int, arraySize: int) -> list:
        """ Exports a word array """

        array = []
        addr = baseAddr

        for i in range(arraySize):
            x = (idaapi.get_bytes(addr, 2))
            addr += 2
            x = int.from_bytes(x, byteorder="little")
            array.append(x)
        
        return array

    def exportDwordArr(baseAddr: int, arraySize: int) -> list:
        """ Exports a dword array """

        array = []
        addr = baseAddr

        idaapi.msg(f"address type: {type(addr)}")
        for i in range(arraySize):
            x = (idaapi.get_bytes(addr, 4))
            addr += 4
            x = int.from_bytes(x, byteorder="little")
            array.append(x)
        
        return array
    
    def exportQwordArr(baseAddr: int, arraySize: int) -> list:
        """ Exports a qword array """

        array = []
        addr = baseAddr

        for i in range(arraySize):
            x = (idaapi.get_bytes(addr, 8))
            addr += 8
            x = int.from_bytes(x, byteorder="little")
            array.append(x)
        
        return array
    
    def exportStringArr(baseAddr: int, arraySize: int) -> list:
        """ Exports a string array """

        array = []
        addr = baseAddr

        for i in range(arraySize):
            x = (idaapi.get_bytes(addr, 1))
            addr += 1
            x = int.from_bytes(x, byteorder="little")
            array.append(chr(x))
        
        return array
    
    def exportWideString(baseAddr: int, arrayType: str, arraySize: int) -> list:
        """ Exports a Wide String (as seen on Windows) """

        array = []
        addr = baseAddr 

        for i in range(arraySize):
            x = (idaapi.get_bytes(addr, 1))
            addr += 2
            x = int.from_bytes(x, byteorder="little")
            array.append(chr(x))

        return array
        
    def exportArray(baseAddr: int, arrayType: str, arraySize: int) -> list:
        """ Bottleneck function to export the array """

        array = []
        if arrayType == "byte":
            array = ArrayX.exportByteArr(baseAddr, arraySize)
        elif arrayType == "word":
            array = ArrayX.exportWordArr(baseAddr, arraySize)
        elif arrayType == "dword":
            array = ArrayX.exportDwordArr(baseAddr, arraySize)
        elif arrayType == "qword":
            array = ArrayX.exportQwordArr(baseAddr, arraySize)
        elif arrayType == "string":
            array = ArrayX.exportStringArr(baseAddr, arraySize)
        elif arrayType == "wideString":
            array = ArrayX.exportWideString(baseAddr, arraySize)

        return array
    


    def init(self):
        idaapi.msg("ArrayX initialised. Ctrl-Shift-E to export data\n")
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        idaapi.msg("ArrayX is running...\n")
        baseAddr = ArrayX.getArrayAddress()
        arrayType = ArrayX.getArrayType()
        arraySize = ArrayX.getArraySize()
        
        idaapi.msg(f"Exporting -- Base address: {hex(baseAddr)}, Array type: {arrayType}, Array size: {arraySize}\n")
        
        array = ArrayX.exportArray(baseAddr, arrayType, arraySize)
        idaapi.msg(f"Array: {array}\n")

        if ArrayX.EXPORT_TO_FILE:
            fileName = ida_kernwin.ask_str(
            "Byte",
            0,
            "Enter the name of the file you would like to export the results to: ",
            )
            with open(fileName, "w") as f:
                f.write(f"Base address: {baseAddr}\n")
                f.write(f"Array type: {arrayType}\n")
                f.write(f"Array size: {arraySize}\n")
                f.write(f"Array: {array}\n")
        else: 
            idaapi.msg(f"No request for file write made, exiting normally...")


    def term(self):
        idaapi.msg("ArrayX terminated\n")

def PLUGIN_ENTRY():
    return ArrayX()


repository: wobbier/HEDZ
file: Ghidra/Extensions/IDAPro/Python/7xx/plugins/xml_exporter.py
description: A decompile project of the old game HEDZ
wanted_name: XML Exporter
comment: Export database as XML file
created_at: 2023-06-16 08:51:17+00:00
pushed_at: 2024-10-13 07:48:58+00:00
url:
https://github.com/wobbier/HEDZ/blob/ef3d177053730d30cb54df667d5fa28719b50201/Ghidra/Extensions/IDAPro/Python/7xx/plugins/xml_exporter.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.
##
#---------------------------------------------------------------------
# xmlexp.py - IDA XML Exporter plugin
#---------------------------------------------------------------------
"""
Plugin for IDA which exports a XML PROGRAM document file from a database.
This file must be placed in the IDA plugins directory.
The file idaxml.py must be placed in the IDA python directory.
"""

from __future__ import print_function
import ida_auto
import ida_idaapi
import ida_kernwin
import idaxml
import idc
import sys

if sys.version_info.major >= 3:
    from idaxml import _exc_info
    sys.exc_value = lambda: _exc_info()[1]
    sys.exc_type = lambda: _exc_info()[0]

class XmlExporterPlugin(ida_idaapi.plugin_t):
    """
    XML Exporter plugin class
    """
    flags = 0
    comment = "Export database as XML file"
    help = "Export database as XML <PROGRAM> document"
    wanted_name = "XML Exporter"
    wanted_hotkey = "Ctrl-Shift-x"


    def init(self):
        """
        init function for XML Exporter 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 Exporter plugin.
        
        Args:
            arg: Integer, 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.XmlExporter(arg)
        try:
            try:
                xml.export_xml()
            except idaxml.Cancelled:
                ida_kernwin.hide_wait_box()
                msg = "XML Export cancelled!"
                print("\n" + msg)
                idc.warning(msg)
            except:
                ida_kernwin.hide_wait_box()
                msg = "***** Exception occurred: XML Exporter failed! *****"
                print("\n" + msg + "\n", sys.exc_type, sys.exc_value)
                idc.warning(msg)
        finally:
            xml.cleanup()
            ida_auto.set_ida_state(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlExporterPlugin()


repository: Goatman13/ps2_ida_vu_micro
file: _vu_micro.py
description: Plugin that tries to find and disassembly vu microcode in ps2 executable.
wanted_name: Start VUE Analyze
comment: VU.
created_at: 2022-11-19 19:09:46+00:00
pushed_at: 2025-01-11 15:01:33+00:00
url:
https://github.com/Goatman13/ps2_ida_vu_micro/blob/c5f184bfdd3a762f59c65044f3cdcc7494aed9bc/_vu_micro.py

content:
from ida_bytes import *
from idaapi import *
from idc import *
import idaapi
import ida_bytes
import idc
import binascii

MAKE_CREFS = 1

def get_special_bit(instruction):

	sbits = (instruction >> 27) & 0xF
	string = ""
	if ((sbits >> 3) & 1):
		string += "[E]"
	if ((sbits >> 2) & 1):
		string += "[M]"
	if ((sbits >> 1) & 1):
		string += "[D]"
	if (sbits & 1):
		string += "[T]"

	return string
	
def get_2bit_field(bc):

	if (bc == 0):
		return "x"
	if (bc == 1):
		return "y"
	if (bc == 2):
		return "z"
	if (bc == 3):
		return "w"	

def get_4bit_field(field):

	if (field == 1):
		return "w"
	elif (field == 2):
		return "z"
	elif (field == 4):
		return "y"
	elif (field == 8):
		return "x"
	elif (field == 3):
		return "zw"
	elif (field == 5):
		return "yw"
	elif (field == 6):
		return "yz"
	elif (field == 9):
		return "xw"
	elif (field == 10):
		return "xz"
	elif (field == 12):
		return "xy"
	elif (field == 7):
		return "yzw"
	elif (field == 11):
		return "xzw"
	elif (field == 13):
		return "xyw"			
	elif (field == 14):
		return "xyz"
	elif (field == 15):
		return "xyzw"
	else:
		warning("Opcode is missing field!")
		return "xyzw"

def get_vu1_reg_for_vu0(imm):

	spec_regs = ["status", "mac", "clip", "vi19", "R", "I", "Q", "P", "vi24_reserved", "vi25_reserved",
				"TPC", "vi27_reserved", "vi28_reserved", "vi29_reserved", "vi30_reserved", "vi31_reserved"]
	imm &= 0x3F
	if (imm < 0x20):
		return "vu1_vf{:d}".format(imm)
	elif (imm < 0x30):
		imm -= 0x20
		return "vu1_vi{:d}".format(imm)
	else:
		imm -= 0x30
		return "vu1_" + spec_regs[imm]

def itof(address, instr, dest, source, field, sbits):
	
	field2 = get_4bit_field(field)
	
	while len(instr + "." + field2) < 13:
		field2 += " "
	
	string  = instr + "." + field2 + " vf{:d}, vf{:d}" + sbits
	set_manual_insn(address, string.format(dest, source))

def vu_bc(address, instr, dest, reg1, reg2, field, bc, sbits):
	
	bc2 = get_2bit_field(bc)
	field2 = get_4bit_field(field)
	
	if (dest == 34):
		dest_str = " ACC, "
	else: 
		dest_str = " vf{:d}, "
	
	while len(instr + bc2 + "." + field2) < 13:
		field2 += " "
	
	string  = instr + bc2 + "." + field2 + dest_str + "vf{:d}, vf{:d}" + bc2 + sbits
	if (dest == 34): 
		string  = string.format(reg1, reg2)
	else:
		string  = string.format(dest, reg1, reg2)
	
	set_manual_insn(address, string)

def vu_dr1r2f(address, instr, dest, reg1, reg2, field, sbits):
	
	field2 = get_4bit_field(field)
	
	if (reg2 == 32):
		reg2_str = ", Q"
	elif (reg2 == 33):
		reg2_str = ", I"
	else:
		reg2_str = ", vf{:d}"
		
	if (dest == 34):
		dest_str = " ACC, "
	else: 
		dest_str = " vf{:d}, "
		
	while len(instr + "." + field2) < 13:
		field2 += " "
		
	string  = instr + "." + field2 + dest_str + "vf{:d}" + reg2_str + sbits
	
	if (reg2 >= 32 and dest == 34):
		string  = string.format(reg1)
	elif (reg2 >= 32 and dest != 34):
		string  = string.format(dest, reg1)
	elif (reg2 < 32 and dest == 34):
		string  = string.format(reg1, reg2)
	else:
		string  = string.format(dest, reg1, reg2)
		
	set_manual_insn(address, string)	

def upper(address, instruction):

	if ((instruction >> 31) == 1):
		loi_addr = address - 4
		val = get_dword(loi_addr)
		val = struct.pack('>I', val)
		val = struct.unpack('>f', val)[0]
		val = str(val)
		#print(val)
		string = "loi           " + val
		set_manual_insn(loi_addr, string)

	op = instruction & 0x3F

	if (op <= 0x03):
		addbc(address, instruction)
	elif (op >= 0x04 and op <= 0x07):
		subbc(address, instruction)
	elif (op >= 0x08 and op <= 0x0B):
		maddbc(address, instruction)
	elif (op >= 0x0C and op <= 0x0F):
		msubbc(address, instruction)
	elif (op >= 0x10 and op <= 0x13):
		maxbc(address, instruction)
	elif (op >= 0x14 and op <= 0x17):
		minibc(address, instruction)
	elif (op >= 0x18 and op <= 0x1B):
		mulbc(address, instruction)
	elif (op == 0x1C):
		mulq(address, instruction)
	elif (op == 0x1D):
		maxi(address, instruction)
	elif (op == 0x1E):
		muli(address, instruction)
	elif (op == 0x1F):
		minii(address, instruction)
	elif (op == 0x20):
		addq(address, instruction)
	elif (op == 0x21):
		maddq(address, instruction)
	elif (op == 0x22):
		addi(address, instruction)
	elif (op == 0x23):
		maddi(address, instruction)
	elif (op == 0x24):
		subq(address, instruction)
	elif (op == 0x25):
		msubq(address, instruction)
	elif (op == 0x26):
		subi(address, instruction)
	elif (op == 0x27):
		msubi(address, instruction)
	elif (op == 0x28):
		add(address, instruction)
	elif (op == 0x29):
		madd(address, instruction)
	elif (op == 0x2A):
		mul(address, instruction)
	elif (op == 0x2B):
		_max(address, instruction)
	elif (op == 0x2C):
		sub(address, instruction)
	elif (op == 0x2D):
		msub(address, instruction)
	elif (op == 0x2E):
		opmsub(address, instruction)
	elif (op == 0x2F):
		mini(address, instruction)
	elif (op >= 0x3C and op <= 0x3F):
		upper_special(address, instruction)

def addbc(address, instruction):

	bc = instruction & 0x3
	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "add", dest, source, bc_reg, field, bc, sbits)


def subbc(address, instruction):

	bc = instruction & 0x3
	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "sub", dest, source, bc_reg, field, bc, sbits)

def maddbc(address, instruction):

	bc = instruction & 0x3
	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "madd", dest, source, bc_reg, field, bc, sbits)

def msubbc(address, instruction):

	bc = instruction & 0x3
	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "msub", dest, source, bc_reg, field, bc, sbits)

def maxbc(address, instruction):

	bc = instruction & 0x3
	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "max", dest, source, bc_reg, field, bc, sbits)

def minibc(address, instruction):

	bc = instruction & 0x3
	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "mini", dest, source, bc_reg, field, bc, sbits)

def mulbc(address, instruction):

	bc = instruction & 0x3
	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "mul", dest, source, bc_reg, field, bc, sbits)

def mulq(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "mulq", dest, source, 32, field, sbits)

def maxi(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "maxi", dest, source, 33, field, sbits)

def muli(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "muli", dest, source, 33, field, sbits)

def minii(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "minii", dest, source, 33, field, sbits)

def addq(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "addq", dest, source, 32, field, sbits)

def maddq(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "maddq", dest, source,32, field, sbits)

def addi(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "addi", dest, source, 33, field, sbits)

def maddi(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "maddi", dest, source, 33, field, sbits)

def subq(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "subq", dest, source, 32, field, sbits)

def msubq(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "msubq", dest, source, 32, field, sbits)

def subi(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "subi", dest, source, 33, field, sbits)

def msubi(address, instruction):

	dest = (instruction >> 6) & 0x1F
	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "msubi", dest, source, 33, field, sbits)

def add(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "add", dest, reg1, reg2, field, sbits)

def madd(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "madd", dest, reg1, reg2, field, sbits)

def mul(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "mul", dest, reg1, reg2, field, sbits)

def _max(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "max", dest, reg1, reg2, field, sbits)

def sub(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "sub", dest, reg1, reg2, field, sbits)

def msub(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "msub", dest, reg1, reg2, field, sbits)

def opmsub(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	string = "opmsub.xyz    vf{:d}, vf{:d}, vf{:d}"
	string += get_special_bit(instruction)
	set_manual_insn(address, string.format(dest, reg1, reg2))	

def mini(address, instruction):

	dest = (instruction >> 6) & 0x1F
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "mini", dest, reg1, reg2, field, sbits)
	
##########################################################
def iadd(address, instruction):

	dest = (instruction >> 6) & 0xF
	reg1 = (instruction >> 11) & 0xF
	reg2 = (instruction >> 16) & 0xF
	string = "iadd          vi{:d}, vi{:d}, vi{:d}"
	set_manual_insn(address, string.format(dest, reg1, reg2))	

def isub(address, instruction):

	dest = (instruction >> 6) & 0xF
	reg1 = (instruction >> 11) & 0xF
	reg2 = (instruction >> 16) & 0xF
	string = "isub          vi{:d}, vi{:d}, vi{:d}"
	set_manual_insn(address, string.format(dest, reg1, reg2))	

def iaddi(address, instruction):

	reg1 = (instruction >> 11) & 0xF
	dest = (instruction >> 16) & 0xF
	imm5 = (instruction >> 6) & 0x1F
	sign = ""
	if imm5 > 0xF:
		imm5 = ~imm5
		imm5 &= 0xF
		imm5 += 1
		sign = "-"
	string = "iaddi         vi{:d}, vi{:d}, " + sign + "0x{:X}"
	set_manual_insn(address, string.format(dest, reg1, imm5))	

def iand(address, instruction):

	dest = (instruction >> 6) & 0xF
	reg1 = (instruction >> 11) & 0xF
	reg2 = (instruction >> 16) & 0xF
	string = "iand          vi{:d}, vi{:d}, vi{:d}"
	set_manual_insn(address, string.format(dest, reg1, reg2))

def ior(address, instruction):

	dest = (instruction >> 6) & 0xF
	reg1 = (instruction >> 11) & 0xF
	reg2 = (instruction >> 16) & 0xF
	string = "ior           vi{:d}, vi{:d}, vi{:d}"
	set_manual_insn(address, string.format(dest, reg1, reg2))	

def upper_special(address, instruction):

	op = (instruction & 0x3) | ((instruction >> 4) & 0x7C)
	
	if (op <= 0x03):
		addabc(address, instruction)
	elif (op >= 0x04 and op <= 0x07):
		subabc(address, instruction)
	elif (op >= 0x08 and op <= 0x0B):
		maddabc(address, instruction)
	elif (op >= 0x0C and op <= 0x0F):
		msubabc(address, instruction)
	elif (op == 0x10):
		itof0(address, instruction)
	elif (op == 0x11):
		itof4(address, instruction)
	elif (op == 0x12):
		itof12(address, instruction)
	elif (op == 0x13):
		itof15(address, instruction)
	elif (op == 0x14):
		ftoi0(address, instruction)
	elif (op == 0x15):
		ftoi4(address, instruction)
	elif (op == 0x16):
		ftoi12(address, instruction)
	elif (op == 0x17):
		ftoi15(address, instruction)
	elif (op >= 0x18 and op <= 0x1B):
		mulabc(address, instruction)
	elif (op == 0x1C):
		mulaq(address, instruction)
	elif (op == 0x1D):
		_abs(address, instruction)
	elif (op == 0x1E):
		mulai(address, instruction)
	elif (op == 0x1F):
		clip(address, instruction)
	elif (op == 0x20):
		addaq(address, instruction)
	elif (op == 0x21):
		maddaq(address, instruction)
	elif (op == 0x22):
		addai(address, instruction)
	elif (op == 0x23):
		maddai(address, instruction)
	elif (op == 0x25):
		msubaq(address, instruction)
	elif (op == 0x26):
		subai(address, instruction)
	elif (op == 0x27):
		msubai(address, instruction)
	elif (op == 0x28):
		adda(address, instruction)
	elif (op == 0x29):
		madda(address, instruction)
	elif (op == 0x2A):
		mula(address, instruction)
	elif (op == 0x2C):
		suba(address, instruction)
	elif (op == 0x2D):
		msuba(address, instruction)
	elif (op == 0x2E):
		opmula(address, instruction)
	elif (op == 0x2F):
		nop(address,instruction)
		

def lower(address, instruction):

	# Early check for I bit to prevent
	# branch code from creating xref.
	is_Ibit = get_dword(address + 4)
	if ((is_Ibit >> 31) == 1):
		return

	if (instruction == 0x8000033C):
		set_manual_insn(address,"nop")
		return
	if (instruction & (1 << 31)):
		lower1(address, instruction)
	else:
		lower2(address, instruction)

def lower1(address, instruction):
	
	op = (instruction & 0x3F)
	
	if (op == 0x30):
		iadd(address, instruction)
	elif (op == 0x31):
		isub(address, instruction)
	elif (op == 0x32):
		iaddi(address, instruction)
	elif (op == 0x34):
		iand(address, instruction)
	elif (op == 0x35):
		ior(address, instruction)
	elif (op >= 0x3C and op <= 0x3F):
		lower1_special(address, instruction)
	else:
		print("[unknown lower1]")

def lower1_special(address, instruction):	

	op = (instruction & 0x3) | ((instruction >> 4) & 0x7C)
	
	if (op == 0x30):
		move(address, instruction)
	elif (op == 0x31):
		mr32(address, instruction)
	elif (op == 0x34):
		lqi(address, instruction)
	elif (op == 0x35):
		sqi(address, instruction)
	elif (op == 0x36):
		lqd(address, instruction)
	elif (op == 0x37):
		sqd(address, instruction)
	elif (op == 0x38):
		div(address, instruction)
	elif (op == 0x39):
		sqrt(address, instruction)
	elif (op == 0x3A):
		rsqrt(address, instruction)
	elif (op == 0x3B):
		waitq(address, instruction)
	elif (op == 0x3C):
		mtir(address, instruction)
	elif (op == 0x3D):
		mfir(address, instruction)
	elif (op == 0x3E):
		ilwr(address, instruction)
	elif (op == 0x3F):
		iswr(address, instruction)
	elif (op == 0x40):
		rnext(address, instruction)
	elif (op == 0x41):
		rget(address, instruction)
	elif (op == 0x42):
		rinit(address, instruction)
	elif (op == 0x43):
		rxor(address, instruction)
	elif (op == 0x64):
		mfp(address, instruction)
	elif (op == 0x68):
		xtop(address, instruction)
	elif (op == 0x69):
		xitop(address, instruction)
	elif (op == 0x6C):
		xgkick(address, instruction)
	elif (op == 0x70):
		esadd(address, instruction)
	elif (op == 0x71):
		ersadd(address, instruction)
	elif (op == 0x72):
		eleng(address, instruction)
	elif (op == 0x73):
		erleng(address, instruction)
	elif (op == 0x74): 
		eatanxy(address, instruction)
	elif (op == 0x75):
		eatanxz(address, instruction)
	elif (op == 0x76):
		esum(address, instruction)
	elif (op == 0x78):
		esqrt(address, instruction)
	elif (op == 0x79):
		ersqrt(address, instruction)
	elif (op == 0x7A):
		ercpr(address, instruction)
	elif (op == 0x7B):
		set_manual_insn(address, "waitp")
	elif (op == 0x7C):
		esin(address, instruction)
	elif (op == 0x7D):
		eatan(address, instruction)
	elif (op == 0x7E):
		eexp(address, instruction)
	else:
		print("Bad op cop2_special2 0x{:X}".format(op))

def mfp(address, instruction):

	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	
	while len("mfp." + field2) < 13:
		field2 += " "	
	
	string = "mfp." + field2 + " vf{:d}, P"
	set_manual_insn(address, string.format(dest))


def xtop(address, instruction):

	it = (instruction >> 16) & 0x1F
	string = "xtop          vi{:d}"
	set_manual_insn(address, string.format(it))


def xitop(address, instruction):

	it = (instruction >> 16) & 0x1F
	string = "xitop         vi{:d}"
	set_manual_insn(address, string.format(it))


def xgkick(address, instruction):

	_is = (instruction >> 11) & 0x1F
	string = "xgkick        vi{:d}"
	set_manual_insn(address, string.format(_is))


def esadd(address, instruction):

	source = (instruction >> 11) & 0x1F
	string = "esadd         P, vf{:d}"
	set_manual_insn(address, string.format(source))


def ersadd(address, instruction):

	source = (instruction >> 11) & 0x1F
	string = "ersadd        P, vf{:d}"
	set_manual_insn(address, string.format(source))


def eleng(address, instruction):

	source = (instruction >> 11) & 0x1F
	string = "eleng         P, vf{:d}"
	set_manual_insn(address, string.format(source))


def esum(address, instruction):

	source = (instruction >> 11) & 0x1F
	string = "esum          P, vf{:d}"
	set_manual_insn(address, string.format(source))


def ercpr(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "ercpr         P, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))


def erleng(address, instruction):

	source = (instruction >> 11) & 0x1F
	string = "erleng        P, vf{:d}"
	set_manual_insn(address, string.format(source))


def esqrt(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "esqrt         P, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))


def ersqrt(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "ersqrt        P, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))


def esin(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "esin          P, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))


def eatan(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "eatan         P, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))


def eatanxy(address, instruction):

	source = (instruction >> 11) & 0x1F
	string = "eatanxy       P, vf{:d}"
	set_manual_insn(address, string.format(source))


def eatanxz(address, instruction):

	source = (instruction >> 11) & 0x1F
	string = "eatanxz       P, vf{:d}"
	set_manual_insn(address, string.format(source))


def eexp(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "eexp          P, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))


def lower2(address, instruction):

	op = ((instruction >> 25) & 0x7F)
	
	if (op == 0x00):
		lq(address, instruction)
	elif (op  == 0x01):
		sq(address, instruction)
	elif (op  == 0x04):
		loadstore_imm(address, "ilw", instruction)
	elif (op  == 0x05):
		loadstore_imm(address, "isw", instruction)
	elif (op  == 0x08):
		arithu(address, "iaddiu", instruction)
	elif (op  == 0x09):
		arithu(address, "isubiu", instruction)
	elif (op  == 0x10):
		fceq(address, instruction)
	elif (op  == 0x11):
		fcset(address, instruction)
	elif (op  == 0x12):
		fcand(address, instruction)
	elif (op  == 0x13):
		fcor(address, instruction)
	elif (op  == 0x14):
		fseq(address, instruction)
	elif (op  == 0x15):
		fsset(address, instruction)
	elif (op  == 0x16):
		fsand(address, instruction)
	elif (op  == 0x17):
		fsor(address, instruction)
	elif (op  == 0x18):
		fmeq(address, instruction)
	elif (op  == 0x1A):
		fmand(address, instruction)
	elif (op  == 0x1B):
		fmor(address, instruction)
	elif (op  == 0x1C):
		fcget(address, instruction)
	elif (op  == 0x20):
		b(address, instruction)
	elif (op  == 0x21):
		bal(address, instruction)
	elif (op  == 0x24):
		jr(address, instruction)
	elif (op  == 0x25):
		jalr(address, instruction)
	elif (op  == 0x28):
		branch(address, "ibeq", instruction)
	elif (op  == 0x29):
		branch(address, "ibne", instruction)
	elif (op  == 0x2C):
		branch_zero(address, "ibltz", instruction)
	elif (op  == 0x2D):
		branch_zero(address, "ibgtz", instruction)
	elif (op  == 0x2E):
		branch_zero(address, "iblez", instruction)
	elif (op  == 0x2F):
		branch_zero(address, "ibgez", instruction)
	else:
		print("[unknown lower2]")
	


def loadstore_imm(address, string, instruction):

	sign = ""
	_is = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	imm = instruction & 0x7FF
	if (imm > 0x3FF and _is != 0):
		imm = ~imm
		imm &= 0x3FF
		imm += 1
		sign = "-"
	elif (imm > 0x3FF and _is == 0):
		vu1_reg = get_vu1_reg_for_vu0(imm)
		while len(string + "." + field2) < 13:
			field2 += " "	
		string = string + "." + field2 + " vi{:d}, " + vu1_reg

		#set_cmt(address, "load store " + vu1_reg, 0)
		set_manual_insn(address, string.format(it))
		return
	imm *= 16
	while len(string + "." + field2) < 13:
		field2 += " "	
	
	string = string + "." + field2 + " vi{:d}, " + sign + "0x{:X}(vi{:d})"
	set_manual_insn(address, string.format(it, imm, _is))


def arithu(address, string, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	imm = instruction & 0x7FF
	imm |= ((instruction >> 21) & 0xF) << 11
	while len(string) < 13:
		string += " "
	
	string = string + " vi{:d}, vi{:d}, 0x{:X}"
	set_manual_insn(address, string.format(dest, source, imm))


def branch(address, string, instruction):

	imm = instruction & 0x7FF
	_is = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0x1F
	if (imm > 0x3FF):
		imm &= 0x3FF
		imm = ~imm
		imm &= 0x3FF
		imm *= 8
		addr = (address - imm)
		addr -= skip_vif_data(addr, address)
		
		while len(string) < 13:
			string += " "
		
		string = string + " vi{:d}, vi{:d}, 0x{:X}"
		set_manual_insn(address, string.format(it, _is, addr))
		if(MAKE_CREFS == 1):
			add_cref(address, addr, fl_JN | XREF_USER)
		return
		
	imm *= 8
	addr = address + imm + 8
	addr += skip_vif_data(address, addr)
	if(MAKE_CREFS == 1):
		add_cref(address, addr, fl_JN | XREF_USER)
	while len(string) < 13:
		string += " "	
	
	string = string + " vi{:d}, vi{:d}, 0x{:X}"
	set_manual_insn(address, string.format(it, _is, addr))


def branch_zero(address, string, instruction):

	imm = instruction & 0x7FF
	reg = (instruction >> 11) & 0x1F
	if (imm > 0x3FF):
		imm &= 0x3FF
		imm = ~imm
		imm &= 0x3FF
		imm *= 8
		addr = (address - imm) #+ 8
		addr -= skip_vif_data(addr, address)
	
		while len(string) < 13:
			string += " "
		
		string = string + " vi{:d}, 0x{:X}"
		set_manual_insn(address, string.format(reg, addr))
		if(MAKE_CREFS == 1):
			add_cref(address, addr, fl_JN | XREF_USER)
		return
	
	imm *= 8
	addr = address + imm + 8
	addr += skip_vif_data(address, addr)
	if(MAKE_CREFS == 1):
		add_cref(address, addr, fl_JN | XREF_USER)	
	while len(string) < 13:
		string += " "
	
	string = string + " vi{:d}, 0x{:X}"
	set_manual_insn(address, string.format(reg, addr))

def b(address, instruction):

	imm = instruction & 0x7FF
	if (imm > 0x3FF):
		imm &= 0x3FF
		imm = ~imm
		imm &= 0x3FF
		imm *= 8
		addr = (address - imm)
		addr -= skip_vif_data(addr, address)
		if(MAKE_CREFS == 1):
			add_cref(address, addr, fl_JN | XREF_USER)
		string = "b             0x{:X}"
		set_manual_insn(address, string.format(addr))
		return
	
	imm *= 8
	addr = address + imm + 8
	addr += skip_vif_data(address, addr)
	if(MAKE_CREFS == 1):
		add_cref(address, addr, fl_JN | XREF_USER)	
	string = "b             0x{:X}"
	set_manual_insn(address, string.format(addr))


def bal(address, instruction):

	imm = instruction & 0x7FF
	link_reg = (instruction >> 16) & 0x1F
	if (imm > 0x3FF):
		imm &= 0x3FF
		imm = ~imm
		imm &= 0x3FF
		imm *= 8
		addr = (address - imm)
		addr -= skip_vif_data(addr, address)
		if(MAKE_CREFS == 1):
			add_cref(address, addr, fl_CN | XREF_USER)		
		string = "bal           vi{:d} 0x{:X}"
		set_manual_insn(address, string.format(link_reg, addr))
		return
	
	imm *= 8	
	addr = address + imm + 8
	addr += skip_vif_data(address, addr)
	if(MAKE_CREFS == 1):
		add_cref(address, addr, fl_JN | XREF_USER)		
	string = "bal           vi{:d} 0x{:X}"
	set_manual_insn(address, string.format(link_reg, addr))


def jr(address, instruction):

	addr_reg = (instruction >> 11) & 0x1F
	#if(MAKE_CREFS == 1):
	#	add_cref(address, addr, fl_JN | XREF_USER)
	string = "jr            vi{:d}"
	set_manual_insn(address, string.format(addr_reg))


def jalr(address, instruction):

	addr_reg = (instruction >> 11) & 0x1F
	link_reg = (instruction >> 16) & 0x1F
	#if(MAKE_CREFS == 1):
	#	add_cref(address, addr, fl_CN | XREF_USER)
	string = "jalr          vi{:d}, vi{:d}"
	set_manual_insn(address, string.format(link_reg, addr_reg))


def lq(address, instruction):

	sign = ""
	_is = (instruction >> 11) & 0x1F
	ft = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)	
	imm = instruction & 0x7FF
	if (imm > 0x3FF and _is != 0):
		imm = ~imm
		imm &= 0x3FF
		imm += 1
		sign = "-"
	elif (imm > 0x3FF and _is == 0):
		vu1_reg = get_vu1_reg_for_vu0(imm)
		while len("lq." + field2) < 13:
			field2 += " "
		string = "lq." + field2 + " vf{:d}, " + vu1_reg
		set_manual_insn(address, string.format(ft))
		#set_cmt(address, "download from " + vu1_reg, 0)
		return
	imm *= 16 

	while len("lq." + field2) < 13:
		field2 += " "
		
	string = "lq." + field2 + " vf{:d}, " + sign + "0x{:X}(vi{:d})"
	set_manual_insn(address, string.format(ft, imm, _is))


def sq(address, instruction):

	sign = ""
	it = (instruction >> 16) & 0x1F
	fs = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	imm = instruction & 0x7FF
	if (imm > 0x3FF and it != 0):
		imm = ~imm
		imm &= 0x3FF
		imm += 1
		sign = "-"
	elif (imm > 0x3FF and it == 0):
		vu1_reg = get_vu1_reg_for_vu0(imm)
		while len("sq." + field2) < 13:
			field2 += " "
		string = "sq." + field2 + " vf{:d}, " + vu1_reg
		set_manual_insn(address, string.format(fs))
		#set_cmt(address, "upload to " + vu1_reg, 0)
		return
	imm *= 16 

	while len("sq." + field2) < 13:
		field2 += " "
		
	string = "sq." + field2 + " vf{:d}, " + sign + "0x{:X}(vi{:d})"
	set_manual_insn(address, string.format(fs, imm, it))


def fceq(address, instruction):

	imm = instruction & 0xFFFFFF
	string = "fceq          vi1, 0x{:X}"
	set_manual_insn(address, string.format(imm))


def fcset(address, instruction):

	imm = instruction & 0xFFFFFF
	string = "fcset         0x{:X}"
	set_manual_insn(address, string.format(imm))


def fcand(address, instruction):

	imm = instruction & 0xFFFFFF
	string = "fcand         vi1, 0x{:X}"
	set_manual_insn(address, string.format(imm))


def fcor(address, instruction):

	imm = instruction & 0xFFFFFF
	string = "fcor          vi1, 0x{:X}"
	set_manual_insn(address, string.format(imm))


def fseq(address, instruction):

	imm = ((instruction >> 10) & 0x800) | (instruction & 0x7FF)
	dest = (instruction >> 16) & 0x1F
	string = "fseq          vi{:d}, 0x{:X}"
	set_manual_insn(address, string.format(dest, imm))


def fsset(address, instruction):

	imm = ((instruction >> 10) & 0x800) | (instruction & 0x7FF)
	string = "fsset         0x{:X}"
	set_manual_insn(address, string.format(imm))


def fsand(address, instruction):

	imm = ((instruction >> 10) & 0x800) | (instruction & 0x7FF)
	dest = (instruction >> 16) & 0x1F
	string = "fsand         vi{:d}, 0x{:X}"
	set_manual_insn(address, string.format(dest, imm))


def fsor(address, instruction):

	imm = ((instruction >> 10) & 0x800) | (instruction & 0x7FF)
	dest = (instruction >> 16) & 0x1F
	string = "fsor          vi{:d}, 0x{:X}"
	set_manual_insn(address, string.format(dest, imm))


def fmeq(address, instruction):

	_is = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0x1F
	string = "fmeq          vi{:d}, vi{:d}"
	set_manual_insn(address, string.format(it, _is))


def fmand(address, instruction):

	_is = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0x1F
	string = "fmand         vi{:d}, vi{:d}"	
	set_manual_insn(address, string.format(it, _is))


def fmor(address, instruction):

	_is = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0x1F
	string = "fmor          vi{:d}, vi{:d}"	
	set_manual_insn(address, string.format(it, _is))


def fcget(address, instruction):

	it = (instruction >> 16) & 0x1F
	string = "fcget         vi{:d}"
	set_manual_insn(address, string.format(it))


def new_skip_vif_data(start, end):
	
	to_skip = 0
	addr = start
	while (start < end):
		if start < ida_idaapi.BADADDR:
			if get_dword(addr + 4) >> 31 == 1: #loi
				addr += 4
				continue
			instruction = get_dword(addr)
			instruction >>= 24
			instruction &= 0x7F
			if instruction == 0:
				end += 4
			elif instruction > 0x5F: # unpack
				end += 4
			elif instruction == 0x4A:
				prev_ins = get_dword(addr - 4) >> 24
				next_ins = get_dword(addr + 4) >> 24
				if prev_ins != 0 or prev_ins < 0x60 or next_ins != 0 or next_ins < 0x60:
					continue
				
				


def skip_vif_data(start, end):
	
	to_skip = 0

	# Check for addresses in VU range.
	# Better avoid misdetection early.
	if start < 0x4000 or start >= 0x11000000:
		return to_skip

	while (start < end):
		if start < ida_idaapi.BADADDR:
			if ((is_code(ida_bytes.get_flags(start)) == 1)):
				print("wrong branch address, wrap around VU mem?")
				return to_skip
			instruction = get_dword(start)
			end_instr = get_dword(end)
			# add check for loi here, 0.0 is not that rare
			while end_instr == 0x70000000 or end_instr == 0:
				end += 4
				end_instr = get_dword(end)
			if instruction == 0x70000000 or instruction == 0 and start < end:
				if get_dword(start + 4) >> 31 == 1:# and get_manual_insn(start)[0:3] == "loi":
				#if get_manual_insn(start) == "loi           0.0":
					start += 4
					continue
				while (get_dword(start) >> 24) != 0x4A and start < end:
					start += 4
					to_skip += 4
					if ((is_code(ida_bytes.get_flags(start)) == 1)):
						print("wrong branch address, wrap?")
						return to_skip
				to_skip += 4 # Skip that 4A too.
				start += 4
				#to_skip += 4
				continue
	
		start += 4

	return to_skip

#####################################################
def addabc(address, instruction):

	bc = instruction & 0x3
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "adda", 34, source, bc_reg, field, bc, sbits)

def subabc(address, instruction):

	bc = instruction & 0x3
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "suba", 34, source, bc_reg, field, bc, sbits)

def maddabc(address, instruction):

	bc = instruction & 0x3
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "madda", 34, source, bc_reg, field, bc, sbits)

def msubabc(address, instruction):

	bc = instruction & 0x3
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "msuba", 34, source, bc_reg, field, bc, sbits)

def itof0(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "itof0", dest, source, field, sbits)

def itof4(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "itof4", dest, source, field, sbits)

def itof12(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "itof12", dest, source, field, sbits)

def itof15(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "itof15", dest, source, field, sbits)

def ftoi0(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "ftoi0", dest, source, field, sbits)

def ftoi4(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "ftoi4", dest, source, field, sbits)

def ftoi12(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "ftoi12", dest, source, field, sbits)

def ftoi15(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "ftoi15", dest, source, field, sbits)

def mulabc(address, instruction):

	bc = instruction & 0x3
	source = (instruction >> 11) & 0x1F
	bc_reg = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_bc(address, "mula", 34, source, bc_reg, field, bc, sbits)

def mulaq(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "mulaq", 34, source, 32, field, sbits)

def _abs(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	itof(address, "abs", dest, source, field, sbits)

def mulai(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "mulai", 34, source, 33, field, sbits)

def clip(address, instruction):

	instr = "clipw."
	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field = get_4bit_field(field)
	if field != "xyz":
		print("mVU WARNING: Clip instruction with malfolmed dest at 0x{:X}!".format(address))
	while len(instr + field) < 13:
		field += " "
	string = instr + field + " vf{:d}, vf{:d}w"
	set_manual_insn(address, string.format(reg1, reg2))	

def addaq(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "addaq", 34, source, 32, field, sbits)

def maddaq(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "maddaq", 34, source, 32, field, sbits)

def maddai(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "maddai", 34, source, 33, field, sbits)

def msubaq(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "msubaq", 34, source, 32, field, sbits)

def subai(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "subai", 34, source, 33, field, sbits)

def msubai(address, instruction):

	source = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "msubai", 34, source, 33, field, sbits)

def adda(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "adda", 34, reg1, reg2, field, sbits)

def addai(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "addai", 34, reg1, 33, field, sbits)

def madda(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "madda", 34, reg1, reg2, field, sbits)

def mula(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "mula", 34, reg1, reg2, field, sbits)

def suba(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "suba", 34, reg1, reg2, field, sbits)

def msuba(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	sbits = get_special_bit(instruction)
	vu_dr1r2f(address, "msuba", 34, reg1, reg2, field, sbits)

def opmula(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	string = "opmula.xyz    ACC, vf{:d}, vf{:d}"
	set_manual_insn(address, string.format(reg1, reg2))

def nop(address, instruction):

	sbits = get_special_bit(instruction)
	string = "nop           " + sbits
	set_manual_insn(address, string)

def move(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	itof(address, "move", dest, source, field, "")

def mr32(address, instruction):

	source = (instruction >> 11) & 0x1F
	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	itof(address, "mr32", dest, source, field, "")

def lqi(address, instruction):

	_is = (instruction >> 11) & 0xF
	ft = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	
	while len("lqi." + field2) < 13:
		field2 += " "	
	
	string = "lqi." + field2 + " vf{:d}, (vi{:d}++)"
	set_manual_insn(address, string.format(ft, _is))

def sqi(address, instruction):

	fs = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0xF
	dest_field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(dest_field)
	
	while len("sqi." + field2) < 13:
		field2 += " "	
	
	string = "sqi." + field2 + " vf{:d}, (vi{:d}++)"
	set_manual_insn(address, string.format(fs, it))

def lqd(address, instruction):

	_is = (instruction >> 11) & 0xF
	ft = (instruction >> 16) & 0x1F
	dest_field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(dest_field)
	
	while len("lqd." + field2) < 13:
		field2 += " "	
	
	string = "lqd." + field2 + " vf{:d}, (--vi{:d})"
	set_manual_insn(address, string.format(ft, _is))

def sqd(address, instruction):

	fs = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0xF
	dest_field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(dest_field)
	
	while len("sqd." + field2) < 13:
		field2 += " "	
	
	string = "sqd." + field2 + " vf{:d}, (--vi{:d})"
	set_manual_insn(address, string.format(fs, it))

def div(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	fsf = (instruction >> 21) & 0x3
	ftf = (instruction >> 23) & 0x3	
	fsf2 = get_2bit_field(fsf)
	ftf2 = get_2bit_field(ftf)
	string = "div           Q, vf{:d}" + fsf2 + " vf{:d}" + ftf2
	
	set_manual_insn(address, string.format(reg1, reg2))

def sqrt(address, instruction):

	source = (instruction >> 16) & 0x1F
	ftf = (instruction >> 23) & 0x3	
	ftf2 = get_2bit_field(ftf)
	string = "sqrt          Q, vf{:d}" + ftf2

	set_manual_insn(address, string.format(source))

def rsqrt(address, instruction):

	reg1 = (instruction >> 11) & 0x1F
	reg2 = (instruction >> 16) & 0x1F
	fsf = (instruction >> 21) & 0x3
	ftf = (instruction >> 23) & 0x3
	fsf2 = get_2bit_field(fsf)
	ftf2 = get_2bit_field(ftf)
	string = "rsqrt         Q, vf{:d}" + fsf2 + " vf{:d}" + ftf2
	
	set_manual_insn(address, string.format(reg1, reg2))

def waitq(address, instruction):

	set_manual_insn(address, "waitq")

def mtir(address, instruction):

	fs = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0xF
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "mtir          vi{:d}, vf{:d}" + fsf2
	set_manual_insn(address, string.format(it, fs))

def mfir(address, instruction):

	_is = (instruction >> 11) & 0x1F
	ft = (instruction >> 16) & 0x1F
	dest_field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(dest_field)
	
	while len("mfir." + field2) < 13:
		field2 += " "	
	
	string = "mfir." + field2 + " vf{:d}, vi{:d}"
	set_manual_insn(address, string.format(ft, _is))

def ilwr(address, instruction):

	_is = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	
	while len("ilwr." + field2) < 13:
		field2 += " "		
	
	string = "ilwr." + field2 + " vi{:d}, (vi{:d})" + field2
	set_manual_insn(address, string.format(it, _is))

def iswr(address, instruction):

	_is = (instruction >> 11) & 0x1F
	it = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	
	while len("iswr." + field2) < 13:
		field2 += " "		
	
	string = "iswr." + field2 + " vi{:d}, (vi{:d})" + field2
	set_manual_insn(address, string.format(it, _is))

def rnext(address, instruction):

	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	
	while len("rnext." + field2) < 13:
		field2 += " "		
	
	string = "rnext." + field2 + " vf{:d}, R"
	set_manual_insn(address, string.format(dest))

def rget(address, instruction):

	dest = (instruction >> 16) & 0x1F
	field = (instruction >> 21) & 0xF
	field2 = get_4bit_field(field)
	
	while len("rget." + field2) < 13:
		field2 += " "			
	
	string = "rget." + field2 + " vf{:d}, R"
	set_manual_insn(address, string.format(dest))

def rinit(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "rinit         R, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))

def rxor(address, instruction):

	source = (instruction >> 11) & 0x1F
	fsf = (instruction >> 21) & 0x3
	fsf2 = get_2bit_field(fsf)
	string = "rxor          R, vf{:d}." + fsf2
	set_manual_insn(address, string.format(source))

def calculate_mpg_size(start):
	print("Entering calculate_mpg_size")
	address = start
	size = 0
	while 1:
		if (get_dword(address) >> 24) == 0x4A:
			add_size = (get_dword(address) >> 16) & 0xFF
			if add_size == 0:
				add_size = 0x100
			add_size *= 8
			size += add_size
			size += 8
			#address = address + size
			#print(hex(size))
			#print(hex(address))
			address = start + size
			continue
		else:
			stopper = address + 0x50
			counter = 0
			while address < stopper:
				if ((get_dword(address) >> 24) & 0xFF) == 0x4A:
					size += counter
					#print(hex(address))
					#print("2")
					break
				else:
					if ((get_dword(address) >> 24) & 0xFF) == 0x4A: 
						tpr = get_dword(address) >> 16 & 0xFF
						#print("tpr %x", tpr)
						counter += 4
						break
					address += 4
					counter += 4
					#print(hex(address))
					#print("3")
					if address == stopper - 4:
						#print(hex(address))
						#print("4")
						return size

	#return size
			
def mark_code(address, end):
	
	print("Entering mark_code")
	counter = 0
	while (address < end):
		if address < ida_idaapi.BADADDR:	
			if ((address & 3) != 0):
				address = ((address & 0xFFFFFFFC) + 4)
				continue
			
			#print(hex(address))
			#Try detect DMA and VIF code and skip it.
			#At the same time we can comment it for readability.
			previous_op = get_dword(address - 4)
			instruction = get_dword(address)
			old = address
			if (instruction >> 24) == 0x70 or (instruction >> 24) == 0x60 and instruction & 0xFFF != 0x2FF:
				TAG_TYPE = (instruction >> 28) & 7
				if TAG_TYPE == 7:
					TAG_TYPE = "END"
				elif TAG_TYPE == 6:
					TAG_TYPE = "RET"
				else:
					TAG_TYPE = "UNK" # Unreachable for now.
				del_items(address)
				create_dword(address)
				del_items(address + 4)
				create_dword(address + 4)
				dma_qwc = (instruction & 0xFFFF)
				dma_addr = get_dword(address + 4)
				set_cmt(address, "DMA TAG      " + TAG_TYPE + " QWC=0x{:X}".format(dma_qwc), 0)
				set_cmt(address + 4, "DMA TAG      ADDR=0x{:X}".format(dma_addr), 0)
				address += 8
				#Skip as non VU code everything until VIF MPG is found,
				#or until we reach end address.
				while (get_dword(address) >> 24) != 0x4A and address <= end:
					create_dword(address)
					instruction = get_dword(address)
					vif_code = (instruction >> 24) & 0x7F
					if vif_code == 0x00:
						vif_cmd = "VIF NOP"
					elif vif_code == 0x01:
						cycle_cl = instruction & 0xF
						cycle_wl = (instruction >> 8) & 0xF
						vif_cmd = "VIF STCYCL   WL=0x{:X}, CL=0x{:X}".format(cycle_wl, cycle_cl)
					elif vif_code == 0x02:
						vif_ofst = instruction & 0x3FF
						vif_cmd = "VIF OFFSET   OFST=0x{:X}".format(vif_ofst)
					elif vif_code == 0x03:
						vif_base = instruction & 0x3FF
						vif_cmd = "VIF BASE     BASE=0x{:X}".format(vif_base)
					elif vif_code == 0x04:
						vif_itop = instruction & 0x3FF
						vif_cmd = "VIF ITOP     ITOP=0x{:X}".format(vif_itop)
					elif vif_code == 0x05:
						stmod_mode = instruction & 0x3
						vif_cmd = "VIF STMOD    MODE=0x{:X}".format(stmod_mode)
					elif vif_code == 0x06:
						vif_cmd = "VIF MSKPATH3"
					elif vif_code == 0x07:
						vif_cmd = "VIF MARK"
					elif vif_code == 0x10:
						#Likely not DMA TAG
						if instruction == 0x10000000:
							vif_cmd = "VIF FLUSHE"
						#likely DMA TAG
						else:
							dma_qwc = (instruction & 0xFFFF) 
							set_cmt(address, "DMA TAG      CNT QWC=0x{:X}".format(dma_qwc), 0)
							address += 4
							dma_addr = get_dword(address)
							vif_cmd = "DMA TAG      ADDR=0x{:X}".format(dma_addr)
					elif vif_code == 0x11:
						vif_cmd = "VIF FLUSH"
					elif vif_code == 0x13:
						vif_cmd = "VIF FLUSHA"
					elif vif_code == 0x14:
						vif_cmd = "VIF MSCAL"
					elif vif_code == 0x15:
						vif_cmd = "VIF MSCALF"
					elif vif_code == 0x17:
						vif_cmd = "VIF MSCNT"
					elif vif_code == 0x20:
						set_cmt(address, "VIF STMASK", 0)
						address += 4
						stmask_data = get_dword(address)
						vif_cmd = "VIF STMASK   DATA=0x{:08X}".format(stmask_data)
					elif vif_code == 0x30:
						vif_cmd = "VIF STROW"
					elif vif_code == 0x31:
						vif_cmd = "VIF STCOL"
					elif vif_code == 0x50:
						vif_cmd = "VIF DIRECT"
					elif vif_code == 0x51:
						vif_cmd = "VIF DIRECTHL"
					else:
						vif_cmd = ""
					set_cmt(address, vif_cmd, 0)
					address += 4
				continue
				
			if instruction == 0 :
				create_dword(address)
				address += 4
				continue
			
			#Todo check jalr
			if (instruction >> 24) == 0x4A and previous_op == 0:
				create_dword(address)
				vif_size = (instruction >> 16) & 0xFF
				if vif_size == 0:
					vif_size = 0x100
				vif_size *= 8
				vif_addr = (instruction & 0xFFFF) * 8
				vif_addr &= 0x3FF8 #Wrong for VU0, but better than nothing. 
				comment = "VIF MPG      SIZE=0x{:X}, ADDRESS=0x{:X}"
				set_cmt(address, comment.format(vif_size, vif_addr), 0)
				create_dword(address - 4)
				set_cmt(address - 4, "VIF NOP", 0)				
				address += 4
				continue
			
			del_items(address)
			create_dword(address)
			if ((address & 4) == 4):
				set_color(address, CIC_ITEM, 0x383838)
				#print(hex(color))
				upper(address, instruction)
			else:
				set_color(address, CIC_ITEM, 0x2B2B2B)
				lower(address, instruction)
			
		address += 4
	return 0

def vu_single_line():
	
	start_addr = read_selection_start()
	end_addr = read_selection_end()
	if(start_addr == BADADDR):
		start_addr = get_screen_ea();
		end_addr = start_addr + 4;
		
	mark_code(start_addr, end_addr)

def vu_mpg_4A():
	
	ea = get_screen_ea()
	if (get_dword(ea) >> 24) == 0x4A:	
		end = ea + calculate_mpg_size(ea)
		print(hex(ea))
		print(hex(end))
		mark_code(ea, end)
		
	else:
		print("To start plugin you need to specify line with VIF MPG command")

class ActionHandler(idaapi.action_handler_t):

    def __init__(self, callback):
        
        idaapi.action_handler_t.__init__(self)
        self.callback = callback
    
    def activate(self, ctx):

        self.callback()
        return 1

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

def register_actions():   

    actions = [
        {
            'id': 'start:vu_single_line',
            'name': 'Disassemble Marked VU Code',
            'hotkey': 'F10',
            'comment': 'Disassemble Marked VU Code',
            'callback': vu_single_line,
            'menu_location': ''
        },
        {
            'id': 'start:vu_mpg_4A',
            'name': 'Disassemble VU Code using VIF MPG',
            'hotkey': 'Alt-Shift-5',
            'comment': 'Disassemble VU Code using VIF MPG',
            'callback': vu_mpg_4A,
            'menu_location': ''
        }
    ]


    for action in actions:

        if not idaapi.register_action(idaapi.action_desc_t(
            action['id'], # Must be the unique item
            action['name'], # The name the user sees
            ActionHandler(action['callback']), # The function to call
            action['hotkey'], # A shortcut, if any (optional)
            action['comment'] # A comment, if any (optional)
        )):

            print('Failed to register ' + action['id'])

        #if not idaapi.attach_action_to_menu(
        #    action['menu_location'], # The menu location
        #    action['id'], # The unique function ID
        #    0):
		#
        #    print('Failed to attach to menu '+ action['id'])

class vu_helper_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_HIDE
	comment = "VU. "
	help = "Analyze VU"
	wanted_name = "Start VUE Analyze"
	wanted_hotkey = "Alt-Shift-5"

	def init(self):
		if (idaapi.ph.id == idaapi.PLFM_MIPS and ida_ida.inf_get_procname() == 'r5900l'):
			register_actions()
			idaapi.msg("Analyze VU loaded.\n")
			return idaapi.PLUGIN_KEEP

		return idaapi.PLUGIN_SKIP
	
	def run(self, arg):
		idaapi.msg("Analyze VU run.\n")
	
	def term(self):
		pass

def PLUGIN_ENTRY():
	return vu_helper_t()


repository: CrackerCat/MemoryDump
file: MemoryDump.py
description: memory dump Plugin for ida pro 7.0 7.1
wanted_name: MemoryDump
comment: Memory Dump for IDA Pro 7.0 and 7.1
created_at: 2018-04-23 03:38:39+00:00
pushed_at: 2018-04-20 13:46:05+00:00
url:
https://github.com/CrackerCat/MemoryDump/blob/b02277b910d496f06e3524d598aec6594db8b490/MemoryDump.py

content:
# -*- coding: UTF-8 -*-
import idaapi


class MemoryDumpForm(idaapi.Form):
    """Simple Form to test  and combo box controls"""
    def __init__(self):
        idaapi.Form.__init__(self, r"""STARTITEM 0
        MemoryDump
        {FormChangeCb}
        Please Input Addr:
        <#Hint1#StartAddr  :{StartAddress}>
        <#Hint2#EndAddr/Len:{EndAddress}>
        <##Option##Len:{rLen}>
        <EndAddr:{rEndAddr}>{cGroup2}>
        """, {
            'StartAddress': idaapi.Form.StringInput(width=50, swidth=15),
            'EndAddress': idaapi.Form.StringInput(width=50, swidth=15),
            'FormChangeCb': idaapi.Form.FormChangeCb(self.OnFormChange),
            'cGroup2': idaapi.Form.RadGroupControl(("rLen", "rEndAddr")),
        })

    def OnFormChange(self, fid):
        if fid == -2:
            # print "start save"
            self.start = self.GetControlValue(self.StartAddress)
            self.endorlen = self.GetControlValue(self.EndAddress)
            self.dumptype = self.GetControlValue(self.cGroup2)
            if len(self.start) == 0 or len(self.endorlen) == 0:
                idaapi.warning("addr or len is null")
                return -1
            else:
                self.StartDump()
        return 1

    def StartDump(self):
        # print self.start
        # print self.endorlen
        self.filepath = idaapi.ask_file(1, "*.dump", "save dump file")
        if self.dumptype == 0:
            ea = self.getHexNum(self.start)
            len = self.getHexNum(self.endorlen)
            if not idaapi.is_loaded(ea) or not idaapi.is_loaded(ea + len):
                idaapi.warning("arrary is out of bound")
                return -1
            if len <= 0:
                idaapi.warning("len is <= 0")
                return -1
            print("start read bytes")
            self.Close(0)
            idaapi.show_wait_box("read bytes")
            self.memdata = idaapi.get_many_bytes(ea, len)
            print("read bytes end")
            idaapi.hide_wait_box("read end")
        elif self.dumptype == 1:
            ea = self.getHexNum(self.start)
            len = self.getHexNum(self.endorlen) - self.getHexNum(self.start)
            if not idaapi.is_loaded(ea) or not idaapi.is_loaded(ea + len):
                idaapi.warning("arrary is out of bound")
                return -1
            if len <= 0:
                idaapi.warning("len is <= 0")
                return -1
            print("start read bytes")
            self.Close(0)
            idaapi.show_wait_box("read bytes")
            self.memdata = idaapi.get_many_bytes(ea, len)
            print("read bytes end")
            idaapi.hide_wait_box("read end")
        fp = open(self.filepath, 'wb')
        fp.write(self.memdata)
        fp.close()
        idaapi.msg("save:" + self.filepath)
        return 1

    def getHexNum(self, nums):
        return long(nums, 16)


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

    def activate(self, ctx):
        #print "start show"
        form = MemoryDumpForm()
        form.Compile()
        form.Execute()
        return 1

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


class MemoryDump(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX | idaapi.PLUGIN_HIDE
    comment = "Memory Dump for IDA Pro 7.0 and 7.1"
    help = "Memory Dump"
    wanted_name = "MemoryDump"
    wanted_hotkey = ""

    def init(self):
        idaapi.msg("Ida plugin init called.\n")
        idaapi.register_action(
            idaapi.action_desc_t("dump:memoryDump", "MemoryDump", memory_dump_handle(), "Alt+D", "", -1))
        return idaapi.PLUGIN_KEEP

    def term(self):
        idaapi.unregister_action("dump:memoryDump")
        idaapi.msg("term was called \n")

    def run(self):
        idaapi.msg("run was called \n")
        pass


def PLUGIN_ENTRY():
    return MemoryDump()


repository: techbliss/Python_editor
file: 7.2/plugins/Code editor/template/Plugin_temp.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/Code%20editor/template/Plugin_temp.py

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


PLUGIN_VERSION = "1.4"
IDAVERISONS    = "IDA PRO 7.0+"
AUTHORS        = "Storm Shadow"
DATE           = "2017"
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(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    # Run editor when invoked.
    def activate(self, ctx):
        g = globals()
        idahome = idaapi.idadir("plugins\\Code editor")
        IDAPython_ExecScript(idahome + "\\pyeditor.py", g)

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

class ripeye(idaapi.plugin_t):
    flags = 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 = idaapi.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)
            idaapi.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
        idaapi.register_action(action_desc)

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

        form = idaapi.get_current_tform()
        idaapi.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 = idaapi.BADADDR
        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            form = idaapi.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
        idaapi.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: acoderly/gaiya
file: ida_plugin/GaiyaPlugin.py
description: Gaiya is a toolkit to extract C&C from elf.
wanted_name: GaiyaSigGen
created_at: 2021-04-23 08:17:28+00:00
pushed_at: 2021-07-29 01:53:06+00:00
url:
https://github.com/acoderly/gaiya/blob/389aeb8dad929fcc460db1defd2daf71572ce306/ida_plugin/GaiyaPlugin.py

content:
import idaapi
import idc
import time
from idaapi import PluginForm
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import Qt


class ScrollWidget(QtWidgets.QWidget):

    def __init__(self, parent=None, frame=QtWidgets.QFrame.Box):
        super(ScrollWidget, self).__init__()

        # Container Widget
        widget = QtWidgets.QWidget()
        # Layout of Container Widget
        self.layout = QtWidgets.QVBoxLayout(self)
        self.layout.setContentsMargins(0, 0, 0, 0)
        widget.setLayout(self.layout)

        scroll = QtWidgets.QScrollArea()
        scroll.setFrameShape(frame)
        scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
        scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        scroll.setWidgetResizable(True)
        scroll.setWidget(widget)

        scroll_layout = QtWidgets.QVBoxLayout(self)
        scroll_layout.addWidget(scroll)
        scroll_layout.setContentsMargins(0, 0, 0, 0)
        self.setLayout(scroll_layout)

    def addWidget(self, widget):
        self.layout.addWidget(widget)

    def addLayout(self, layout):
        self.layout.addLayout(layout)


class GaiyaPluginFormClass(PluginForm):
    def __init__(self):
        super(GaiyaPluginFormClass, self).__init__()
        self.parent = None
        self.sig_pattern = """
import idc
import idaapi
import idautils

# func_sig_dump
{FUNC_SIG_DUMP}

# code_sig_dump
{CODE_SIG_DUMP}

def sig_parser(ea_list):
    print("[>>>>>>] sig_parser is called.",[hex(i) for i in ea_list])
    print("[>>>>>>] Function name is:",idc.GetFunctionName(ea_list[0]))

    current_file = idaapi.get_root_filename()
    return False, current_file, None, None


sig = {BEGIN}"name": "{SIG_NAME}",
       "md5": "{FILE_NAME}",
       "date": "{DATE}",
       "code_sig_str": "{CODE_SIG}",
       "code_sig_start_ea": "{CODE_SIG_START_EA}",
       "code_sig_end_ea": "{CODE_SIG_END_EA}",
       "func_sig_str": "{FUNC_SIG}",
       "func_sig_start_ea": "{FUNC_SIG_START_EA}",
       "func_sig_end_ea": "{FUNC_SIG_END_EA}",
       "parser": sig_parser{END}
{HELPER_CODE_STR}
        """
        self.HELPER_CODE_STR = """
####################################################################################################################
####################################################################################################################
####################################################################################################################
g_sigs = [sig, ]
def get_target_func(sig):
    sig_name = sig["name"]
    func_sig = sig["func_sig_str"]
    func_ea = sig["func_sig_start_ea"]
    func_len = len(func_sig.split("_"))
    func_sig = "".join(func_sig.split("_"))
    print("[******] The sig_name:<{}> func_sig:<{}> func_ea:<{}>".format(sig_name, func_sig, func_ea))

    target_func = []
    for func_ea in idautils.Functions():
        flags = idc.GetFunctionFlags(func_ea)
        if flags & FUNC_LIB or flags & FUNC_THUNK:
            continue
        dism_addr = list(idautils.FuncItems(func_ea))
        for idx in range(len(dism_addr)):
            try:
                dis_str = "".join(map(idc.GetMnem, dism_addr[idx:idx + func_len]))
                if dis_str == func_sig:
                    target_func.append(func_ea)
                    break
            except IndexError:
                break

    return target_func


def get_target_code(sig, target_func):
    file_name = idaapi.get_root_filename()
    # logger.info("[>>>>>>] process {} start".format(file_name))
    sig_name = sig["name"]
    code_sig = sig["code_sig_str"]
    code_len = len(code_sig.split("_"))
    code_ea = sig["code_sig_start_ea"]
    code_sig = "".join(code_sig.split("_"))
    print("[******] The sig_name:<{}> code_sig:<{}> code_ea:<{}>".format(sig_name, code_sig, code_ea))
    ret = False
    print("[******] get_target_code param target_func:{}".format(target_func))
    for func_ea in target_func:
        dism_addr = list(idautils.FuncItems(func_ea))
        for idx in range(len(dism_addr)):
            try:
                dism_str = "".join(map(idc.GetMnem, dism_addr[idx:idx + code_len]))
                if dism_str == code_sig:
                    print("[******] get_target_code matched.")
                    parser = sig["parser"]
                    if parser != None:
                        ret, current_file, ip, port = parser(dism_addr[idx:idx + code_len])
                        if ret is not False:
                            print("{}_{}_{}".format(current_file.split(".")[0], ip, port))
                        # TODO How to log the result.
            except IndexError:
                break
    # logger.info("[>>>>>>] process {} end".format(file_name))
    return ret


def parser():
    for sig in g_sigs:
        target_func = get_target_func(sig)
        print("[******] target_func:{}".format(target_func))
        for item in target_func:
            print("[******] target_func_name:{}".format(idc.GetFunctionName(item)))
        ret = get_target_code(sig, target_func)
        # TODO Need log fail stuff.
        if ret is True:
            print("[******] process success.")
        else:
            print("[******] process failed.")


idaapi.autoWait()
parser()
"""

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

        self.form = form
        self.parent = PluginForm.FormToPyQtWidget(form)
        self.populate_form()

    def OnClose(self, form):
        """
        Called when the widget is closed
        """
        pass

    def populate_form(self):
        layout = self.view_configuration_info()
        widget = QtWidgets.QWidget()
        widget.setLayout(layout)
        outer_layout = QtWidgets.QHBoxLayout()
        outer_layout.addWidget(widget)
        self.parent.setLayout(outer_layout)

    def Show(self, caption, options=0):
        options |= PluginForm.WOPN_CENTERED
        super(GaiyaPluginFormClass, self).Show(caption, options)

    def gaiya_layout(self, outer_layout):
        groupbox = QtWidgets.QGroupBox()
        vbox = QtWidgets.QVBoxLayout(groupbox)

        grid_layout = QtWidgets.QGridLayout(groupbox)
        vbox.addLayout(grid_layout)

        self.sig_name_le = QtWidgets.QLineEdit()

        self.func_start_ea_le = QtWidgets.QLineEdit()
        self.func_end_ea_le = QtWidgets.QLineEdit()

        self.code_start_ea_le = QtWidgets.QLineEdit()
        self.code_end_ea_le = QtWidgets.QLineEdit()

        self.output_le = QtWidgets.QTextEdit()

        layout = QtWidgets.QHBoxLayout()
        # self.server_message = QtWidgets.QLabel()
        # layout.addWidget(self.server_message)
        layout.addStretch()
        vbox.addSpacing(20)
        vbox.addLayout(layout)

        grid_layout.addWidget(QtWidgets.QLabel('Sig Name:'), 0, 0)
        grid_layout.addWidget(self.sig_name_le, 0, 1)
        grid_layout.addWidget(QtWidgets.QLabel('Target Function Start EA:'), 1, 0)
        grid_layout.addWidget(self.func_start_ea_le, 1, 1)
        grid_layout.addWidget(QtWidgets.QLabel('Target Function End EA:'), 2, 0)
        grid_layout.addWidget(self.func_end_ea_le, 2, 1)
        grid_layout.addWidget(QtWidgets.QLabel('Target Code Start EA:'), 3, 0)
        grid_layout.addWidget(self.code_start_ea_le, 3, 1)
        grid_layout.addWidget(QtWidgets.QLabel('Target Code End EA:'), 4, 0)
        grid_layout.addWidget(self.code_end_ea_le, 4, 1)
        grid_layout.addWidget(QtWidgets.QLabel('Output:'), 5, 0)
        grid_layout.addWidget(self.output_le, 5, 1)

        grid_layout.setColumnMinimumWidth(0, 75)
        grid_layout.setSpacing(10)
        grid_layout.setContentsMargins(10, 10, 10, 10)

        outer_layout.addWidget(groupbox)

    def view_configuration_info(self):
        container = QtWidgets.QVBoxLayout()

        layout = QtWidgets.QHBoxLayout()
        self.message = QtWidgets.QLabel()
        layout.addWidget(self.message)
        layout.addStretch()
        gen_button = QtWidgets.QPushButton('Gen')
        layout.addWidget(gen_button)
        gen_button.clicked.connect(self.gen_output)

        scroll_layout = ScrollWidget(frame=QtWidgets.QFrame.NoFrame)
        self.gaiya_layout(scroll_layout)
        container.addWidget(scroll_layout)
        container.addLayout(layout)

        return container

    def gen_sig(self):

        self.func_start_ea = self.func_start_ea_le.text()
        self.func_end_ea = self.func_end_ea_le.text()
        self.code_start_ea = self.code_start_ea_le.text()
        self.code_end_ea = self.code_end_ea_le.text()
        res = "func_start_ea:{}\nfunc_end_ea:{}\ncode_start_ea:{}\ncode_end_ea:{}\n".format(
            self.func_start_ea, self.func_end_ea, self.code_start_ea, self.code_end_ea
        )
        print("gen_sig button is clicked {}".format(res))
        self.output_le.setText(res)

    def get_sig(self, start_ea, end_ea):
        ea = start_ea
        end = end_ea
        sig = []
        while ea < end:
            sig.append(idc.GetMnem(ea))
            ea = ea + idc.get_item_size(ea)
        return "_".join(sig)

    def get_disasm(self, start_ea, end_ea):
        ea = start_ea
        end = end_ea
        dis = []
        while ea < end:
            dis.append(idc.GetDisasm(ea))
            ea = ea + idc.get_item_size(ea)
        tmp = []
        for idx in range(0, len(dis)):
            item = "# {} {}".format(idx, dis[idx])
            tmp.append(item)
        return "\n".join(tmp)

    def gen_output(self):
        self.sig_name = self.sig_name_le.text()
        self.func_start_ea = int(self.func_start_ea_le.text(), base=16)
        self.func_end_ea = int(self.func_end_ea_le.text(), base=16)
        self.code_start_ea = int(self.code_start_ea_le.text(), base=16)
        self.code_end_ea = int(self.code_end_ea_le.text(), base=16)
        self.output = self.sig_pattern.format(BEGIN="{",
                                              FUNC_SIG_DUMP=self.get_disasm(self.func_start_ea,
                                                                            self.func_end_ea),
                                              CODE_SIG_DUMP=self.get_disasm(self.code_start_ea,
                                                                            self.code_end_ea),
                                              SIG_NAME=self.sig_name + "_" + time.strftime('%Y-%m-%d_%H',
                                                                                           time.localtime(time.time())),
                                              FILE_NAME=idaapi.get_root_filename(),
                                              DATE=time.strftime('%Y-%m-%d_%H:%M:%S', time.localtime(time.time())),
                                              CODE_SIG=self.get_sig(self.code_start_ea, self.code_end_ea),
                                              CODE_SIG_START_EA=hex(self.code_start_ea),
                                              CODE_SIG_END_EA=hex(self.code_end_ea),
                                              FUNC_SIG=self.get_sig(self.func_start_ea, self.func_end_ea),
                                              FUNC_SIG_START_EA=hex(self.func_start_ea),
                                              FUNC_SIG_END_EA=hex(self.func_end_ea),
                                              END="}",
                                              HELPER_CODE_STR=self.HELPER_CODE_STR)
        self.output_le.setText(self.output)


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

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

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

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

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

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

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


    class Gen(Gen_Menu_Context):
        def activate(self, ctx):
            self.plugin.run()
            return 1
except:
    pass

p_initialized = False


class gaiyasigen_t(idaapi.plugin_t):  # pragma: no cover
    comment = ""
    help = "todo"
    wanted_name = "GaiyaSigGen"
    wanted_hotkey = "Ctrl-Alt-G"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        global p_initialized

        # register popup menu handlers
        try:
            Gen.register(self, "GaiyaSigGen")
        except:
            pass

        if p_initialized is False:
            p_initialized = True
            idaapi.register_action(idaapi.action_desc_t(
                "GaiyaSigGen",
                "Gaiya sig generator",
                Gen(),
                None,
                None,
                0))
            idaapi.attach_action_to_menu("Gen", "GaiyaSigGen", idaapi.SETMENU_APP)

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        plg = GaiyaPluginFormClass()
        plg.Show("Gaiya Sig Generator")


def PLUGIN_ENTRY():  # pragma: no cover
    return gaiyasigen_t()
# /Applications/IDAPro7.0/ida64.app/Contents/MacOS/plugins/GaiyaPlugin.py
# /Applications/IDAPro7.0/ida.app/Contents/MacOS/plugins/GaiyaPlugin.py


repository: m00dy-/IDAPlugins
file: amoco/AmocoPlugin.py
description: Various plugins for IDA pro
wanted_name: Amoco IDA Plugin
comment: Amoco is really cool !
created_at: 2015-10-07 09:03:42+00:00
pushed_at: 2015-10-09 08:14:55+00:00
url:
https://github.com/m00dy-/IDAPlugins/blob/ffc966c3f07e820711355c4b4cb38c3854711438/amoco/AmocoPlugin.py

content:
# -*- coding: utf-8 -*-
# version 0: POC 

'''
Plugin IDA qui une fois copié dans le repertoire
plugins d'IDA est accessible depuis le menu 
Edit --> plugins 
'''

from idaapi import *
import amoco

class MyGraph(GraphViewer):
	def __init__(self, flow_graph,title):
		GraphViewer.__init__(self, title)
		print '[+] Flow Graph (GraphViewer Class) init'
		self._title = title
		self.SetCurrentRendererType(idaapi.TCCRT_GRAPH)
		self.flow_graph = flow_graph
		#self.result = result
		#self.names = {}
        
	def OnRefresh(self):
		print 'refresh'
		self.Clear()
		addr_in = {}
		test = idaapi.COLSTR("test0	test2	",idaapi.SCOLOR_INSN)
		self.AddNode(test)
		for i in self.flow_graph:
			id = self.AddNode(i)
			
		return True


	def OnSelect(self,node_id):
		print("OnSelect Event .")
		#self.Select(node_id)
		return True

	def OnClick(self,node_id):
		print("OnClick() Event .")
		#print self[node_id]
		print self
		print self._title
		frm = idaapi.find_tform(self._title)
		print frm
		#self.Select(node_id)
		return True

	def OnGetText(self, node_id):
		#self.Refresh()
		return str(self[node_id])

	def OnCommand(self,cmd_id):
		if self.cmd_test == cmd_id:
			print("Amoco : Symbolic form ")
                        # access flow graph blocks and print their symbolic form
                        fl = list() 
                        for node in self.flow_graph:
                            idaapi.msg(" Symbolic form is \n %s \n"%node.map)
                            fl.append(node.map)
                        
                        for i in fl:
                            self.flow_graph.append(i)
                        self.Refresh() 

			return
		if self.cmd_test2 == cmd_id:
			print("Amoco Menu :  Item2")
			return
		print("command :%s",cmd_id)

	def Show(self):
	
		if not GraphViewer.Show(self):
			return False
		self.cmd_test  = self.AddCommand("Amoco Menu :  Symbolic form","F2")
		self.cmd_test2 = self.AddCommand("Amoco Menu :  Item2","F3")

		if self.cmd_test == 0 :
			print("Failed to add popup menu Item1 !")
		if self.cmd_test2 == 0:
			print("Failed to add popup menu Item2 !")
		return True




#Configuration Handler
class HTooltipC(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        idaapi.warning("Configuration : Tooltip-providing action triggered")
        return 1

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


#Analyse Handler 
class HTooltipDisassemble(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        idaapi.warning("Disassemble : Tooltip-providing action triggered")
        idaapi.msg(" Disassemble handler %s \n"%type(ctx))
        idaapi.msg(" Disassemble handler current ea %s \n"%hex(ctx.cur_ea).rstrip('L'))
        amocoform = idaapi.get_current_tform() 
        idaapi.msg(" Disassemble current form is %s\n"%idaapi.get_tform_title(amocoform))
        idaapi.msg(" Disassemble current form type is %s\n"%idaapi.get_tform_type(amocoform))
        highlighted_item = "None"
        #highlighted_item = idaapi.get_highlighted_identifier() 

        if (idaapi.get_tform_title(amocoform) == "IDA View Amoco Plugin" ):
            if ( idaapi.get_tform_type(amocoform) == 29 ) : 
                highlighted_item = idaapi.get_highlighted_identifier() 

        idaapi.msg(" Disassemble current selected item %s\n"% highlighted_item) 

        # get the path the analyzed file to pass it to amoco 
        filename =  GetInputFilePath() 
        pe  =  amoco.system.loader.load_program(filename)

        # get start and end selection 
        start_selection  = hex(SelStart()) 
        end_selection = hex(SelEnd())
        idaapi.msg(" Start selection  : %s - End selection : %s\n"%(start_selection,end_selection) )  
        
        # apply lbackward algo 
        z = amoco.lbackward(pe)

        # get selected block 
        block =  z.getblock(SelStart())
        idaapi.msg(" Amoco disassembly is \n %s \n"% block)
 
        # create a new graph view for Amoco 
        flow_graph = list()
        flow_graph.append(block)
        g = MyGraph(flow_graph,"Amoco Flow Graph")
        g.Show()
  
  

        return 1

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



class Hooks(idaapi.UI_Hooks):
    def populating_tform_popup(self, form, popup):
        idaapi.msg("Hooks called\n")
        idaapi.attach_action_to_popup(form, popup, "my:tooltip0", "Amoco/Config/", idaapi.SETMENU_APP)
        idaapi.attach_action_to_popup(form, popup, "my:tooltip1", "Amoco/Analyse/", idaapi.SETMENU_APP)


class amoco_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = "Amoco is really cool ! "

    help = "This is help"
    wanted_name = "Amoco IDA Plugin"
    wanted_hotkey = "Alt-F8"
    hooks = None 

    def init(self):
        idaapi.msg("init() called!\n")
        self.hooks = Hooks()
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        idaapi.msg("run() called with %d!\n" % arg)
        tooltip_act0 = idaapi.action_desc_t('my:tooltip0', 'Amoco : Configuration ', HTooltipC(), '', 'This is a tooltip-providing action tooltip', -1)
        tooltip_act1 = idaapi.action_desc_t('my:tooltip1', 'Amoco : Disassemble ', HTooltipDisassemble(), '', 'This is yet another a tooltip-providing action tooltip', -1)
        idaapi.register_action(tooltip_act0)
        idaapi.register_action(tooltip_act1)
        idaapi.attach_action_to_menu("View/", 'my:tooltip0', idaapi.SETMENU_APP)
        idaapi.attach_action_to_menu("View/", 'my:tooltip1', idaapi.SETMENU_APP)

        # we give the focus to IDA View-A 
        ida_view_a_form =  idaapi.find_tform("IDA View-A")
        idaapi.switchto_tform(ida_view_a_form,1)
        idaapi.open_disasm_window("Amoco Plugin ")

        #hooks = Hooks()
        idaapi.msg("Calling Hooks\n")
        self.hooks.hook()

    def term(self):
        idaapi.msg("term() called!\n")
        if self.hooks is not None :
            self.hooks.unhook()
            self.hooks =  None 
    
        #TO DO unhook actions 

def PLUGIN_ENTRY():
    return amoco_plugin_t()


repository: 1398440232/FridaStudy
file: ida AIanalyse/WPeChatGPT.py
created_at: 2024-09-03 11:08:00+00:00
pushed_at: 2024-09-04 07:56:25+00:00
url:
https://github.com/1398440232/FridaStudy/blob/811e089b3f7b1478d2a1c35ab17a617a1001adc6/ida%20AIanalyse/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
# Plugin information, you can change the model here.
PLUGIN_NAME = 'WPeChatGPT'
PROD_NAME = 'ChatGPT'
MODEL = 'qwen-plus'
# Set your API key here, or put in in the OPENAI_API_KEY environment variable.
openai_api_key = "https://dashscope.aliyuncs.com/compatible-mode/v1"
openai.api_key = "sk-3d9aab13f1314xxxxxxxxxxxxxxxxxx"
# Set your forward-proxy if necessary. (e.g. Clash = http://127.0.0.1:7890)
proxy = ""
# Set reverse-proxy URL if you need. (e.g. Azure OpenAI)
proxy_address = ""
# Create openai client (python openai package version > 1.2)
if proxy:
    client = openai.OpenAI(http_client=httpx.Client(proxies=proxy, transport=httpx.HTTPTransport(local_address="0.0.0.0")), api_key=openai_api_key)
    print("WPeChatGPT has appointed the forward-proxy.")
elif proxy_address:
    client = openai.OpenAI(base_url=proxy_address, api_key=openai_api_key)
    print("WPeChatGPT has appointed the reverse-proxy.")
else:
    client = openai.OpenAI(base_url=openai_api_key,api_key="sk-3d9aab13f1314xxxxxxxxxxxxxxxxxx")


# 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 "---GPT_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语言伪代码函数，分别分析该函数的预期目的、参数的作用、详细功能，最后取一个新的函数名字。（用简体中文回答我，并且回答开始前加上'---GPT_START---'字符串结束后加上'---GPT_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 '---GPT_START---' string before the beginning of the answer and add '---GPT_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 "---GPT_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代码进行还原。（回答开始前加上'---GPT_Python_START---'字符串结束后加上'---GPT_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 '---GPT_Python_START---' string before the beginning of the answer and add '---GPT_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 "---GPT_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语言伪代码函数的漏洞并提出可能的利用方法。（用简体中文回答我，并且回答开始前加上'---GPT_VulnFinder_START---'字符串结束后加上'---GPT_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 '---GPT_VulnFinder_START---' the end of the string plus '---GPT_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 "---GPT_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构造代码来利用下面函数中的漏洞。（用简体中文回答我，并且回答开始前加上'---GPT_VulnPython_START---'字符串结束后加上'---GPT_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 '---GPT_VulnPython_START---' the end of the string plus '---GPT_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")
            print(str(e))
        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.5 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.5 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():
    if openai.api_key == "ENTER_OPEN_API_KEY_HERE":
        openai.api_key = os.getenv("OPENAI_API_KEY")
        if not openai.api_key:
            print("未找到 API_KEY，请在脚本中填写 openai.api_key! :(@WPeace")
            raise ValueError("No valid OpenAI API key found")
    return type(f"myplugin_{PLUGIN_NAME}", (myplugin_WPeChatGPT, ), dict())()


repository: vient/ida_movbe_lifter
file: movbe.py
url: https://github.com/vient/ida_movbe_lifter/blob/fa4b98c7c4a213d92a94112c179713bc998f47ce/movbe.py
wanted_name: MOVBE
comment: MOVBE support for the Hex-Rays x64 Decompiler
created_at: 2021-12-31 16:39:15+00:00
pushed_at: 2022-01-06 19:28:42+00:00
content:
import ida_allins
import ida_hexrays
import ida_ida
import ida_idaapi
import ida_idp
import ida_kernwin
import ida_loader
import ida_typeinf
import ida_ua


def get_reg_size(op):
  sizes = {
    ida_ua.dt_word: 2,
    ida_ua.dt_dword: 4,
    ida_ua.dt_qword: 8,
  }
  return sizes[op.dtype]


def get_helper_name(size):
  names = {
    2: '_byteswap_ushort',
    4: '_byteswap_ulong',
    8: '_byteswap_uint64',
  }
  return names[size]


def get_type_info(size):
  types = {
    2: ida_typeinf.BTF_UINT16,
    4: ida_typeinf.BTF_UINT32,
    8: ida_typeinf.BTF_UINT64,
  }
  return ida_typeinf.tinfo_t(types[size])


class MOVBELifter(ida_hexrays.microcode_filter_t):
  def __init__(self):
    super().__init__()

  def install(self):
    ida_hexrays.install_microcode_filter(self, True)
    print("MOVBE lifter installed")

  def remove(self):
    ida_hexrays.install_microcode_filter(self, False)
    print("MOVBE lifter removed")

  def match(self, cdg):
    return cdg.insn.itype == ida_allins.NN_movbe

  def apply(self, cdg):
    insn = cdg.insn
    to_reg = insn.Op1.type == ida_ua.o_reg
    if to_reg:
      size = get_reg_size(insn.Op1)
      arg = cdg.load_operand(1)
    else:
      size = get_reg_size(insn.Op2)
      arg = ida_hexrays.reg2mreg(insn.Op2.reg)

    call_arg = ida_hexrays.mcallarg_t()
    call_arg.set_regarg(arg, size, get_type_info(size))

    call_info = ida_hexrays.mcallinfo_t()
    call_info.cc = ida_typeinf.CM_CC_FASTCALL
    call_info.callee = ida_idaapi.BADADDR
    call_info.role = ida_hexrays.ROLE_UNK
    call_info.flags = ida_hexrays.FCI_SPLOK | ida_hexrays.FCI_FINAL | ida_hexrays.FCI_PROP | ida_hexrays.FCI_PURE
    call_info.solid_args = 1
    call_info.args.push_back(call_arg)
    call_info.return_type = get_type_info(size)

    call_insn = ida_hexrays.minsn_t(cdg.insn.ea)
    call_insn.opcode = ida_hexrays.m_call
    call_insn.l.make_helper(get_helper_name(size))
    call_insn.d.t = ida_hexrays.mop_f
    call_insn.d.f = call_info
    call_insn.d.size = size

    if to_reg:
      call_result = cdg.mba.alloc_kreg(size)

      mov_insn = ida_hexrays.minsn_t(cdg.insn.ea)
      mov_insn.opcode = ida_hexrays.m_mov
      mov_insn.l.t = ida_hexrays.mop_d
      mov_insn.l.d = call_insn
      mov_insn.l.size = size
      mov_insn.d.t = ida_hexrays.mop_r
      mov_insn.d.r = call_result
      mov_insn.d.size = size

      cdg.mb.insert_into_block(mov_insn, cdg.mb.tail)
      dest = ida_hexrays.reg2mreg(insn.Op1.reg)
      cdg.emit(ida_hexrays.m_mov, size, call_result, 0, dest, 0)
      cdg.mba.free_kreg(call_result, size)
    else:
      call_mop = ida_hexrays.mop_t()
      call_mop.t = ida_hexrays.mop_d
      call_mop.d = call_insn
      call_mop.size = size
      cdg.store_operand(0, call_mop)

    return ida_hexrays.MERR_OK


def PLUGIN_ENTRY():
  return MOVBE()


class MOVBE(ida_idaapi.plugin_t):
  flags = ida_idaapi.PLUGIN_PROC | ida_idaapi.PLUGIN_HIDE
  comment = "MOVBE support for the Hex-Rays x64 Decompiler"
  help = ""
  wanted_name = "MOVBE"
  wanted_hotkey = ""
  loaded = False

  def init(self):
    if ida_idp.ph.id != ida_idp.PLFM_386:
      return ida_idaapi.PLUGIN_SKIP

    ida_loader.load_plugin("hexx64")
    if not ida_hexrays.init_hexrays_plugin():
      return ida_idaapi.PLUGIN_SKIP
    assert ida_hexrays.init_hexrays_plugin(), "Missing Hexx64 Decompiler..."

    self.movbe_lifter = MOVBELifter()
    self.movbe_lifter.install()

    return ida_idaapi.PLUGIN_KEEP

  def run(self, arg):
    ida_kernwin.warning("%s cannot be run as a script in IDA." % self.wanted_name)

  def term(self):
    self.movbe_lifter = None


repository: OALabs/hashdb-ida
file: hashdb.py
description: HashDB API hash lookup plugin for IDA Pro
url: https://github.com/OALabs/hashdb-ida/blob/89f3c519ece5b7b76e58392881d31f3e7eef3296/hashdb.py
comment: HashDB Lookup Service
created_at: 2021-09-14 20:04:37+00:00
pushed_at: 2024-10-15 01:58:16+00:00
content:
########################################################################################
##
## This plugin is the client for the HashDB lookup service operated by OALABS:
##
## https://hashdb.openanalysis.net/
##
##   _   _           _    ____________ 
##  | | | |         | |   |  _  \ ___ \ 
##  | |_| | __ _ ___| |__ | | | | |_/ /
##  |  _  |/ _` / __| '_ \| | | | ___ \ 
##  | | | | (_| \__ \ | | | |/ /| |_/ /
##  \_| |_/\__,_|___/_| |_|___/ \____/ 
##
## HashDB is a community-sourced library of hashing algorithms used in malware.
## New hash algorithms can be added here: https://github.com/OALabs/hashdb
##
## Updated for IDA 7.xx and Python 3
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\<ida version>\plugins
##
## To run:
##      Configure Settings:
##          Edit->Plugins->HashDB
##          click `Refresh Algorithms` to pull a list of hash algorithms
##          select the hash algorithm you need from the drop-down
##          OK
##      Lookup Hash:
##          Highlight constant in IDA disassembly or psuedocode view
##          Right-click -> HashDB Lookup
##          If a hash is found it will be added to an enum controlled in the settings
##          Right-click on the constant again -> Enum -> Select new hash enum
##
########################################################################################

__AUTHOR__ = '@herrcore'

PLUGIN_NAME = "HashDB"
PLUGIN_HOTKEY = 'Alt+`'
VERSION = '1.10.0'

import sys
import time

import idaapi

#--------------------------------------------------------------------------
# IDA Python version madness
#--------------------------------------------------------------------------

major, minor = map(int, idaapi.get_kernel_version().split("."))
assert (major > 6),"ERROR: HashDB plugin requires IDA v7+"
assert (sys.version_info >= (3, 5)), "ERROR: HashDB plugin requires Python 3.5"
# We need to make some adjustments for IDA 9
IDA_9 = major >= 9

if IDA_9:
    import ida_ida
else:
    import ida_enum


import idc
import ida_kernwin
import ida_name
import ida_bytes
import ida_netnode
import ida_typeinf

# Imports for the exception handler
import traceback
import json
import webbrowser
import urllib.parse


#--------------------------------------------------------------------------
# Setup IDA 9 globals 
#--------------------------------------------------------------------------
if IDA_9:
    BITS =  32 if not ida_ida.inf_is_64bit() else 64
    BWN_DISASM = ida_kernwin.BWN_DISASM
    BWN_PSEUDOCODE = ida_kernwin.BWN_PSEUDOCODE
    SETMENU_APP = ida_kernwin.SETMENU_APP
else:
    BITS = 32 if not idaapi.get_inf_structure().is_64bit() else 64
    BWN_DISASM = idaapi.BWN_DISASMS
    BWN_PSEUDOCODE = idaapi.BWN_PSEUDOCODE
    SETMENU_APP = idaapi.SETMENU_APP


def get_enum(enum_name: str):
    if IDA_9:
        return idc.get_enum(enum_name)
    else:
        return ida_enum.get_enum(enum_name)


def import_type(idx: int, enum_name: str):
    if IDA_9:
        return idc.import_type(idx, enum_name)
    else:
        return ida_typeinf.import_type(ida_typeinf.get_idati(), idx, enum_name, 0)
    

def attach_action_to_popup(widget, popup_handle, name, popuppath=None, flags=0):
    if IDA_9:
        return ida_kernwin.attach_action_to_popup(widget, popup_handle, name, popuppath, flags)
    else:
        return idaapi.attach_action_to_popup(widget, popup_handle, name, popuppath, flags)


def get_enum_member_by_name(enum_value_name: str):
    if IDA_9:
        return idc.get_enum_member_by_name(enum_value_name)
    else:
        return ida_enum.get_enum_member_by_name(enum_value_name)


#--------------------------------------------------------------------------
# Global exception hook to detect plugin exceptions until
#  we implement a proper test-driven development setup
# Note: minimum Python version support is 3.5 
#--------------------------------------------------------------------------
HASHDB_REPORT_BUG_URL = "https://github.com/OALabs/hashdb-ida/issues/new"
def hashdb_exception_hook(exception_type, value, traceback_object):
    is_hashdb_exception = False

    frame_data = {
        "user_data": {
            "platform": sys.platform,
            "python_version": '.'.join([str(sys.version_info.major), str(sys.version_info.minor), str(sys.version_info.micro)]),
            "plugin_version": VERSION,
            "ida": {
                "kernel_version": ida_kernwin.get_kernel_version(),
                "bits": BITS
            }
        },
        "exception_data": {
            "exception_type": exception_type.__name__,
            "exception_value": str(value)
        },
        "frames": []}
    frame_summaries = traceback.StackSummary.extract(traceback.walk_tb(traceback_object), capture_locals=True)
    for frame_index, frame_summary in enumerate(frame_summaries):
        file_name = frame_summary.filename
        if "__file__" in globals():
            if not file_name == __file__:
                continue
        is_hashdb_exception = True

        # Save frame data
        frame_data["frames"].append({
            "frame_index": frame_index,
            "line_number": frame_summary.lineno,
            "function_name": frame_summary.name,
            "line": frame_summary.line,
            "locals": frame_summary.locals
        })

    if is_hashdb_exception:
        class crash_detection_form(ida_kernwin.Form):
            def __init__(self):
                form = "BUTTON YES* Yes\nBUTTON CANCEL No\nHashDB Error!\n\n{format}"
                controls = {
                    "format": super().StringLabel(
                    """<center>
                        <p style="margin: 0; font-size: 20px; color: #F44336"><b>HashDB has detected an internal error.</b><p>
                        <p style="margin: 0; font-size: 12px">Would you like to submit a stack trace to the developers?</p>
                        <ol style="font-size: 11px; text-align: left">
                            <li>Selecting "Yes" will open a feedback dialogue and redirect you to:
                                <p style="margin: 0 4px 0 0;"><b><i>github.com/OALabs/hashdb-ida</i></b></p>
                            </li>
                            <li>All personally identifiable information will be removed.</li>
                            <li>Afterwards, you will be asked if you want to unload the plugin.</li>
                        </ol>
                    </center>""", super().FT_HTML_LABEL)
                }
                super().__init__(form, controls)
                
                # Compile
                self.Compile()

        # Execute the crash detection form on the main thread
        crash_form = crash_detection_form()
        crash_button_selected = ida_kernwin.execute_sync(crash_form.Execute, ida_kernwin.MFF_FAST)
        crash_form.Free()

        # Did the user allow us to submit a request?
        if crash_button_selected == 1: # Yes button
            # Setup the body
            body = "## Steps to reproduce:\n1. \n\n## Stack trace:\n```\n{}\n```".format(json.dumps(frame_data))
            
            # Open the tab
            global HASHDB_REPORT_BUG_URL
            webbrowser.open_new_tab(HASHDB_REPORT_BUG_URL + "?" + urllib.parse.urlencode({
                "title": "[BUG]: ",
                "body": body
            }))
    
        # Ask the user if they want to terminate the plugin
        class unload_plugin_form(ida_kernwin.Form):
            def __init__(self):
                form = "BUTTON YES* Yes\nBUTTON CANCEL No\nHashDB\n\n{format}"
                controls = {
                    "format": super().StringLabel(
                    """<center>
                        <p style="margin: 0; font-size: 20px; color: #F44336"><b>Would you like to unload the plugin?</b><p>
                        <p>This action will make the plugin unusable until IDA is restarted.</p>
                    </center>""", super().FT_HTML_LABEL)
                }
                super().__init__(form, controls)
                
                # Compile
                self.Compile()
        unload_form = unload_plugin_form()
        unload_button_selected = ida_kernwin.execute_sync(unload_form.Execute, ida_kernwin.MFF_FAST)
        unload_form.Free()
        
        if unload_button_selected == 1: # Yes button
            global HASHDB_PLUGIN_OBJECT
            ida_kernwin.execute_sync(HASHDB_PLUGIN_OBJECT.term, ida_kernwin.MFF_FAST)

    sys.__excepthook__(exception_type, value, traceback_object)
sys.excepthook = hashdb_exception_hook

# Rest of the imports
import re
import functools
import requests
from typing import Union

# These imports are specific to the Worker implementation
import inspect
import logging
import threading
from threading import Thread
from dataclasses import dataclass, field
from typing import Callable

#--------------------------------------------------------------------------
# Global settings/variables
#--------------------------------------------------------------------------

HASHDB_API_URL ="https://hashdb.openanalysis.net"
HASHDB_USE_XOR = False
HASHDB_XOR_VALUE = 0
HASHDB_ALGORITHM = None
HASHDB_ALGORITHM_SIZE = 0
ENUM_PREFIX = "hashdb_strings"
NETNODE_NAME = "$hashdb"

# Variables for async operations
HASHDB_REQUEST_TIMEOUT = 15 # Limit to 15 seconds
HASHDB_REQUEST_LOCK = threading.Lock()

#--------------------------------------------------------------------------
# Setup Icon
#--------------------------------------------------------------------------

HASH_ICON_DATA = b"".join([b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x10\x00\x00\x00\x10\x08',
                          b'\x04\x00\x00\x00\xb5\xfa7\xea\x00\x00\x00\x04gAMA\x00\x00\xb1\x8f\x0b\xfca',
                          b'\x05\x00\x00\x00 cHRM\x00\x00z&\x00\x00\x80\x84\x00\x00\xfa\x00\x00\x00\x80',
                          b'\xe8\x00\x00u0\x00\x00\xea`\x00\x00:\x98\x00\x00\x17p\x9c\xbaQ<\x00\x00\x00',
                          b'\x02bKGD\x00\xff\x87\x8f\xcc\xbf\x00\x00\x00\tpHYs\x00\x00\x0b\x13\x00\x00',
                          b'\x0b\x13\x01\x00\x9a\x9c\x18\x00\x00\x00\x07tIME\x07\xe5\t\x18\x12\x18(\xba',
                          b'\xecz-\x00\x00\x01#IDAT(\xcfm\xd1\xbdJ\x9ba\x18\xc6\xf1_\xde<\xd5d\x08\xc1',
                          b'\xb46\x967!\x1d,\x88\xd0\xa1P\xe8\x01\x14\x0c\xb8\xbbt\xa9\xa3\x07\xd0\xb9',
                          b'\xab \x1e\x83s\x87R\xa4]K\xe8".*NEpJZL\x9b\xa2V\x90\xc6\xa4\xc6\xc7%\x92\xa0',
                          b'\xfe\xd7\xeb\xe6\xe6\xfa`\x9c\x8c\x82\x04\xe4\xe4\xdd\xc3\xb4\x0fV\x95\xf0',
                          b'\xd6\x17\x0bw\x0f\xeaz\xf6<\xf4\xc0\xa6h\x05\xc3\x877,\x98\xf0\xd5\xb1g^i\xfb',
                          b'\x06\x01AY\x10\x15\xbdv\xe9\xbb\x19\x8bf4\x0c\xa4~g\x90\xfa\xa8\xeaJ\xd6c\x89',
                          b'\x8e\xbe\xa2\xa2s\x7f\xb5\xbcI\xc6\x12\x94\x04\'\xfa\xf2\x8azNen\xa4\xac\'*^8',
                          b'\xd0\xb5\xa4\xec\xbd\xe8\xb3\xa7\xaaR!\x08\xca\x12\x03\xb3j\x9a\x0e\xe5\xbc',
                          b'\xc4\x8e\xbe\xa8c@\xcd\x96\x9f\x9a\xfe\x88\xbaZZ.D\x1d?lKG1\'\x94\\:\x11M\x99t',
                          b'\xa6;r\x10\xa4*\x96\xfd\xb7\xef\xb9Y\r\xd1;\xa9\x9aT\x18U\xb4&Z\xc7\x9c#m\xf3',
                          b'\xb7+~dOO\x1d+\xa2M\x93#);\xdc\xae\xec\x97\r\xff\x94L\xf9d\xf7\xeeL\x89\xc2',
                          b'\xd0V^n\\\xb8\x06\xd6\xa1L\xe6_H\xbf\xfc\x00\x00\x00%tEXtdate:create\x00202',
                          b'1-09-24T18:24:40+00:00\xd7;f\xf5\x00\x00\x00%tEXtdate:modify\x002021-09-24T',
                          b'18:24:40+00:00\xa6f\xdeI\x00\x00\x00WzTXtRaw profile type iptc\x00\x00x\x9c',
                          b'\xe3\xf2\x0c\x08qV((\xcaO\xcb\xccI\xe5R\x00\x03#\x0b.c\x0b\x13#\x13K\x93\x14',
                          b'\x03\x13 D\x804\xc3d\x03#\xb3T \xcb\xd8\xd4\xc8\xc4\xcc\xc4\x1c\xc4\x07\xcb',
                          b'\x80H\xa0J.\x00\xea\x17\x11t\xf2B5\x95\x00\x00\x00\x00IEND\xaeB`\x82'])
HASH_ICON = ida_kernwin.load_custom_icon(data=HASH_ICON_DATA, format="png")
XOR_ICON_DATA = b"".join([b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x10\x00\x00\x00\x10\x08\x04',
                         b'\x00\x00\x00\xb5\xfa7\xea\x00\x00\x00\x04gAMA\x00\x00\xb1\x8f\x0b\xfca\x05\x00',
                         b'\x00\x00 cHRM\x00\x00z&\x00\x00\x80\x84\x00\x00\xfa\x00\x00\x00\x80\xe8\x00\x00',
                         b'u0\x00\x00\xea`\x00\x00:\x98\x00\x00\x17p\x9c\xbaQ<\x00\x00\x00\x02bKGD\x00\xff',
                         b'\x87\x8f\xcc\xbf\x00\x00\x00\tpHYs\x00\x00\x0b\x13\x00\x00\x0b\x13\x01\x00\x9a',
                         b'\x9c\x18\x00\x00\x00\x07tIME\x07\xe5\t\x18\x12\x0b";\xd6\xd2\xa1\x00\x00\x00\xc3',
                         b'IDAT(\xcf\xa5\xd01N\x02\x01\x14\x04\xd0\x07B01\x9a\x10b\x89%\t\x8dg\xa0\xd0f-\xb8',
                         b'\x80x\x86\x8d\r\xd9#X\xee\x05(\x94\x0b\xd0\xd0@A\xcb\t4\xdb\x98\xd8Z\x90\xacv\x82',
                         b'Z,\xac\xab1P0\xdd\xfc\xcc\x9f\xcc\x0c\xfb\xa2\xf4\x8bU\x9c \xb5\xfcOPu\xe9F\x0b',
                         b'\x89{\x13\x1f\xd9\xf9 \xff\xbd\x15\x99;\xf2.\x11\xaa\x99\xfb,\x9a_y\x12 \x16#X3',
                         b'\x94\xd7>\xd7F\xc6\xb9|l\xa4\x97\xb9g\x19\xea\xa6^=*\xe9`\xe6K\xdb\xa9\x0b\x8b',
                         b'\x8d\xc3\x16T@*\xf1\xa2\x8f\x18!\xee\x9cI\x7f2\xac\x0cu7\xb1\x10\xe8z\xb0*\xd6',
                         b'|v(\xd2\xd4\xd6p.40\xccj\xee\x1c\xea\xef\xd4\xc7x+N\xbd?\xbe\x01\xa7\xee.6\xd9',
                         b'\xf6\xa5\xd2\x00\x00\x00%tEXtdate:create\x002021-09-24T18:11:34+00:00Vz\xe6\xba',
                         b'\x00\x00\x00%tEXtdate:modify\x002021-09-24T18:11:34+00:00\'\'^\x06\x00\x00\x00',
                         b'WzTXtRaw profile type iptc\x00\x00x\x9c\xe3\xf2\x0c\x08qV((\xcaO\xcb\xccI\xe5R',
                         b'\x00\x03#\x0b.c\x0b\x13#\x13K\x93\x14\x03\x13 D\x804\xc3d\x03#\xb3T \xcb\xd8\xd4',
                         b'\xc8\xc4\xcc\xc4\x1c\xc4\x07\xcb\x80H\xa0J.\x00\xea\x17\x11t\xf2B5\x95\x00\x00',
                         b'\x00\x00IEND\xaeB`\x82'])
XOR_ICON = ida_kernwin.load_custom_icon(data=XOR_ICON_DATA, format="png")
HUNT_ICON_DATA = b"".join([b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x10\x00\x00\x00\x10\x08\x04',
                          b'\x00\x00\x00\xb5\xfa7\xea\x00\x00\x00\x04gAMA\x00\x00\xb1\x8f\x0b\xfca\x05\x00',
                          b'\x00\x00 cHRM\x00\x00z&\x00\x00\x80\x84\x00\x00\xfa\x00\x00\x00\x80\xe8\x00\x00u0',
                          b'\x00\x00\xea`\x00\x00:\x98\x00\x00\x17p\x9c\xbaQ<\x00\x00\x00\x02bKGD\x00\xff\x87',
                          b'\x8f\xcc\xbf\x00\x00\x00\tpHYs\x00\x00\x0b\x13\x00\x00\x0b\x13\x01\x00\x9a\x9c\x18',
                          b'\x00\x00\x00\x07tIME\x07\xe5\t\x1d\x10#"R\xd1XW\x00\x00\x01.IDAT(\xcf\x8d\xd1;K\x9b',
                          b'\x01\x18\xc5\xf1\x9fI\xa8C\xbd\xf1\x0e\xdd\xd2\xc5NJ\x07;h\xd5\xa1\xf1\x0bT\xd4M',
                          b'\x14//\x82\xad\x8b\x83\x93`\xf1\x02~\x84\x08-b\x1c\xea\xe6\xe2 \x08^\x9aAQ\x07\x87',
                          b'R\x9d:\x99Atx\x15\xab`\xbc\xd0\x0eitS\xcf\xf2\xc0\xe1\xf0\x1c\xf8\x9f\x12\x0f*Q!@',
                          b'\xe4\xdc\xdf\x07\xb3xkuz\xe7\x05\xae\xedY\xb0_\x08\x15\x02\t=\x06l\xdap\x89\x97Z4',
                          b'\xfbf\xde-\t\xd0#4\xa1J\xef\xff\x8aE\xab\xc60[\xf8\xf0\xd6W\x93\xde\xfb`\xce!^\xeb',
                          b'\x93\xb5\xed\x8b\x01\xbf\xe2\x18v\xe4T\xbbQ\xcd\xba\xa4\\I\xebw\xe0N\x8d\xb5\x98Ju~h',
                          b'\x93\xd1\xaa\xda\xb8q\xd5>\xcah\x93U\xa72&P\xeaB \xa7\xde\x8cA\x83f4\xc8\t\xfcQ*\x88yB',
                          b'\t\x91\xbc2\x91\xa4]\x9f\xa4\xf1\xd9\x8e\xa4H\xb9\xbc(.\xaf\xd6\x1b\xebBi\xaftK\xf9i',
                          b'\xc9\x88\xef\x1a\xe5,\xc7ql\xc8\x8a;\xa1UK\xb2n\x8c\xc8\xfa\xad\xcb\xb4\x93\x02\xc9PhJ',
                          b'\x95\x8e{Pg\xc6\xcc\x16A\x15Qo\xd9p\x812-\x9a\x8a\xa8\x9f9\xd6#s\xff\x03\xabm^\xab\xaf',
                          b'\xe8z\xc0\x00\x00\x00%tEXtdate:create\x002021-09-29T16:35:34+00:00\xf4Q\xb1\xe8\x00\x00',
                          b'\x00%tEXtdate:modify\x002021-09-29T16:35:34+00:00\x85\x0c\tT\x00\x00\x00WzTXtRaw prof',
                          b'ile type iptc\x00\x00x\x9c\xe3\xf2\x0c\x08qV((\xcaO\xcb\xccI\xe5R\x00\x03#\x0b.c\x0b',
                          b'\x13#\x13K\x93\x14\x03\x13 D\x804\xc3d\x03#\xb3T \xcb\xd8\xd4\xc8\xc4\xcc\xc4\x1c',
                          b'\xc4\x07\xcb\x80H\xa0J.\x00\xea\x17\x11t\xf2B5\x95\x00\x00\x00\x00IEND\xaeB`\x82'])
HUNT_ICON = ida_kernwin.load_custom_icon(data=HUNT_ICON_DATA, format="png")
SCAN_ICON_DATA = b"".join([b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x10\x00\x00\x00\x10\x08\x04\x00\x00',
                          b'\x00\xb5\xfa7\xea\x00\x00\x00\x04gAMA\x00\x00\xb1\x8f\x0b\xfca\x05\x00\x00\x00 cHRM',
                          b'\x00\x00z&\x00\x00\x80\x84\x00\x00\xfa\x00\x00\x00\x80\xe8\x00\x00u0\x00\x00\xea`\x00',
                          b'\x00:\x98\x00\x00\x17p\x9c\xbaQ<\x00\x00\x00\x02bKGD\x00\xff\x87\x8f\xcc\xbf\x00\x00',
                          b'\x00\tpHYs\x00\x00\x0b\x13\x00\x00\x0b\x13\x01\x00\x9a\x9c\x18\x00\x00\x00\x07tIME\x07',
                          b'\xe5\n\x08\x17\x1c\x04\xfd*<n\x00\x00\x01#IDAT(\xcfe\xd1\xb1K\xe2\x01\x18\xc6\xf1\x8f',
                          b'\xf6\xcb\xeb\xd0\x04\xc7 \xa8Ii\x89\xa2\xed~\x16\xc4m\xd6\xbf\xd0\x18\xed\xd1m\x07F-ECK',
                          b'\xa0K\xc49_\rACk\x10\xdc\x12\x04\x91qR\xd2\xd0PS\xd7\xa0h\x905\xe8\x85\xd9\xf7\x1d\xdfgx',
                          b'\xdf\xe7K\x9b@\xa8\xa8\xa2aO\x9f.\x02\x90\xb2l\xc9\x80\xb2s\'ZzH)i\xda\x17\x8a\x8b\x89',
                          b'\xf6\xae\xa3\xd65m\x88\xcbXud\xa57\x12z\xf0[\xc2\xbc\xaa\xbaK?{\x03EO\xa6e\\\xbb\x94',
                          b'\x93\xfcx"\xfc\xf5GB^MN\xca\x8a\x1f\x86\x0c\xf8\xda\x99\xfe\xc0\xb0s\xcf\xa6T\x9dZ\xb4',
                          b'\xe5\xc5\x82G-\x11Q-\xd5\xc0g^}1\xe9\x9f\x8aq\x13\x81;#b\xce|\x97\xb5\x8bW\xbf\xcc*)',
                          b'\xd8q,A\xe1\xfd\xc8\xb2\x9c\xa4\xa49W\xae\xa5\xcdxR t\xdfy\xf3F\xdd\x85\x0bu7\xe6\r:p/',
                          b'\xec.*-\xef\xd0\xa1\xbc\xb4\x84MMk\xedN\xfeW\x9d\x15\x17\x13\x13\x97u\xa0\xa9$E\xe4\x83',
                          b'\xac+\xb7\x185\xa6\xa1h\xdbc\xb7\xd5\xb6\xee\x9a\x9a\x8a\xa2o\x1d\xcf\xde\x00\x9fhY\xc0',
                          b'\x9b\x9d\xab^\x00\x00\x00%tEXtdate:create\x002021-10-08T23:28:04+00:00\xee\x90\xd3~\x00',
                          b'\x00\x00%tEXtdate:modify\x002021-10-08T23:28:04+00:00\x9f\xcdk\xc2\x00\x00\x00WzTXtRaw ',
                          b'profile type iptc\x00\x00x\x9c\xe3\xf2\x0c\x08qV((\xcaO\xcb\xccI\xe5R\x00\x03#\x0b.c\x0b',
                          b'\x13#\x13K\x93\x14\x03\x13 D\x804\xc3d\x03#\xb3T \xcb\xd8\xd4\xc8\xc4\xcc\xc4\x1c\xc4\x07',
                          b'\xcb\x80H\xa0J.\x00\xea\x17\x11t\xf2B5\x95\x00\x00\x00\x00IEND\xaeB`\x82'])
SCAN_ICON = ida_kernwin.load_custom_icon(data=SCAN_ICON_DATA, format="png")

#--------------------------------------------------------------------------
# Error class
#--------------------------------------------------------------------------
class HashDBError(Exception):
    pass


#--------------------------------------------------------------------------
# Worker implementation
#--------------------------------------------------------------------------
@dataclass(unsafe_hash=True)
class Worker(Thread):
    """The worker implementation for multi-threading support."""
    target: Callable
    args: tuple = field(default_factory=tuple, compare=False)
    done_callback: Callable = None
    error_callback: Callable = None

    def __post_init__(self):
        """Required to initialize the base class (Thread)."""
        super().__init__(target=self.__wrapped_target, args=self.args, daemon=True)

    def __wrapped_target(self, *args, **kwargs):
        """
        Wraps the target function to allow callbacks and error handling.
        @raise Exception: if an unhandled exception is encountered it will
                          be raised
        """
        try:
            # Execute the target
            results = self.target(*args, **kwargs)

            # Execute the done callback, if it exists
            if self.done_callback is not None:
                # Call the function based on the amount of arguments it expects
                argument_spec = inspect.getfullargspec(self.done_callback)
                argument_count = len(argument_spec.args)

                if argument_count > 1:
                    self.done_callback(*results)
                elif argument_count == 1 and results is not None:
                    self.done_callback(results)
                else:
                    self.done_callback()
        except Exception as exception:
            # Execute the error callback, if it exits;
            #  otherwise raise the exception (unhandled)
            if self.error_callback is not None:
                # Call the function based on the amount of arguments it expects
                argument_spec = inspect.getfullargspec(self.error_callback)
                argument_count = len(argument_spec.args)

                if argument_count == 1:
                    self.error_callback(exception)
                else:
                    self.error_callback()
            else:
                raise exception
        finally:
            # Cleanup the callbacks (decrease reference counts)
            if self.done_callback is not None:
                del self.done_callback
            if self.error_callback is not None:
                del self.error_callback

                
#--------------------------------------------------------------------------
# HashDB API 
#--------------------------------------------------------------------------

def get_algorithms(api_url='https://hashdb.openanalysis.net', timeout=None):
    # Handle an empty timeout
    global HASHDB_REQUEST_TIMEOUT
    if timeout is None:
        timeout = HASHDB_REQUEST_TIMEOUT

    algorithms_url = api_url + '/hash'
    r = requests.get(algorithms_url, timeout=timeout)
    if not r.ok:
        raise HashDBError("Get algorithms API request failed, status %s" % r.status_code)
    results = r.json()

    algorithms = []
    for algorithm in results.get('algorithms',[]):
        size = determine_algorithm_size(algorithm.get('type', None))
        if size == 'Unknown':
            idaapi.msg("ERROR: Unknown algorithm type encountered when fetching algorithms: %s" % size)
        algorithms.append([algorithm.get('algorithm'), size])
    return algorithms


def get_strings_from_hash(algorithm, hash_value, xor_value=0, api_url='https://hashdb.openanalysis.net', timeout=None):
    # Handle an empty timeout
    global HASHDB_REQUEST_TIMEOUT
    if timeout is None:
        timeout = HASHDB_REQUEST_TIMEOUT

    hash_value ^= xor_value
    hash_url = api_url + '/hash/%s/%d' % (algorithm, hash_value)
    r = requests.get(hash_url, timeout=timeout)
    if not r.ok:
        raise HashDBError("Get hash API request failed, status %s" % r.status_code)
    results = r.json()
    # Remove null bytes from non-api strings
    hashes = results.get('hashes',[])
    out_hashes = []
    for hash_info in hashes:
        if not hash_info.get('string',{}).get('is_api',True):
            hash_info['string']['string'] = hash_info['string']['string'].replace('\x00','')
        out_hashes.append(hash_info)
    return {'hashes':out_hashes}


def get_module_hashes(module_name, algorithm, permutation, api_url='https://hashdb.openanalysis.net', timeout=None):
    # Handle an empty timeout
    global HASHDB_REQUEST_TIMEOUT
    if timeout is None:
        timeout = HASHDB_REQUEST_TIMEOUT
    
    module_url = api_url + '/module/%s/%s/%s' % (module_name, algorithm, permutation)
    r = requests.get(module_url, timeout=timeout)
    if not r.ok:
        raise HashDBError("Get hash API request failed, status %s" % r.status_code)
    results = r.json()
    return results


def hunt_hash(hash_value, api_url='https://hashdb.openanalysis.net', timeout = None):
    # Handle an empty timeout
    global HASHDB_REQUEST_TIMEOUT
    if timeout is None:
        timeout = HASHDB_REQUEST_TIMEOUT
    
    matches = []
    hash_list = [hash_value]
    module_url = api_url + '/hunt'
    r = requests.post(module_url, json={"hashes": hash_list}, timeout=timeout)
    if not r.ok:
        print(module_url)
        print(hash_list)
        print(r.json())
        raise HashDBError("Get hash API request failed, status %s" % r.status_code)
    for hit in r.json().get('hits',[]):
        algo = hit.get('algorithm',None)
        if (algo != None) and (algo not in matches):
            matches.append(algo)
    return matches


#--------------------------------------------------------------------------
# Save and restore settings
#--------------------------------------------------------------------------
def load_settings():
    global HASHDB_API_URL 
    global HASHDB_USE_XOR, HASHDB_XOR_VALUE 
    global HASHDB_ALGORITHM, ENUM_PREFIX
    global NETNODE_NAME
    node = ida_netnode.netnode(NETNODE_NAME)
    if ida_netnode.exist(node):
        if bool(node.hashstr("HASHDB_API_URL")):
            HASHDB_API_URL = node.hashstr("HASHDB_API_URL")
        if bool(node.hashstr("HASHDB_USE_XOR")):
            if node.hashstr("HASHDB_USE_XOR").lower() == "true":
                HASHDB_USE_XOR = True
            else: 
                HASHDB_USE_XOR = False
        if bool(node.hashstr("HASHDB_XOR_VALUE")):
            HASHDB_XOR_VALUE = int(node.hashstr("HASHDB_XOR_VALUE"))
        if bool(node.hashstr("HASHDB_ALGORITHM")) and bool(node.hashstr("HASHDB_ALGORITHM_SIZE")):
            successful = set_algorithm(node.hashstr("HASHDB_ALGORITHM"), node.hashstr("HASHDB_ALGORITHM_SIZE"))
            if not successful:
                idaapi.msg("HashDB failed to set the algorithm when parsing the saved config!\n")
        if bool(node.hashstr("ENUM_PREFIX")):
            ENUM_PREFIX = node.hashstr("ENUM_PREFIX")
        idaapi.msg("HashDB configuration loaded!\n")
    else:
        idaapi.msg("No saved HashDB configuration\n")
    return


def save_settings():
    global HASHDB_API_URL 
    global HASHDB_USE_XOR, HASHDB_XOR_VALUE 
    global HASHDB_ALGORITHM, ENUM_PREFIX
    global NETNODE_NAME

    # Check if our netnode already exists, otherwise create a new one
    node = ida_netnode.netnode(NETNODE_NAME)
    if not ida_netnode.exist(node):
        node = ida_netnode.netnode()
        if not node.create(NETNODE_NAME):
            idaapi.msg("ERROR: Unable to save HashDB settings, failed to create the netnode.\n")
            return

    if HASHDB_API_URL != None:
        node.hashset_buf("HASHDB_API_URL", str(HASHDB_API_URL))
    if HASHDB_USE_XOR != None:
        node.hashset_buf("HASHDB_USE_XOR", str(HASHDB_USE_XOR))
    if HASHDB_XOR_VALUE != None:
        node.hashset_buf("HASHDB_XOR_VALUE", str(HASHDB_XOR_VALUE))
    if HASHDB_ALGORITHM != None:
        node.hashset_buf("HASHDB_ALGORITHM", str(HASHDB_ALGORITHM))
    if HASHDB_ALGORITHM_SIZE != None:
        node.hashset_buf("HASHDB_ALGORITHM_SIZE", str(HASHDB_ALGORITHM_SIZE))
    if ENUM_PREFIX != None:
        node.hashset_buf("ENUM_PREFIX", str(ENUM_PREFIX))
    idaapi.msg("HashDB settings saved\n")


#--------------------------------------------------------------------------
# Settings form
#--------------------------------------------------------------------------
class hashdb_settings_t(ida_kernwin.Form):
    """Global settings form for hashdb"""

    class algorithm_chooser_t(ida_kernwin.Choose):
        """
        A simple chooser to be used as an embedded chooser
        """
        def __init__(self, algo_list):
            ida_kernwin.Choose.__init__(
                self,
                "",
                [
                    ["Algorithm", 15],
                    ["Size (Bits)", 5]
                ],
                flags=0,
                embedded=True,
                width=30,
                height=6)
            self.items = algo_list
            self.icon = None

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

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


    def __init__(self, algorithms):
        self.__n = 0
        F = ida_kernwin.Form
        F.__init__(self,
r"""BUTTON YES* Ok
BUTTON CANCEL Cancel
HashDB Settings

{FormChangeCb}
<##API URL          :{iServer}>
<##Enum Prefix      :{iEnum}>
<Enable XOR:{rXor}>{cXorGroup}>  |  <##:{iXor}>(hex)
<Select algorithm :{cAlgoChooser}><Refresh Algorithms:{iBtnRefresh}>

""", {      'FormChangeCb': F.FormChangeCb(self.OnFormChange),
            'iServer': F.StringInput(),
            'iEnum': F.StringInput(),
            'cXorGroup': F.ChkGroupControl(("rXor",)),
            'iXor': F.NumericInput(tp=F.FT_RAWHEX),
            'cAlgoChooser' : F.EmbeddedChooserControl(hashdb_settings_t.algorithm_chooser_t(algorithms)),
            'iBtnRefresh': F.ButtonInput(self.OnBtnRefresh),
        })

    def OnBtnRefresh(self, code=0):
        api_url = self.GetControlValue(self.iServer)
        algorithms = []
        try:
            ida_kernwin.show_wait_box("HIDECANCEL\nPlease wait...")
            algorithms = get_algorithms(api_url=api_url)
        except Exception as e:
            idaapi.msg("ERROR: HashDB API request failed: %s\n" % e)
        finally:
            ida_kernwin.hide_wait_box()
        # Sort the algorithms by algorithm name (lowercase)
        sorted_algorithms = sorted(algorithms, key = lambda algorithm: algorithm[0].lower())
        self.cAlgoChooser.chooser.items = sorted_algorithms
        self.RefreshField(self.cAlgoChooser)


    def OnFormChange(self, fid):
        if fid == -1:
            # Form is initialized
            # Hide Xor input if dissabled 
            if self.GetControlValue(self.cXorGroup) == 1:
                self.EnableField(self.iXor, True)
            else:
                self.EnableField(self.iXor, False)
            self.SetFocusedField(self.cAlgoChooser)
        elif fid == self.cXorGroup.id:
            if self.GetControlValue(self.cXorGroup) == 1:
                self.EnableField(self.iXor, True)
            else:
                self.EnableField(self.iXor, False)
        else:
            pass
            #print("Unknown fid %r" % fid)
        return 1

    @staticmethod
    def show(api_url="https://hashdb.openanalysis.net",
             enum_prefix="hashdb_strings",
             use_xor=False,
             xor_value=0,
             algorithms=[]):
        global HASHDB_API_URL
        global HASHDB_USE_XOR
        global HASHDB_XOR_VALUE
        global HASHDB_ALGORITHM
        global ENUM_PREFIX
        # Sort the algorithms
        sorted_algorithms = sorted(algorithms, key = lambda algorithm: algorithm[0].lower())
        f = hashdb_settings_t(sorted_algorithms)
        f, args = f.Compile()
        # Set default values
        f.iServer.value = api_url
        f.iEnum.value = enum_prefix
        if use_xor:
            f.rXor.checked = True
        else:
            f.rXor.checked = False
        f.iXor.value = xor_value
        # Show form
        ok = f.Execute()
        if ok == 1:
            # Save default settings first
            HASHDB_USE_XOR = f.rXor.checked
            HASHDB_XOR_VALUE = f.iXor.value
            HASHDB_API_URL = f.iServer.value
            ENUM_PREFIX = f.iEnum.value
            # Check if algorithm is selected
            if f.cAlgoChooser.selection == None:
                # No algorithm selected bail!
                idaapi.msg("HashDB: No algorithm selected!\n")
                f.Free()
                return False
            # Set the algorithm
            algorithm = f.cAlgoChooser.chooser.items[f.cAlgoChooser.selection[0]]
            set_algorithm(algorithm[0], algorithm[1]) # Error messages handled inside of set_algorithm
            f.Free()
            return True
        else:
            f.Free()
            return False


#--------------------------------------------------------------------------
# Hash collision select form
#--------------------------------------------------------------------------
class match_select_t(ida_kernwin.Form):
    """Simple form to select string match during hash collision"""
    def __init__(self, collision_strings):
        self.__n = 0
        F = ida_kernwin.Form
        F.__init__(self,
r"""BUTTON YES* Ok
HashDB Hash Collision

{FormChangeCb}
More than one string matches this hash!
<Select the correct string :{cbCollisions}>

""", {      'FormChangeCb': F.FormChangeCb(self.OnFormChange),
            'cbCollisions': F.DropdownListControl(
                        items=collision_strings,
                        readonly=True,
                        selval=0),
        })


    def OnFormChange(self, fid):
        if fid == -1:
            # Form is initialized
            self.SetFocusedField(self.cbCollisions)
        elif fid == self.cbCollisions.id:
            sel_idx = self.GetControlValue(self.cbCollisions)
        else:
            pass
            #print("Unknown fid %r" % fid)
        return 1

    @staticmethod
    def show(collision_strings):
        global HASHDB_API_URL
        global HASHDB_USE_XOR
        global HASHDB_XOR_VALUE
        global HASHDB_ALGORITHM
        f = match_select_t(collision_strings)
        f, args = f.Compile()
        # Show form
        ok = f.Execute()
        if ok == 1:
            string_selection = f.cbCollisions[f.cbCollisions.value]
            f.Free()
            return string_selection
        else:
            f.Free()
            return None


#--------------------------------------------------------------------------
# Hash hunt results form
#--------------------------------------------------------------------------
class hunt_result_form_t(ida_kernwin.Form):

    class algorithm_chooser_t(ida_kernwin.Choose):
        """
        A simple chooser to be used as an embedded chooser
        """
        def __init__(self, algo_list):
            ida_kernwin.Choose.__init__(
                self,
                "",
                [
                    ["Algorithm", 10],
                    ["Size (Bits)", 5]
                ],
                flags=0,
                embedded=True,
                width=30,
                height=6)
            self.items = algo_list
            self.icon = None

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

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

    def __init__(self, algo_list, msg):
        self.invert = False
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""BUTTON YES* OK
Matched Algorithms

{FormChangeCb}
{cStrStatus}
<:{cAlgoChooser}>
""", {
            'cStrStatus': F.StringLabel(msg),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange),
            'cAlgoChooser' : F.EmbeddedChooserControl(hunt_result_form_t.algorithm_chooser_t(algo_list))
        })

    def OnFormChange(self, fid):
        if fid == -1:
            # Hide algorithm chooser if empty
            if self.cAlgoChooser.chooser.items == []:
                self.ShowField(self.cAlgoChooser, False)
        return 1

    def show(algo_list):
        global HASHDB_API_URL
        global HASHDB_USE_XOR
        global HASHDB_XOR_VALUE
        global HASHDB_ALGORITHM
        # Set default values
        if len(algo_list) == 0:
            msg = "No algorithms matched the hash."
            f = hunt_result_form_t(algo_list, msg)
        else:
            msg = "The following algorithms contain a matching hash.\nSelect an algorithm to set as the default for HashDB."
            f = hunt_result_form_t(algo_list, msg)
        f, args = f.Compile()
        # Show form
        ok = f.Execute()
        if ok == 1:
            if f.cAlgoChooser.selection == None:
                # No algorithm selected bail!
                f.Free()
                return False
            # Set the algorithm
            algorithm = f.cAlgoChooser.chooser.items[f.cAlgoChooser.selection[0]]
            set_algorithm(algorithm[0], algorithm[1]) # Error messages handled inside of set_algorithm
            f.Free()
            return True
        else:
            f.Free()
            return False


#--------------------------------------------------------------------------
# Module import select form
#--------------------------------------------------------------------------
class api_import_select_t(ida_kernwin.Form):
    """Simple form to select module to import apis from"""
    def __init__(self, string_value, module_list):
        self.__n = 0
        F = ida_kernwin.Form
        F.__init__(self,
r"""BUTTON YES* Import
BUTTON CANCEL No
HashDB Bulk Import

{FormChangeCb}
{cStr1} 
Do you want to import all function hashes from this module?
<Select module :{cbModules}>

""", {      'FormChangeCb': F.FormChangeCb(self.OnFormChange),
            'cStr1': F.StringLabel("<span style='float:left;'>The hash for <b>"+string_value+"</b> is a module function.<span>", tp=F.FT_HTML_LABEL),
            'cbModules': F.DropdownListControl(
                        items=module_list,
                        readonly=True,
                        selval=0),
        })

    def OnFormChange(self, fid):
        if fid == -1:
            # Form is initialized
            self.SetFocusedField(self.cbModules)
        elif fid == self.cbModules.id:
            sel_idx = self.GetControlValue(self.cbModules)
        else:
            pass
            #print("Unknown fid %r" % fid)
        return 1

    @staticmethod
    def show(string_value, module_list):
        f = api_import_select_t(string_value, module_list)
        f, args = f.Compile()
        # Show form
        ok = f.Execute()
        if ok == 1:
            module_selection = f.cbModules[f.cbModules.value]
            f.Free()
            return module_selection
        else:
            f.Free()
            return None


#--------------------------------------------------------------------------
# Unqualified name replacement form
# Logic: When an unqualified name is encountered, the user is asked to
#         provide a replacement.
# Example: "-path" is an unqualified name, the user should replace it with
#           a qualified name such as "_path" 
#--------------------------------------------------------------------------   
class unqualified_name_replace_t(ida_kernwin.Form):
    def __init__(self, unqualified_name: str, invalid_characters: list) -> None:
        form = "BUTTON YES* Replace\n" \
               "BUTTON CANCEL Skip\n" \
               "HashDB: Please replace the invalid characters\n\n" \
               "{form_change_callback}\n" \
               "Some of the characters in the hashed string are invalid (highlighted red):\n" \
               "{unqualified_name}\n" \
               "<##New name\: :{new_name}>"
        
        invalid_characters_html = "<span style=\"font-size: 16px\">{}</span>"
        controls = {
            "form_change_callback": super().FormChangeCb(self.form_changed),
            "unqualified_name": super().StringLabel(
                invalid_characters_html.format(html_format_invalid_characters(
                    unqualified_name, invalid_characters)),
                super().FT_HTML_LABEL),
            # value -> initial value
            "new_name": super().StringInput(value=unqualified_name)
        }
        super().__init__(form, controls)

        # Compile
        self.Compile()
    
    def form_changed(self, field_id: int) -> int:
        # Form initialized, focus to the new name text field
        if field_id == -1:
            self.SetFocusedField(self.new_name)
        return 1

    @staticmethod
    def show(unqualified_name: str, invalid_characters: list) -> str:
        """
        Show the unqualified name replace form and return
         the new user-defined name, or None if
         the user decides to skip.
        """

        # Construct and compile the form
        unqualified_name_form = unqualified_name_replace_t(unqualified_name, invalid_characters)
        # Execute/show the form
        selected_button = unqualified_name_form.Execute()
        new_name = unqualified_name_form.new_name.value

        # Free the form
        unqualified_name_form.Free()

        if selected_button == 1: # Replace button
            return new_name
        return None


#--------------------------------------------------------------------------
# IDA helper functions
#--------------------------------------------------------------------------
def get_invalid_characters(string: str) -> list:
    invalid_characters = []
    # Is the string empty?
    if not string:
        return invalid_characters

    # Is the first character a digit?
    if string[0].isdigit():
        invalid_characters.append(0)

    # Iterate through the characters in the string,
    #  and check if they are valid using
    #  ida_name.is_ident_cp
    for index, character in enumerate(string):
        if not ida_name.is_ident_cp(ord(character)):
            invalid_characters.append(index)

    # Return the invalid characters
    return invalid_characters


def html_format_invalid_characters(string: str, invalid_characters: list, color: str = "#F44336") -> str:
    # Are there any invalid characters in the string?
    if not invalid_characters:
        return string
    
    # Format the invalid characters
    formatted_string = ""
    for index, character in enumerate(string):
        if index in invalid_characters and color:
            formatted_string += "<span style=\"color: {}\">{}</span>".format(color, character)
        else:
            formatted_string += character

    # Return the formatted string
    return formatted_string


def get_existing_enum_values(enum_name):
    # Check if the enum exists
    if get_enum(enum_name) == idaapi.BADNODE:
        return {}

    # Fetch the type definition
    values = ida_typeinf.get_named_type(ida_typeinf.get_idati(), enum_name, ida_typeinf.NTF_TYPE)
    if values is None:
        return {}

    _, type_str, fields_str, _, _, _, _ = values
    type_definition = ida_typeinf.idc_print_type(type_str, fields_str, enum_name, 0)

    # Parse the type definition
    opening_bracket_position = type_definition.find("{")
    closing_bracket_position = type_definition.rfind("}")
    if opening_bracket_position == -1 or closing_bracket_position == -1:  # failed to find the brackets
        return {}

    # Split and parse the entries
    hash_values = {}

    split_entries = type_definition[opening_bracket_position+1:closing_bracket_position-1].split(",")
    for enum_value_name, enum_value_str in (entry.split(" = ") for entry in split_entries):
        base = 16 if enum_value_str.startswith("0x") or enum_value_str.startswith("-0x") else 10
        enum_value_str = enum_value_str.replace("LL", "")
        hash_values[enum_value_name] = int(enum_value_str, base)

    return hash_values


def add_enums(enum_name, hash_list):
    """
    Adds a hash list to an enum by name.
     IMPORTANT: This function should always be executed on the main thread.

    The hash list should be a list of tuples with three values:
     name: str, value: int, is_api: bool
    """
    fixed_hash_values = get_existing_enum_values(enum_name)  # dict[str, int]
    for enum_value_name, enum_value, is_api in hash_list:
        # Fixup the name before appending it to the list
        for index, character in enumerate(enum_value_name):
            if character.isspace() or character == ".":
                # Count not specified to replace all occurrences at once
                enum_value_name = enum_value_name.replace(character, '_')

        # Check if the name is valid and ask the user to replace invalid characters
        skip = False
        invalid_characters = get_invalid_characters(enum_value_name)
        while invalid_characters:
            # Open the unqualified name form
            new_member_name = unqualified_name_replace_t.show(enum_value_name, invalid_characters)

            # Did the user skip, or provide an empty string?
            if not new_member_name:
                skip = True
                break

            enum_value_name = new_member_name
            # Check if the user provided an invalid name
            invalid_characters = get_invalid_characters(enum_value_name)
        if skip:
            idaapi.msg("HashDB: Skipping hash result \"{}\" with value: {}\n".format(enum_value_name, hex(enum_value)))
            continue

        # Check if the name already exists in the database
        maximum_attempts = 256  # ENUM_MEMBER_ERROR_VALUE -> only allows 256 members with this value
        original_enum_value_name = enum_value_name
        for i in range(maximum_attempts):
            # Handle name suffixes
            if is_api:
                enum_value_name = f"{original_enum_value_name}_{i}"
            elif i:  # non API names that do not conflict with other names should not have a number suffix
                enum_value_name = f"{original_enum_value_name}_{i - 1}"  # begin indices with 0

            # Check if the enum name and value already exist (prevent duplicates)
            if (enum_value_name, enum_value) in fixed_hash_values.items():
                break

            # Check if the enum name already exists in the database
            if ida_name.get_name_ea(idaapi.BADADDR, enum_value_name) != idaapi.BADADDR or \
               get_enum_member_by_name(enum_value_name) != idaapi.BADADDR:
                continue

            # Add the values to the fixed hash list
            fixed_hash_values[enum_value_name] = enum_value
            break

    # Generate the enum definition
    enum_values = ",".join(f"{enum_value_name} = {enum_value:#x}"
                           for enum_value_name, enum_value in fixed_hash_values.items())

    # Support for 64-bit enums
    global HASHDB_ALGORITHM_SIZE
    is_64_bit_enum = BITS == 64 or HASHDB_ALGORITHM_SIZE == 64
    suffix = ": unsigned __int64" if is_64_bit_enum else ""  # handle 64-bit IDBs
    enum_definition = f"enum {enum_name + suffix} {{ {enum_values} }};"

    # Parse and import the enum into the database
    if ida_typeinf.idc_parse_types(enum_definition, 0):
        idaapi.msg(f"HashDB: Failed to parse {enum_name} into the database, please contact the developers.\n")
        return None

    enum_id = import_type(-1, enum_name)
    if enum_id == idaapi.BADNODE:
        idaapi.msg(f"HashDB: Failed to import {enum_name} into the database, please contact the developers.\n")
        return None

    # Return the enum id
    return enum_id


def generate_enum_name(prefix: str) -> str:
    """
    Generates an enum name from a prefix
    """
    global HASHDB_ALGORITHM
    return prefix + '_' + HASHDB_ALGORITHM


def make_const_enum(enum_id, hash_value):
    # We are in the disassembler we can set the enum directly
    ea = idc.here()
    start = idaapi.get_item_head(ea)
    # Determind if this is code or data/undefined
    if idc.is_code(idc.get_full_flags(ea)):
        # Find the operand position
        if idc.get_operand_value(ea,0) == hash_value:
            ida_bytes.op_enum(start, 0, enum_id, 0)
            return True
        elif idc.get_operand_value(ea,1) == hash_value:
            ida_bytes.op_enum(start, 1, enum_id, 0)
            return True
        else:
            return False
    else:
        ida_bytes.op_enum(start, 0, enum_id, 0)


def parse_highlighted_value():
    current_viewer = ida_kernwin.get_current_viewer()
    is_range_selected, _, _ = ida_kernwin.read_range_selection(current_viewer)
    if is_range_selected:
        raise HashDBError("Ranged selections are not supported when parsing highlighted values. "
                          "Were multiple lines selected?")

    highlight_result = ida_kernwin.get_highlight(current_viewer)
    if highlight_result is None:
        raise HashDBError("Invalid selection; nothing was selected.")

    identifier, flags = highlight_result
    if not flags:
        raise HashDBError("Invalid selection; did you manually range-select a value, or double-click a value?")

    # Handle various number formats
    match = re.match(r"^(?P<sign>[-+])?(?P<prefix>0|0x)?(?P<number>[0-9A-F]+)(?P<suffix>b|o|h|u|i64|ui64)?$",
                     identifier)
    if match is None:
        raise HashDBError(f"Failed to parse the value {identifier!r}, please submit a bug report.")

    # Determine the radix and number value
    sign, prefix, number, suffix = match.groups()
    if sign is None:
        sign = "+"

    radix = 10
    if suffix == "b":
        radix = 2  # binary
    elif suffix == "o" or prefix == "0" and not suffix == "h":
        radix = 8  # octal
    elif prefix == "0x" or suffix == "h":
        radix = 16  # hexadecimal

    try:
        return int(f"{sign}{number}", radix)
    except ValueError:
        raise HashDBError(f"Failed to parse the value {identifier!r} with a radix of {radix}.")


def determine_highlighted_type_size(ea: int) -> int:
    '''Guess the highlighted type and return the size in bytes.'''
    type = idaapi.idc_guess_type(ea)
    if type == '__int64':
        return 8
    if type == 'int':
        return 4
    if type == '__int16':
        return 2
    if type == 'char':
        return 1
    # If IDA couldn't guess the type (undefined, etc.) type will be empty
    return 0


def read_integer_from_db(ea: int, default_size: int = 0) -> int:
    '''
    Read the highlighted data from the database.
    Returns: [value, size, was_type_valid]
    '''
    type_size = determine_highlighted_type_size(ea)
    # 64-bit
    if type_size == 8 or (not type_size and default_size == 8):
        return [ida_bytes.get_64bit(ea), 8, bool(type_size)]
    # 32-bit
    if type_size == 4 or (not type_size and default_size == 4):
        return [ida_bytes.get_32bit(ea), 4, bool(type_size)]
    # 16-bit
    if type_size == 2 or (not type_size and default_size == 2):
        return [ida_bytes.get_16bit(ea), 2, bool(type_size)]
    # 8-bit and "undefined" values
    if type_size == 1 or (not type_size and not default_size) or (not type_size and default_size == 1):
        return [ida_bytes.get_byte(ea), 1, bool(type_size)]

    # Should never get executed
    raise HashDBError("Failed to read integer from database at location: {} with size {}.".format(hex(ea), default_size))


def convert_data_to_integer(ea, size: int = 0) -> int:
    '''
    Converts the data into a QWORD, DWORD, WORD, or BYTE based on the size provided
    '''
    global HASHDB_ALGORITHM_SIZE
    if not size:
        size = HASHDB_ALGORITHM_SIZE // 8

    if size == 8:
        ida_bytes.create_qword(ea, size, True)
    elif size == 4:
        ida_bytes.create_dword(ea, size, True)
    elif size == 2:
        ida_bytes.create_word(ea, size, True)
    elif size == 1:
        ida_bytes.create_byte(ea, size, True)
    return size


#--------------------------------------------------------------------------
# Global settings
#--------------------------------------------------------------------------
def global_settings():
    global HASHDB_API_URL
    global HASHDB_USE_XOR
    global HASHDB_XOR_VALUE
    global HASHDB_ALGORITHM
    global ENUM_PREFIX
    if HASHDB_ALGORITHM != None:
        algorithms = [[HASHDB_ALGORITHM, str(HASHDB_ALGORITHM_SIZE)]]
    else:
        algorithms = []
    settings_results = hashdb_settings_t.show(api_url=HASHDB_API_URL, 
                                              enum_prefix=ENUM_PREFIX,
                                              use_xor=HASHDB_USE_XOR,
                                              xor_value=HASHDB_XOR_VALUE,
                                              algorithms=algorithms)
    if settings_results:
        idaapi.msg("HashDB configured successfully!\nHASHDB_API_URL: %s\nHASHDB_USE_XOR: %s\nHASHDB_XOR_VALUE: %s\nHASHDB_ALGORITHM: %s\nHASHDB_ALGORITHM_SIZE: %s\n" % 
                   (HASHDB_API_URL, HASHDB_USE_XOR, hex(HASHDB_XOR_VALUE), HASHDB_ALGORITHM, HASHDB_ALGORITHM_SIZE))
    else:
        idaapi.msg("HashDB configuration cancelled!\n")
    return 


#--------------------------------------------------------------------------
# Set the algorithm and its size
#--------------------------------------------------------------------------
def set_algorithm(algorithm: str, size: int) -> bool:
    global HASHDB_ALGORITHM
    global HASHDB_ALGORITHM_SIZE

    # Type checks to prevent accidental errors
    if not isinstance(algorithm, str):
        idaapi.msg("HashDB encountered an error while trying to set the algorithm: provided algorithm is a string type: %s\n", algorithm)
        return False
    
    if isinstance(size, str):
        size = int(size)
    if not isinstance(size, int):
        idaapi.msg("HashDB encountered an error while trying to set the algorithm: provided size is not an integer: %s\n" % size)
        return False
    
    # More checks for supported sizes
    supported_algorithm_sizes = [32, 64]
    if size not in supported_algorithm_sizes or size % 8 != 0:
        idaapi.msg("HashDB encountered an error while trying to set the algorithm: the size provided was invalid: %s\n" % size)
        return False
    
    # Set the algorithm and size
    HASHDB_ALGORITHM = algorithm
    HASHDB_ALGORITHM_SIZE = size
    return True


def determine_algorithm_size(algorithm_type: str) -> str:
    size = 'Unknown'
    if algorithm_type is None:
        return size
    
    if algorithm_type == 'unsigned_int':
        size = '32'
    elif algorithm_type == 'unsigned_long':
        size = '64'
    return size


#--------------------------------------------------------------------------
# Set xor key
#--------------------------------------------------------------------------
def set_xor_key():
    """
    Set xor key from selection
    """
    global HASHDB_USE_XOR
    global HASHDB_XOR_VALUE
    try:
        xor_value = parse_highlighted_value()
        ida_kernwin.msg(f"HashDB: Set xor value to {xor_value:#x}\n")
    except HashDBError as exception:
        ida_kernwin.msg(f"HashDB ERROR: {exception}\n")
        return False
    HASHDB_XOR_VALUE = xor_value
    HASHDB_USE_XOR = True
    idaapi.msg("XOR key set: {}\n".format(hex(xor_value)))
    return True
    

#--------------------------------------------------------------------------
# Hash lookup
#--------------------------------------------------------------------------
def hash_lookup_done_handler(hash_list: Union[None, list], hash_value: int = None):
    global ENUM_PREFIX
    def add_enums_wrapper(enum_name, hash_list):
        nonlocal enum_id
        enum_id = add_enums(enum_name, hash_list)
        return 0 # execute_sync dictates an int return value
    
    if hash_list is None or hash_value is None:
        return

    # Parse the hash list
    hash_string = None
    if len(hash_list) == 1:
        hash_string = hash_list[0].get("string", {})
    else:
        # Multiple hashes found, allow the user to
        #  select the best match
        collisions = {}
        for string_match in hash_list:
            string_value = string_match.get("string", {})
            if string_value.get('is_api', False):
                collisions[string_value.get("api", "")] = string_value
            else:
                collisions[string_value.get("string", "")] = string_value
        
        # Execute the match_select_t form on the main thread
        def match_select_show(collision_strings):
            nonlocal selected_string
            selected_string = match_select_t.show(collision_strings)
            return 0 # execute_sync dictates an int return value

        selected_string = None
        match_select_callable = functools.partial(match_select_show, [*collisions.keys()])
        ida_kernwin.execute_sync(match_select_callable, ida_kernwin.MFF_FAST)
        if selected_string is None:
            return
        
        hash_string = collisions[selected_string]
    
    # Parse the string from the hash_string match
    string_value = ""
    if hash_string.get("is_api", False):
        string_value = hash_string.get("api", "")
    else:
        string_value = hash_string.get("string", "")

    # Handle empty string values
    if not len(string_value):
        string_value = "empty_string"

    # Hash found!
    idaapi.msg("HashDB: Hash match found: {}\n".format(string_value))

    # Add the hash to the global enum, and exit if we can't create it
    enum_id = None
    add_enums_callable = functools.partial(add_enums_wrapper, generate_enum_name(ENUM_PREFIX), [(string_value, hash_value, hash_string.get("is_api", False))])
    ida_kernwin.execute_sync(add_enums_callable, ida_kernwin.MFF_FAST)
    if enum_id is None:
        idaapi.msg("ERROR: Unable to create or find enum: {}\n".format(generate_enum_name(ENUM_PREFIX)))
        return
    
    # If the hash was pulled from the disassembly window
    # make the constant an enum 
    # TODO: I don't know how to do this in the decompiler window
    def make_const_enum_wrapper(enum_id, hash_value):
        if ida_kernwin.get_viewer_place_type(ida_kernwin.get_current_viewer()) == ida_kernwin.TCCPT_IDAPLACE:
            make_const_enum(enum_id, hash_value)
        return 0 # execute_sync dictates an int return value
    
    make_const_enum_wrapper_callable = functools.partial(make_const_enum_wrapper, enum_id, hash_value)
    ida_kernwin.execute_sync(make_const_enum_wrapper_callable, ida_kernwin.MFF_FAST)

    # Handle API hashes
    if not hash_string.get("is_api", False):
        return

    # Execute the api_import_select_t form on the main thread
    def api_import_select_show(string_value, module_list) -> int:
        nonlocal module_name
        module_name = api_import_select_t.show(string_value, module_list)
        return 0 # execute_sync dictates an int return value

    module_name = None
    api_import_select_callable = functools.partial(api_import_select_show, string_value, hash_string.get("modules", []))
    ida_kernwin.execute_sync(api_import_select_callable, ida_kernwin.MFF_FAST)
    if module_name is None:
        return

    # Import all of the hashes from the module and permutation
    module_hash_list = None
    try:
        global HASHDB_ALGORITHM, HASHDB_API_URL, HASHDB_REQUEST_TIMEOUT
        module_hash_list = get_module_hashes(module_name, HASHDB_ALGORITHM, hash_string.get("permutation", ""), HASHDB_API_URL, timeout=HASHDB_REQUEST_TIMEOUT)
    except requests.Timeout:
        idaapi.msg("ERROR: HashDB API module hashes request timed out.\n")
        logging.exception("API request to {} timed out.".format(HASHDB_API_URL))
        return

    # Add the hash list to the global enum
    global HASHDB_USE_XOR, HASHDB_XOR_VALUE
    enum_list = []
    for hash_entry in module_hash_list.get("hashes", []):
        hash = hash_entry.get("hash", 0)
        string_object = hash_entry.get("string", {})
        enum_list.append((string_object.get("api", string_object.get("string", "")), # name
                         hash ^ HASHDB_XOR_VALUE if HASHDB_USE_XOR else hash, # hash_value
                         True)) # is_api
    
    # Add hashes to enum
    enum_id = None
    add_enums_callable = functools.partial(add_enums_wrapper, generate_enum_name(ENUM_PREFIX), enum_list)
    ida_kernwin.execute_sync(add_enums_callable, ida_kernwin.MFF_FAST)
    if enum_id is None:
        idaapi.msg("ERROR: Unable to create or find enum: {}\n".format(generate_enum_name(ENUM_PREFIX)))
    else:
        idaapi.msg("HashDB: Added {} hashes for module {}\n".format(len(enum_list), module_name))


def hash_lookup_done(hash_list: Union[None, list] = None, hash_value: int = None):
    global HASHDB_REQUEST_LOCK
    hash_lookup_done_handler(hash_list, hash_value)

    # Release the lock
    HASHDB_REQUEST_LOCK.release()


def hash_lookup_error(exception: Exception):
    global HASHDB_REQUEST_LOCK
    exception_string = traceback.format_exc()
    logging.critical("hash_lookup_request errored: {}".format(exception_string))
    idaapi.msg("ERROR: HashDB hash scan failed: {}\n".format(exception_string))
    HASHDB_REQUEST_LOCK.release()


def hash_lookup_request(api_url: str, algorithm: str,
                        hash_value: int, xor_value: Union[None, int],
                        timeout: Union[int, float]):
    # Perform the request
    hash_results = None
    try:
        hash_results = get_strings_from_hash(algorithm, hash_value, xor_value if xor_value is not None else 0, api_url, timeout)
    except requests.Timeout:
        idaapi.msg("ERROR: HashDB API lookup hash request timed out.\n")
        logging.exception("API request to {} timed out:".format(HASHDB_API_URL))
        return None, None

    hash_list = hash_results.get("hashes", [])
    # Did `hashes` exist, was the array empty?
    if not hash_list:
        idaapi.msg("HashDB: No hash found for {}\n".format(hex(hash_value)))
        return None, None
    return hash_list, hash_value


def hash_lookup_run(timeout: Union[int, float] = 0) -> bool:
    # Check if an algorithm is selected
    global HASHDB_ALGORITHM, HASHDB_ALGORITHM_SIZE, HASHDB_API_URL, \
           ENUM_PREFIX, HASHDB_USE_XOR, HASHDB_XOR_VALUE
    if HASHDB_ALGORITHM is None:
        idaapi.warning("Please select a hash algorithm before using HashDB.")
        settings_results = hashdb_settings_t.show(api_url=HASHDB_API_URL, 
                                                  enum_prefix=ENUM_PREFIX,
                                                  use_xor=HASHDB_USE_XOR,
                                                  xor_value=HASHDB_XOR_VALUE)
        if settings_results:
            idaapi.msg("HashDB configured successfully!\n" +
                       "HASHDB_API_URL:        {}\n".format(HASHDB_API_URL) +
                       "HASHDB_USE_XOR:        {}\n".format(HASHDB_USE_XOR) +
                       "HASHDB_XOR_VALUE:      {}\n".format(hex(HASHDB_XOR_VALUE)) +
                       "HASHDB_ALGORITHM:      {}\n".format(HASHDB_ALGORITHM) +
                       "HASHDB_ALGORITHM_SIZE: {}\n".format(HASHDB_ALGORITHM_SIZE))
        else:
            idaapi.msg("HashDB configuration cancelled!\n")
            return True # Release the lock
    
    # Get the selected hash value
    try:
        hash_value = parse_highlighted_value()
        ida_kernwin.msg(f"HashDB: Found hash value: {hash_value:#x}\n")
    except HashDBError as exception:
        ida_kernwin.msg(f"HashDB ERROR: {exception}\n")
        return True  # release the lock

    # Lookup the hash and show a match select form
    worker = Worker(target=hash_lookup_request, args=(
        HASHDB_API_URL, HASHDB_ALGORITHM, hash_value, HASHDB_XOR_VALUE if HASHDB_USE_XOR else None, timeout),
        done_callback=hash_lookup_done, error_callback=hash_lookup_error)
    worker.start()
    return False # Do not release the lock


def hash_lookup():
    """
    Lookup a hash value from the highlighted text.

    The function will spawn a new thread with a timeout (`HASHDB_REQUEST_TIMEOUT`).
     While executing, the request lock is acquired.
    """
    # Check if we're already running a request
    global HASHDB_REQUEST_LOCK, HASHDB_REQUEST_TIMEOUT
    timeout_string = "{}".format(HASHDB_REQUEST_TIMEOUT) + " second{}".format('s' if HASHDB_REQUEST_TIMEOUT != 1 else "")
    if HASHDB_REQUEST_LOCK.locked():
        logging.debug("An async operation was requested, but the response lock was locked. Aborting.")
        ida_kernwin.info("Please wait until the previous request is finished.\n" +
                         "Requests timeout after {}.".format(timeout_string))
        return

    # Acquire the lock and execute the request
    HASHDB_REQUEST_LOCK.acquire()
    idaapi.msg("HashDB: Searching for a hash, please wait! Timeout: {}.\n".format(timeout_string))
    release_lock = hash_lookup_run(timeout=HASHDB_REQUEST_TIMEOUT)
    if release_lock:
        HASHDB_REQUEST_LOCK.release()


#--------------------------------------------------------------------------
# Dynamic IAT hash scan
# TODO: convert_values should be fetched from the UI (add a checkbox)
#--------------------------------------------------------------------------
def hash_scan_done(convert_values: bool = False, hash_list: Union[None, list] = None):
    global HASHDB_REQUEST_LOCK
    logging.debug("hash_scan_done callback invoked, result: {}".format("none" if hash_list is None else "{}".format(hash_list)))

    global ENUM_PREFIX
    def add_enums_wrapper(enum_name: str, hash_list):
        nonlocal enum_id
        enum_id = add_enums(enum_name, hash_list)
        return 0 # execute_sync dictates an int return value
    
    # Check if the `hash_scan_request` function failed (a caught exception should return `None`)
    if hash_list is not None:
        for hash_entry in hash_list:
            hashes = hash_entry["hashes"]
            hash_string_object = {}

            entries_count = len(hashes)
            if not entries_count:
                idaapi.msg("HashDB: Couldn't find any matches for hash value {} ({} bytes) at {}\n".format(hex(hash_entry["hash_value"]), hash_entry["size"], hex(hash_entry["ea"])))
                continue

            # Resolve the hash string object from the response
            if entries_count == 1:
                hash_string_object = hashes[0].get("string", {})
            else:
                # Check for collisions
                collisions = {}
                for entry in hashes:
                    string_object = entry.get("string", {})
                    if string_object.get("is_api", False):
                        collisions[string_object.get("api", "")] = string_object
                    else:
                        collisions[string_object.get("string", "")] = string_object
                
                # Execute the match_select_t form on the main thread
                def match_select_show(collision_strings):
                    nonlocal selected_string
                    selected_string = match_select_t.show(collision_strings)
                    return 0 # execute_sync dictates an int return value

                selected_string = None
                match_select_callable = functools.partial(match_select_show, [*collisions.keys()])
                ida_kernwin.execute_sync(match_select_callable, ida_kernwin.MFF_FAST)
                if selected_string is None:
                    HASHDB_REQUEST_LOCK.release() # Release the lock
                    return
                
                hash_string_object = collisions[selected_string]
            
            # Parse the string from hash_string_object
            hash_string_value = ""
            if hash_string_object.get("is_api", False):
                hash_string_value = hash_string_object.get("api", "")
            else:
                hash_string_value = hash_string_object.get("string", "")
            
            # Handle empty string values
            if not len(hash_string_value):
                hash_string_value = "empty_string"
            
            # Add hash to enum
            enum_id = None
            add_enums_callable = functools.partial(add_enums_wrapper, generate_enum_name(ENUM_PREFIX), [(hash_string_value, hash_entry["hash_value"], hash_string_object.get("is_api", False))])
            ida_kernwin.execute_sync(add_enums_callable, ida_kernwin.MFF_FAST)
            if enum_id is None:
                idaapi.msg("ERROR: Unable to create or find enum: {}\n".format(generate_enum_name(ENUM_PREFIX)))
                HASHDB_REQUEST_LOCK.release() # Release the lock
                return
            
            # Should we convert the values in the database?
            if convert_values:
                # Convert to integer (this step is required due to an IDA api bug - `ida_bytes.op_enum` will set the wrong size)
                def convert_to_integer(ea: int, size: int):
                    convert_data_to_integer(ea, size)

                convert_to_integer_callable = functools.partial(convert_to_integer, hash_entry["ea"], hash_entry["size"])
                ida_kernwin.execute_sync(convert_to_integer_callable, ida_kernwin.MFF_FAST)

                # Convert to enum
                def convert_to_enum(ea: int, enum_id: int):
                    NUMBER_OF_OPERANDS = 0
                    SERIAL = 0
                    ida_bytes.op_enum(ea, NUMBER_OF_OPERANDS, enum_id, SERIAL)
                    return 0 # execute_sync dictates an int return value
                
                convert_to_enum_callable = functools.partial(convert_to_enum, hash_entry["ea"], enum_id)
                ida_kernwin.execute_sync(convert_to_enum_callable, ida_kernwin.MFF_FAST)

                # Add a label
                def set_name(ea: int, name: str):
                    if not name: # is the name empty?
                        return 0 # execute_sync dictates an int return value
                    
                    # Does the name already exist? If so, modify it!
                    index = 1
                    suffix = ""
                    while idc.get_name_ea_simple(name + suffix) != idaapi.BADADDR:
                        suffix = "_{}".format(index)
                        index += 1

                    idc.set_name(ea, name + suffix, idc.SN_CHECK)
                    return 0 # execute_sync dictates an int return value
                
                set_name_callable = functools.partial(set_name, hash_entry["ea"], "ptr_" + hash_string_value)
                ida_kernwin.execute_sync(set_name_callable, ida_kernwin.MFF_FAST)
    
    # Release the lock
    HASHDB_REQUEST_LOCK.release()


def rate_limit_hit(seconds):
    spinner_chars = "|/-\\"
    for _ in range(seconds):
        for char in spinner_chars:
            sys.stdout.write(f'\r{char} Waiting 65 seconds to respect the HashDB API limit...')
            sys.stdout.flush()
            time.sleep(.25)


def hash_scan_error(exception: Exception):
    global HASHDB_REQUEST_LOCK
    exception_string = traceback.format_exc()
    logging.critical("hash_scan_request errored: {}".format(exception_string))
    idaapi.msg("ERROR: HashDB hash scan failed: {}\n".format(exception_string))
    HASHDB_REQUEST_LOCK.release()


def hash_scan_request(convert_values: bool, hash_list: list,
                            api_url: str, algorithm: str, xor_value: int,
                            timeout: Union[int, float]) -> Union[None, list]:
    api_calls = 0
    for hash_entry in hash_list:
        if api_calls == 99:
            print('\n')
            rate_limit_hit(65) # Wait 65 seconds to respect the API limit
            print('\n')
            api_calls = 0
        api_calls += 1
        try:
            hash_results = get_strings_from_hash(algorithm, hash_entry["hash_value"], xor_value if xor_value is not None else 0, api_url, timeout)
        except requests.Timeout:
            idaapi.msg("ERROR: HashDB API lookup scan request timed out.\n")
            logging.exception("API request to {} timed out:".format(HASHDB_API_URL))
            return None, None
        
        hash_entry["hashes"] = hash_results.get("hashes", [])
    return convert_values, hash_list


def hash_scan_run(convert_values: bool, timeout: Union[int, float] = 0) -> bool:
    # Only scan for data in the dissassembler
    if ida_kernwin.get_viewer_place_type(ida_kernwin.get_current_viewer()) != ida_kernwin.TCCPT_IDAPLACE:
        idaapi.msg("ERROR: Scan only available in dissassembler.\n")
        return True # Release the lock
    
    # Get the highlighted range
    start = idc.read_selection_start()
    end = idc.read_selection_end()
    if idaapi.BADADDR in (start, end):
        ea = idc.here()
        start = idaapi.get_item_head(ea)
        end = idaapi.get_item_end(ea)
    
    # If an algorithm isn't selected, give the user a chance to choose one
    global HASHDB_ALGORITHM, HASHDB_ALGORITHM_SIZE, HASHDB_API_URL, \
           ENUM_PREFIX, HASHDB_USE_XOR, HASHDB_XOR_VALUE
    if HASHDB_ALGORITHM is None:
        idaapi.warning("Please select a hash algorithm before using HashDB.")
        settings_results = hashdb_settings_t.show(api_url=HASHDB_API_URL, 
                                                  enum_prefix=ENUM_PREFIX,
                                                  use_xor=HASHDB_USE_XOR,
                                                  xor_value=HASHDB_XOR_VALUE)
        if settings_results:
            idaapi.msg("HashDB configured successfully!\n" +
                       "HASHDB_API_URL:        {}\n".format(HASHDB_API_URL) +
                       "HASHDB_USE_XOR:        {}\n".format(HASHDB_USE_XOR) +
                       "HASHDB_XOR_VALUE:      {}\n".format(hex(HASHDB_XOR_VALUE)) +
                       "HASHDB_ALGORITHM:      {}\n".format(HASHDB_ALGORITHM) +
                       "HASHDB_ALGORITHM_SIZE: {}\n".format(HASHDB_ALGORITHM_SIZE))
        else:
            idaapi.msg("HashDB configuration cancelled!\n")
            return True # Release the lock
    
    # Check for a valid algorithm size
    if not HASHDB_ALGORITHM_SIZE == 32 and not HASHDB_ALGORITHM_SIZE == 64:
        idaapi.msg("ERROR: Unexpected algorithm size provided: {}\n".format(HASHDB_ALGORITHM_SIZE))
        return True # Release the lock
    
    # Look through the selected range and lookup each (valid) entry
    def scan_range(start: int, end: int) -> list:
        """
        Find hash values in a given (highlighted) range.
         This function won't modify any data types in the database.
        
        Undefined types will be interpreted appropriately.
         As a result, the user is required to define the types if they
         expect valid results.
        """
        hash_values = []
        ea = start
        while ea < end:
            # Read the hash value and determine the step size:
            [hash_value, step_size, was_type_valid] = read_integer_from_db(ea)
            # If the type wasn't valid (undefined), convert it in the database
            #   and modify the hash value and step size accordingly:
            if convert_values and not was_type_valid:
                [hash_value, step_size, was_type_valid] = read_integer_from_db(ea, HASHDB_ALGORITHM_SIZE // 8)

            # Insert the hash value into the list
            hash_values.append({"ea": ea, "hash_value": hash_value, "size": step_size})

            # Next hash
            ea += step_size
        return hash_values
    
    hash_list = scan_range(start, end)
    for index, hash_entry in enumerate(hash_list, start=1):
        idaapi.msg("HashDB: [{}] Found hash value {} ({} bytes) at {}\n".format(index, hex(hash_entry["hash_value"]), hash_entry["size"], hex(hash_entry["ea"])))
    
    # Hunt all hashes, and provide the `hash_scan_done` callback with the results
    worker = Worker(target=hash_scan_request, args=(convert_values, hash_list,
                                                    HASHDB_API_URL, HASHDB_ALGORITHM,
                                                    HASHDB_XOR_VALUE if HASHDB_USE_XOR else None,
                                                    timeout),
                                                    done_callback=hash_scan_done, error_callback=hash_scan_error)
    worker.start()
    return False # Do not release the lock


def hash_scan(convert_values = True):
    """
    Scan for a dynamic hash table.

    The function will spawn a new thread with a timeout (`HASHDB_REQUEST_TIMEOUT`).
     While executing, the request lock is acquired.
    """
    # Check if we're already running a request
    global HASHDB_REQUEST_LOCK, HASHDB_REQUEST_TIMEOUT
    timeout_string = "{}".format(HASHDB_REQUEST_TIMEOUT) + " second{}".format('s' if HASHDB_REQUEST_TIMEOUT != 1 else "")
    if HASHDB_REQUEST_LOCK.locked():
        logging.debug("An async operation was requested, but the response lock was locked. Aborting.")
        ida_kernwin.info("Please wait until the previous request is finished.\n" +
                         "Requests timeout after {}.".format(timeout_string))
        return

    # Acquire the lock and execute the request
    HASHDB_REQUEST_LOCK.acquire()
    idaapi.msg("HashDB: Scanning for hashes, please wait! Timeout: {}.\n".format(timeout_string))
    release_lock = hash_scan_run(convert_values=convert_values, timeout=HASHDB_REQUEST_TIMEOUT)
    if release_lock:
        HASHDB_REQUEST_LOCK.release()


#--------------------------------------------------------------------------
# Algorithm search function
#--------------------------------------------------------------------------
def hunt_algorithm_done(response: Union[None, list] = None):
    global HASHDB_REQUEST_LOCK
    logging.debug("hunt_algorithm_done callback invoked, result: {}".format("none" if response is None else "{}".format(response)))

    # Display the result
    if response is not None:
        logging.debug("Displaying hash_result_form_t.")
        hunt_result_form_callable = functools.partial(hunt_result_form_t.show, response)
        ida_kernwin.execute_sync(hunt_result_form_callable, ida_kernwin.MFF_FAST)
    else:
        logging.debug("Couldn't find any algorithms that match the provided hash.")
        idaapi.msg("HashDB: Couldn't find any algorithms that match the provided hash.")
    
    # Release the lock
    HASHDB_REQUEST_LOCK.release()


def hunt_algorithm_error(exception: Exception):
    global HASHDB_REQUEST_LOCK
    exception_string = traceback.format_exc()
    logging.critical("hunt_algorithm_request errored: {}".format(exception_string))
    idaapi.msg("ERROR: HashDB hash scan failed: {}\n".format(exception_string))
    HASHDB_REQUEST_LOCK.release()


def hunt_algorithm_request(hash_value: int, timeout=None) -> Union[None, list]:
    """
    Perform the actual request, and provide the results to the
     `hunt_algorithm_done` callback.
    
    This function is required to be a coroutine for seamless timeout handling.
    """
    global HASHDB_REQUEST_LOCK, HASHDB_API_URL

    # Attempt to find matches
    match_results = None
    try:
        # Send the hunt request
        match_results = hunt_hash(hash_value, api_url=HASHDB_API_URL, timeout=timeout)
    except (requests.exceptions.Timeout, requests.exceptions.ReadTimeout):
        idaapi.msg("ERROR: HashDB API hunt hash request timed out.\n")
        logging.exception("API request to {} timed out.".format(HASHDB_API_URL))
        return None
    
    # Fix the results (algorithm sizes)
    # TODO: At the moment we have to fetch the algorithms again to determine their sizes
    #       (the hunt_result_form_t form expects the algorithm name and size)
    algorithms = None
    try:
        # Send the hunt request
        algorithms = get_algorithms(api_url=HASHDB_API_URL, timeout=timeout)
    except (requests.exceptions.Timeout, requests.exceptions.ReadTimeout):
        idaapi.msg("ERROR: HashDB API algorithms request timed out.\n")
        logging.exception("API request to {} timed out.".format(HASHDB_API_URL))
        return None
    results = []
    for match in match_results:
        for algorithm in algorithms:
            if match == algorithm[0]:
                results.append(algorithm)
                break
    
    # Return the results
    return results


def hunt_algorithm_run(timeout: Union[int, float] = 0) -> bool:
    global HASHDB_REQUEST_LOCK, HASHDB_USE_XOR, HASHDB_XOR_VALUE
    
    # Get the selected hash value
    try:
        hash_value = parse_highlighted_value()
    except HashDBError as exception:
        ida_kernwin.msg(f"HashDB ERROR: {exception}\n")
        logging.warning("Failed to parse a hash value from the highlighted text.")
        return True  # Release the lock
    
    # Xor option
    if HASHDB_USE_XOR:
        hash_value ^= HASHDB_XOR_VALUE

    # Hunt the algorithm and show the hunt result form
    worker = Worker(target=hunt_algorithm_request, args=(hash_value, timeout),
                    done_callback=hunt_algorithm_done, error_callback=hunt_algorithm_error)
    worker.start()
    return False # Do not release the lock


def hunt_algorithm():
    """
    Search for an algorithm using a hash value.

    The function will spawn a new thread with a timeout (`HASHDB_REQUEST_TIMEOUT`).
     While executing, the request lock is acquired.
    """
    # Check if we're already running a request
    global HASHDB_REQUEST_LOCK, HASHDB_REQUEST_TIMEOUT
    timeout_string = "{}".format(HASHDB_REQUEST_TIMEOUT) + " second{}".format('s' if HASHDB_REQUEST_TIMEOUT != 1 else "")
    if HASHDB_REQUEST_LOCK.locked():
        logging.debug("An async operation was requested, but the response lock was locked. Aborting.")
        ida_kernwin.info("Please wait until the previous request is finished.\n" +
                         "Requests timeout after {}.".format(timeout_string))
        return

    # Acquire the lock and execute the request
    HASHDB_REQUEST_LOCK.acquire()
    idaapi.msg("HashDB: Hunting for a hash algorithm, please wait! Timeout: {}.\n".format(timeout_string))
    release_lock = hunt_algorithm_run(timeout=HASHDB_REQUEST_TIMEOUT)
    if release_lock:
        HASHDB_REQUEST_LOCK.release()


#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class HashDB_Plugin_t(idaapi.plugin_t):
    """
    IDA Plugin for HashDB lookup service
    """
    comment = "HashDB Lookup Service"
    help = ""
    wanted_name = PLUGIN_NAME
    # We only want a hotkey for the actual hash lookup
    wanted_hotkey = ''
    flags = idaapi.PLUGIN_KEEP
    terminated = False

    #--------------------------------------------------------------------------
    # Plugin Overloads
    #--------------------------------------------------------------------------
    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        global p_initialized, HASHDB_PLUGIN_OBJECT

        # Check if already initialized 
        if p_initialized is False:
            p_initialized = True
            ## Print a nice header
            print("=" * 80)
            print(r"   _   _           _    ____________ ")
            print(r"  | | | |         | |   |  _  \ ___ \ ")
            print(r"  | |_| | __ _ ___| |__ | | | | |_/ /")
            print(r"  |  _  |/ _` / __| '_ \| | | | ___ \ ")
            print(r"  | | | | (_| \__ \ | | | |/ /| |_/ /")
            print(r"  \_| |_/\__,_|___/_| |_|___/ \____/ ")
            print("")                 
            print("\nHashDB v{0} by @herrcore".format(VERSION))
            print("\nHashDB search shortcut key is {0}".format(PLUGIN_HOTKEY))
            print("=" * 80)
            # Load saved settings if they exist
            load_settings()
            # initialize the menu actions our plugin will inject
            self._init_action_hash_lookup()
            self._init_action_set_xor()
            self._init_action_hunt()
            self._init_action_iat_scan()
            # initialize plugin hooks
            self._init_hooks()

            HASHDB_PLUGIN_OBJECT = self
            return idaapi.PLUGIN_KEEP


    def run(self, arg):
        """
        This is called by IDA when the plugin is run from the plugins menu
        """
        global_settings()
    

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """
        # Already terminated?
        if self.terminated:
            return
        
        # Save settings
        save_settings()

        # Unhook our plugin hooks
        self._hooks.unhook()

        # Unregister our actions & free their resources
        self._del_action_hash_lookup()
        self._del_action_set_xor()
        self._del_action_hunt()
        self._del_action_iat_scan()

        # Done
        self.terminated = True
        idaapi.msg("HashDB: {} terminated...\n".format(self.wanted_name))


    #--------------------------------------------------------------------------
    # IDA Actions
    #--------------------------------------------------------------------------
    ACTION_HASH_LOOKUP  = "hashdb:hash_lookup"
    ACTION_SET_XOR  = "hashdb:set_xor"
    ACTION_HUNT  = "hashdb:hunt"
    ACTION_IAT_SCAN = "hashdb:iat_scan"

    def _init_action_hash_lookup(self):
        """
        Register the hash lookup action with IDA.
        """
        action_desc = idaapi.action_desc_t( self.ACTION_HASH_LOOKUP,         # The action name.
                                            "HashDB Lookup",                     # The action text.
                                            IDACtxEntry(hash_lookup),        # The action handler.
                                            PLUGIN_HOTKEY,                  # Optional: action shortcut
                                            "Lookup hash",   # Optional: tooltip
                                            HASH_ICON
                                            )
        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"


    def _init_action_set_xor(self):
        """
        Register the set xor action with IDA.
        """
        action_desc = idaapi.action_desc_t(
            self.ACTION_SET_XOR,         # The action name.
            "HashDB set XOR key",                     # The action text.
            IDACtxEntry(set_xor_key),        # The action handler.
            None,                  # Optional: action shortcut
            "Set XOR key",   # Optional: tooltip
            XOR_ICON
        )
        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"


    def _init_action_hunt(self):
        """
        Register the hunt action with IDA.
        """
        action_desc = idaapi.action_desc_t(
            self.ACTION_HUNT,         # The action name.
            "HashDB Hunt Algorithm",                     # The action text.
            IDACtxEntry(hunt_algorithm),        # The action handler.
            None,                  # Optional: action shortcut
            "Identify algorithm based on hash",   # Optional: tooltip
            HUNT_ICON
        )
        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"


    def _init_action_iat_scan(self):
        """
        Register the scan action with IDA.
        """
        action_desc = idaapi.action_desc_t(
            self.ACTION_IAT_SCAN,         # The action name.
            "HashDB Scan IAT",                     # The action text.
            IDACtxEntry(hash_scan),        # The action handler.
            None,                  # Optional: action shortcut
            "Scan dynamic import address table hashes",   # Optional: tooltip
            SCAN_ICON
        )
        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"

    
    def _del_action_hash_lookup(self):
        idaapi.unregister_action(self.ACTION_HASH_LOOKUP)


    def _del_action_set_xor(self):
        idaapi.unregister_action(self.ACTION_SET_XOR)


    def _del_action_hunt(self):
        idaapi.unregister_action(self.ACTION_HUNT)

    def _del_action_iat_scan(self):
        idaapi.unregister_action(self.ACTION_IAT_SCAN)


    #--------------------------------------------------------------------------
    # Initialize Hooks
    #--------------------------------------------------------------------------

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


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


#------------------------------------------------------------------------------
# Plugin Hooks
#------------------------------------------------------------------------------
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_actions(widget, popup, idaapi.get_widget_type(widget))
        return 0

    def hxe_callback(self, event, *args):
        """
        HexRays event callback.
        We lump this under the (UI) Hooks class for organizational reasons.
        """

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

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

            attach_action_to_popup(
                form,
                popup,
                HashDB_Plugin_t.ACTION_HASH_LOOKUP,
                "HashDB Lookup",
                SETMENU_APP,
            )
            attach_action_to_popup(
                form,
                popup,
                HashDB_Plugin_t.ACTION_SET_XOR,
                "HashDB set XOR key",
                SETMENU_APP,
            )
            attach_action_to_popup(
                form,
                popup,
                HashDB_Plugin_t.ACTION_HUNT,
                "HashDB Hunt Algorithm",
                SETMENU_APP,
            )
            attach_action_to_popup(
                form,
                popup,
                HashDB_Plugin_t.ACTION_IAT_SCAN,
                "HashDB Scan IAT",
                SETMENU_APP,
            )

        # done
        return 0

#------------------------------------------------------------------------------
# Prefix Wrappers
#------------------------------------------------------------------------------
def inject_actions(form, popup, form_type):
    """
    Inject actions to popup menu(s) based on context.
    """

    #
    # disassembly window
    #

    if (form_type == BWN_DISASM) or (form_type == BWN_PSEUDOCODE):
        # insert the action entry into the menu
        #

        attach_action_to_popup(
            form,
            popup,
            HashDB_Plugin_t.ACTION_HASH_LOOKUP,
            "HashDB Lookup",
            SETMENU_APP
        )

        attach_action_to_popup(
            form,
            popup,
            HashDB_Plugin_t.ACTION_SET_XOR,
            "HashDB set XOR key",
            SETMENU_APP
        )

        attach_action_to_popup(
            form,
            popup,
            HashDB_Plugin_t.ACTION_HUNT,
            "HashDB Hunt Algorithm",
            SETMENU_APP
        )
        if form_type != BWN_PSEUDOCODE:
            attach_action_to_popup(
                form,
                popup,
                HashDB_Plugin_t.ACTION_IAT_SCAN,
                "HashDB Scan IAT",
                SETMENU_APP
            )

    # done
    return 0

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

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


#--------------------------------------------------------------------------
# Plugin Registration
#--------------------------------------------------------------------------

# Global flag to ensure plugin is only initialized once
p_initialized = False

# Global plugin object
HASHDB_PLUGIN_OBJECT = None

# Register IDA plugin
def PLUGIN_ENTRY():
    return HashDB_Plugin_t()


repository: ramikg/uncertaintifier
file: uncertaintifier.py
description: Add question marks to Hex-Rays local variable names
wanted_name: Uncertaintifier
comment: Uncertaintifier
created_at: 2021-07-19 19:01:19+00:00
pushed_at: 2023-02-21 14:30:46+00:00
url:
https://github.com/ramikg/uncertaintifier/blob/676a484cd847bccfef3a3972f8b82a09f38ab674/uncertaintifier.py

content:
import ida_idaapi
import ida_hexrays
import ida_kernwin

REGISTER_CONTEXT_MENU_ITEMS = True

INC_ACTION_NAME = INC_ACTION_DESCRIPTION = 'Increase uncertainty'
INC_ACTION_HOTKEY = 'Ctrl-/'
DEC_ACTION_NAME = DEC_ACTION_DESCRIPTION = 'Decrease uncertainty'
DEC_ACTION_HOTKEY = 'Ctrl-Alt-/'


class UncertaintifierFailedGettingVariableError(Exception):
    pass


def _get_lvar(widget):
    vdui = ida_hexrays.get_widget_vdui(widget)
    if vdui is not None:
        lvar = vdui.item.get_lvar()
        if lvar is not None:
            return lvar

    raise UncertaintifierFailedGettingVariableError()


def _is_pseudocode_lvar(widget):
    if ida_kernwin.get_widget_type(widget) != ida_kernwin.BWN_PSEUDOCODE:
        return False

    try:
        _get_lvar(widget)
        return True
    except UncertaintifierFailedGettingVariableError:
        return False


class QuestionMarkAdderRemover(ida_kernwin.action_handler_t):
    QUESTION_MARK = '?'

    def update(self, ctx):
        if self._should_enable_action(ctx):
            return ida_kernwin.AST_ENABLE
        else:
            return ida_kernwin.AST_DISABLE

    def _should_enable_action(self, ctx):
        return _is_pseudocode_lvar(ctx.widget)

    @staticmethod
    def _get_lvar_name(ctx):
        try:
            return _get_lvar(ctx.widget).name
        except UncertaintifierFailedGettingVariableError:
            print('Failed getting highlighted variable name. Please try again.')
            raise

    @staticmethod
    def _rename_lvar(ctx, new_name):
        IDAAPI_RENAME_LVAR_IS_USER_NAME = 1

        lvar = _get_lvar(ctx.widget)
        vdui = ida_hexrays.get_widget_vdui(ctx.widget)

        vdui.rename_lvar(lvar, new_name, IDAAPI_RENAME_LVAR_IS_USER_NAME)


class QuestionMarkAdder(QuestionMarkAdderRemover):
    def activate(self, ctx):
        try:
            old_name = self._get_lvar_name(ctx)
        except UncertaintifierFailedGettingVariableError:
            return

        self._rename_lvar(ctx, old_name + self.QUESTION_MARK)


class QuestionMarkRemover(QuestionMarkAdderRemover):
    def activate(self, ctx):
        try:
            old_name = self._get_lvar_name(ctx)
        except UncertaintifierFailedGettingVariableError:
            return

        assert old_name.endswith(self.QUESTION_MARK)
        self._rename_lvar(ctx, old_name[:-1])

    def _should_enable_action(self, ctx):
        # super() not used for Python 2 compatibility
        if not super(type(self), self)._should_enable_action(ctx):
            return False

        old_name = _get_lvar(ctx.widget).name
        return old_name.endswith(self.QUESTION_MARK)


class ContextMenuHooks(ida_kernwin.UI_Hooks):
    def finish_populating_widget_popup(self, widget, popup):
        if _is_pseudocode_lvar(widget):
            ida_kernwin.attach_action_to_popup(widget, popup, INC_ACTION_NAME, None)
            ida_kernwin.attach_action_to_popup(widget, popup, DEC_ACTION_NAME, None)


class Uncertaintifier(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    comment = 'Uncertaintifier'
    help = ''
    wanted_name = 'Uncertaintifier'
    wanted_hotkey = ''

    _hooks = ContextMenuHooks()

    def init(self):
        question_mark_adder_desc = ida_kernwin.action_desc_t(INC_ACTION_NAME,
                                                             INC_ACTION_DESCRIPTION,
                                                             QuestionMarkAdder(),
                                                             INC_ACTION_HOTKEY)
        question_mark_remover_desc = ida_kernwin.action_desc_t(DEC_ACTION_NAME,
                                                               DEC_ACTION_DESCRIPTION,
                                                               QuestionMarkRemover(),
                                                               DEC_ACTION_HOTKEY)

        ida_kernwin.register_action(question_mark_adder_desc)
        ida_kernwin.register_action(question_mark_remover_desc)

        if REGISTER_CONTEXT_MENU_ITEMS:
            self._hooks.hook()

        return ida_idaapi.PLUGIN_KEEP

    def run(self):
        pass

    def term(self):
        self._hooks.unhook()


def PLUGIN_ENTRY():
    return Uncertaintifier()


repository: WanQingGit/IDA7script
file: oregami/typeregter_plugin.py
description: 常用的比较实用的ida插件脚本
wanted_name: TypeREGter
comment: TypeREGter
created_at: 2021-03-21 01:50:41+00:00
pushed_at: 2021-04-03 13:35:03+00:00
url:
https://github.com/WanQingGit/IDA7script/blob/c81b5fcb613201b7f03a75e5865cfbd00c845b5c/oregami/typeregter_plugin.py

content:
import ida_nalt
import ida_kernwin
import sark
import re
import idc
import oregami.oregami_gen as oregami_gen
import ida_frame
import ida_funcs
from idaapi import PluginForm
from PyQt5 import QtCore, QtGui, QtWidgets
import _ida_kernwin
import idaapi
from oregami.oregami_gen import *
from oregami.regname import *
import ida_struct

#TODO: if rw, dont change type
  
class TyperegterPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "TypeREGter"
    help = "Set type for regs in their usage frame - only when used as a specific variable"
    wanted_name = "TypeREGter"
    wanted_hotkey = "Shift+T"

    def init(self):
        self._prev_struct_name = ""
        
        proc_name = sark.idaapi.get_inf_structure().procName
        if proc_name not in supported_procs:
            print 'Typeregter plugin: No full support for this processor. Will do our best :)'
            #return idaapi.PLUGIN_SKIP

        return idaapi.PLUGIN_OK
        
    def term(self):
        pass

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


        typeregterPluginStarter(start)
        #regRefPluginStarter(start, tabMode=True, recursive_bool=False)
        

def PLUGIN_ENTRY():
    return TyperegterPlugin()
    
    
    
def AskStr(def_val, q_str):
    #without this line, waiting a few seconds for user input will cause the 'Please wait...' msgbox to jump and make the system freeze.
    orig_timeout = idaapi.set_script_timeout(0x7fffffff)
    ret = idc.AskStr(def_val, q_str)
    idaapi.set_script_timeout(orig_timeout)
    return ret
    
def get_regs(ea, all=True):
    reg = get_cursor_reg(ea)    

    if reg is None:
        return {}
        
    reg_set = set([reg])
    
    regs = proc_ops.get_basic_regs(ea, reg_set)
    
    
    if all:
        return regs
        
    if len(regs)<=1:
        return regs
        
    #need to pick only one reg
    reg = sorted(list(regs))[0] #first reg (alphabetically)
    reg = get_reg_user_name(ea, reg) #if renamed before, use it as default
    reg = AskStr(reg, 'Pick register')
    

    if reg is None: #exit
        return None
        
    #TODO: add check for user named reg
    if reg in regs:
        return [reg]
        
    for e_reg in regs:
        if reg==e_reg or reg==get_reg_user_name(ea, e_reg):
            return [e_reg]
    else:
        print 'No such reg in line'
        return []

def get_type():
    #without this line, waiting a few seconds for user input will cause the 'Please wait...' msgbox to jump and make the system freeze.
    orig_timeout = idaapi.set_script_timeout(0x7fffffff)
    
    struc_type = ida_kernwin.choose_struc('Choose type')
    idaapi.set_script_timeout(orig_timeout)
    
    if struc_type is None:
        return None
        
    #print struc_type
    type_name = ida_struct.get_struc_name(struc_type.id)
    return type_name
    
    print type(a)
    type_name = AskStr('', 'Type for %s' % get_reg_full_name(ea, reg))
    if type_name is None:
        print 'Escaped'
        return None
        
    if ida_struct.get_struc_id(type_name)==0xffffffff:
        print 'Not a type'
        return None
        
    return type_name

def typeregterPluginStarter(ea):
    #Get wanted reg
    regs = get_regs(ea, False)
    
    if regs is None:
        print 'Escaped'
        return
        
    regs = sorted(list(regs))
    
    assert len(regs)<=1
    if len(regs)==0:
        print 'No regs'
        return
    
    reg = regs[0]
    
    #Get type name    
    type_name = get_type()
    
    if type_name is None:
        return
        

    #Do renaming
    typeset_reg(ea, reg, type_name)
    

    
########################             
#### Type set logic ####
########################
        
def get_cur_type(ea, reg):
    opnds = sark.Line(ea).insn.operands
    for i in range(len(opnds)):
        if '[%s]' % get_reg_user_name(ea, reg) in opnds[i].text:
            str_id = ida_struct.get_struc_id(type_name)
            idc.op_stroff(ea, i, str_id, 0)
        
def typeset_reg_for_address(ea, reg, type_name):
    opnds = sark.Line(ea).insn.operands
    for i in range(len(opnds)):
        if ('[%s]' % get_reg_user_name(ea, reg) in opnds[i].text) or ('[%s,' % get_reg_user_name(ea, reg) in opnds[i].text):
            str_id = ida_struct.get_struc_id(type_name)
            idc.op_stroff(ea, i, str_id, 0)
            ida_nalt.set_aflags(ea, ida_nalt.get_aflags(ea) | ida_nalt.AFL_ZSTROFF)
    
def typeset_reg(ea, reg, type_name):
    found_lines, found_breaks, found_outbreaks = oregami_gen.get_refs(ea, reg)
    #eas = list(set(found_lines.keys()) - set(found_breaks.keys()))
    eas = found_lines.keys()
    
    t_eas = []
    for ea in found_lines.keys():
        if found_lines[ea]=='r':
            t_eas += [ea]
            
    
    for ea in t_eas:
        typeset_reg_for_address(ea, reg, type_name)
    
    
###########
## debug ##
###########

def debug(s):
    #return
    print s
    

#if reg was renamed - get full name    
def get_reg_full_name(ea, reg):
    if reg not in get_reg_names(ea):
        return reg
        
    return '%s {%s}' % (get_reg_names(ea)[reg], reg)
    
    
def get_reg_names(ea):
    from idautils import GetRegisterList
    from ida_frame import find_regvar
    from ida_funcs import get_func
    
    d = {}

    for canon in GetRegisterList():
        regvar = find_regvar(get_func(ea), ea, canon)
        if regvar:
            d[canon] = regvar.user
            
    return d  
    
    


repository: Sandhya-hub/langflow
file: venv/Lib/site-packages/qianfan/common/client/plugin.py
created_at: 2024-08-12 08:13:38+00:00
pushed_at: 2024-08-12 08:43:39+00:00
url:
https://github.com/Sandhya-hub/langflow/blob/ab71809c697fe0409b4bb05524920c30385e5602/venv/Lib/site-packages/qianfan/common/client/plugin.py

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


import json
import os
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import unquote, urlparse

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

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


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

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


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

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

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

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

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

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

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

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

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

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

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

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

                break

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

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

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

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

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

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

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

                    live.update(Group(*render_list))

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


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


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

    extra_args = {}

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

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

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

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


repository: PwCUK-CTO/SmartJump
file: plugins/SmartJump.py
description: IDA Pro plugin to enhance the 'g' keyboard shortcut
wanted_name: SmartJumper
comment: Smart IDA jumping
created_at: 2020-09-14 16:27:46+00:00
pushed_at: 2023-07-24 09:59:17+00:00
url:
https://github.com/PwCUK-CTO/SmartJump/blob/4e32c6a26bac67045d6ac248e12c5bb8fd60d9c0/plugins/SmartJump.py

content:
'''
   Copyright 2023 PwCIL

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
'''

import idc, idaapi
import SmartJump.lex as lex
import SmartJump.yacc as yacc
import re

#utilities
arch_size = 32
if idaapi.idainfo_is_64bit():
    arch_size = 64

#support x86 and x64 so check which pointers we will be accessing
def get_pointer(address):
    if arch_size == 32:
        return idaapi.get_dword(address)
    else:
        return idaapi.get_qword(address)

tokens = ('HERE', 'HEXADDR', 'ADDR', 'PLUS', 'LBRACKET', 'RBRACKET', 'LPAREN', 'RPAREN', 'MINUS', 'TIMES', 'DIVIDE',
    'EAX', 'EBX', 'ECX', 'EDX', 'EIP', 'ESI', 'EDI', 'EBP', 'ESP',
    'RAX', 'RBX', 'RCX', 'RDX', 'RIP', 'RSI', 'RDI', 'RBP', 'RSP', 'R8', 'R9', 'R10', 'R11', 'R12', 'R13', 'R14', 'R15', 'NAME',
    )

t_PLUS      = r'\+'
t_MINUS     = r'-'
t_LBRACKET  = r'\['
t_RBRACKET  = r'\]'
t_TIMES     = r'\*'
t_DIVIDE    = r'/'
t_LPAREN    = r'\('
t_RPAREN    = r'\)'

def t_HERE(t):
    r'here(\(\))?'
    t.value = idc.here()
    return t

def t_RAX(t):
    r'rax'
    t.value = idaapi.get_reg_val("RAX")
    return t

def t_RBX(t):
    r'rbx'
    t.value = idaapi.get_reg_val("RBX")
    return t

def t_RCX(t):
    r'rcx'
    t.value = idaapi.get_reg_val("RCX")
    return t

def t_RDX(t):
    r'rdx'
    t.value = idaapi.get_reg_val("RDX")
    return t

def t_RIP(t):
    r'rip'
    t.value = idaapi.get_reg_val("RIP")
    return t

def t_RSI(t):
    r'rsi'
    t.value = idaapi.get_reg_val("RSI")
    return t

def t_RDI(t):
    r'rdi'
    t.value = idaapi.get_reg_val("RDI")
    return t

def t_RBP(t):
    r'rbp'
    t.value = idaapi.get_reg_val("RBP")
    return t

def t_RSP(t):
    r'rsp'
    t.value = idaapi.get_reg_val("RSP")
    return t

def t_R8(t):
    r'r8'
    t.value = idaapi.get_reg_val("R8")
    return t

def t_R9(t):
    r'r9'
    t.value = idaapi.get_reg_val("R9")
    return t

def t_R10(t):
    r'r10'
    t.value = idaapi.get_reg_val("R10")
    return t

def t_R11(t):
    r'r11'
    t.value = idaapi.get_reg_val("R11")
    return t

def t_R12(t):
    r'r12'
    t.value = idaapi.get_reg_val("R12")
    return t

def t_R13(t):
    r'r13'
    t.value = idaapi.get_reg_val("R13")
    return t

def t_R14(t):
    r'r14'
    t.value = idaapi.get_reg_val("R14")
    return t

def t_R15(t):
    r'r15'
    t.value = idaapi.get_reg_val("R15")
    return t

def t_EAX(t):
    r'eax'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("EAX")
    else:
        t.value = (idaapi.get_reg_val("RAX")) & 0xffffffff
    return t

def t_EBX(t):
    r'ebx'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("EBX")
    else:
        t.value = (idaapi.get_reg_val("RBX")) & 0xffffffff
    return t

def t_ECX(t):
    r'ecx'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("ECX")
    else:
        t.value = (idaapi.get_reg_val("RCX")) & 0xffffffff
    return t

def t_EDX(t):
    r'edx'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("EDX")
    else:
        t.value = (idaapi.get_reg_val("RDX")) & 0xffffffff
    return t

def t_EIP(t):
    r'eip'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("EIP")
    else:
        t.value = (idaapi.get_reg_val("RIP")) & 0xffffffff
    return t

def t_ESI(t):
    r'esi'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("ESI")
    else:
        t.value = (idaapi.get_reg_val("RSI")) & 0xffffffff
    return t

def t_EDI(t):
    r'edi'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("EDI")
    else:
        t.value = (idaapi.get_reg_val("RDI")) & 0xffffffff
    return t

def t_EBP(t):
    r'ebp'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("EBP")
    else:
        t.value = (idaapi.get_reg_val("RBP")) & 0xffffffff
    return t

def t_ESP(t):
    r'esp'
    if arch_size == 32:
        t.value = idaapi.get_reg_val("ESP")
    else:
        t.value = (idaapi.get_reg_val("RSP")) & 0xffffffff
    return t

def t_NAME(t):
    r'[a-zA-Z_\.][a-zA-Z0-9_\.]*'
    temp = idc.get_name_ea_simple(t.value)
    if temp == idaapi.BADADDR:
        t.value = int(t.value, 16)
    else:
        t.value = temp
    return t

def t_HEXADDR(t):
    r'0x[0-9a-f]+'
    t.value = int(t.value[2:], 16)
    return t

def t_ADDR(t):
    r'[0-9a-f]+'
    t.value = int(t.value, 16)
    return t

t_ignore = " \t"

def t_newline(t):
    r'\n+'
    t.lexer.lineno += t.value.count("\n")

def t_error(t):
    debug_out(f"Illegal character {t.value[0]!r}")
    t.lexer.skip(1)

glob_lex = lex.lex(reflags=re.I)

precedence = (
    ('left','PLUS','MINUS'),
    ('left','TIMES','DIVIDE'),
)

def p_expression_binop(p):
    '''expression : expression PLUS expression
                | expression MINUS expression
                | expression TIMES expression
                | expression DIVIDE expression'''
    if p[2] == '+'  : p[0] = p[1] + p[3]
    elif p[2] == '-': p[0] = p[1] - p[3]
    elif p[2] == '*': p[0] = p[1] * p[3]
    elif p[2] == '/': p[0] = p[1] / p[3]
    
def p_expression_group(p):
    'expression : LPAREN expression RPAREN'
    p[0] = p[2]

def p_expression_deref(p):
    'expression : LBRACKET expression RBRACKET'
    p[0] = get_pointer(p[2])

def p_expression_hexaddr(p):
    'expression : HEXADDR'
    p[0] = p[1]

def p_expression_addr(p):
    'expression : ADDR'
    p[0] = p[1]

def p_expression_reg(p):
    '''expression : RAX
                | RBX
                | RCX
                | RDX
                | RSI
                | RDI
                | RBP
                | RSP
                | RIP
                | R8
                | R9
                | R10
                | R11
                | R12
                | R13
                | R14
                | R15
                | EAX
                | EBX
                | ECX
                | EDX
                | ESI
                | EDI
                | EBP
                | ESP
                | EIP
                | HERE'''
    p[0] = p[1]

def p_expression_name(p):
    'expression : NAME'
    p[0] = p[1]

def p_error(p):
    print(f"Syntax error at {p.value!r}")

glob_parser = yacc.yacc()

def debug_out(str):
    print ("[SMARTJUMPER]: %s" % str)

class SmartJump_t(idaapi.plugin_t):
    flags = 0
    comment = "Smart IDA jumping"
    wanted_hotkey = 'g'
    help = "Runs by replacing Go command when pressing G"
    wanted_name = "SmartJumper"
    lexer = None
    parser = None
    
    def init(self):
        global glob_lex
        global glob_parser
        if idaapi.ph_get_id() != idaapi.PLFM_386:
            return idaapi.PLUGIN_SKIP
        debug_out("Loading Parsers")
        
        self.lexer = glob_lex
        self.parser = glob_parser
        debug_out("Loaded SmartJumper")
        return idaapi.PLUGIN_KEEP

    def run(self, arg=0):
        jump_str = idaapi.ask_str("", 0, "Jump expression...")
        if jump_str != None:
            try:
                open_deref = jump_str.count("[")
                close_deref = jump_str.count("]")
                if close_deref < open_deref:
                    jump_str += "]" * (open_deref-close_deref)
                if open_deref < close_deref:
                    debug_out("mismatched dereferences")
                else:
                    result = self.parser.parse(jump_str, lexer=self.lexer)
                    debug_out("resolved to %08x" % result)
                    idaapi.jumpto(result)
            except:
                debug_out("problem parsing")

    def term(self):
        return

def PLUGIN_ENTRY():
    return SmartJump_t()


repository: Krietz7/IDA-VarViewer
file: src/IDA-VarViewer.py
description: IDA plugin which provided a form to observe stack variable and reg variable more clearly
wanted_name: VarViewer
created_at: 2024-11-08 14:44:46+00:00
pushed_at: 2025-03-30 14:29:05+00:00
url:
https://github.com/Krietz7/IDA-VarViewer/blob/7fdb5a3bf29bc6a7cd3f6e5f6c569050dd6fd365/src/IDA-VarViewer.py

content:
import idaapi

from VarViewer.config import *
from VarViewer.dbg_stack import CpuInfo
from VarViewer.StackViewer import *
from VarViewer.VariableViewer import *

class MenuContext(idaapi.action_handler_t):
    @classmethod
    def get_name(cls):
        return cls.__name__

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

    @classmethod
    def register(cls, plugin, label, hotkey):
        cls.plugin = plugin
        cls.label = label
        instance = cls()
        return idaapi.register_action(idaapi.action_desc_t(cls.get_name(),instance.get_label(),instance,hotkey))

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

    def activate(self,ctx):
        pass

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


class StackViewer_Menu(MenuContext):
    @classmethod
    def activate(cls,ctx):
        k = StackViewer()
        k.Show("Stack Viewer")


class VariableViewer_Menu(MenuContext):
    @classmethod
    def activate(cls,ctx):
        k = VariableViewer()
        k.Show("Variable Viewer")

class VarViewer(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = ""
    help = ""
    wanted_name = "VarViewer"
    wanted_hotkey = ""

    def __init__(self):
        pass

    def init(self):
        CpuInfo.create_instance()
        if idaapi.IDA_SDK_VERSION < 770:
            return idaapi.PLUGIN_SKIP

        StackViewer_Menu.register(self, "Open Stack View",STACK_VIEW_HOTKEY)
        VariableViewer_Menu.register(self, "Open Variable View",VARIABLE_VIEW_HOTKEY)

        idaapi.attach_action_to_menu("Debugger/Debugger windows/Stack Viewer", StackViewer_Menu.get_name(), idaapi.SETMENU_APP)
        idaapi.attach_action_to_menu("Debugger/Debugger windows/Variable Viewer", VariableViewer_Menu.get_name(), idaapi.SETMENU_APP)

        return idaapi.PLUGIN_OK

    def run(self,arg):
        print("Initialized:",CpuInfo.instance.procname,CpuInfo.instance.cpu_struct)

        k = StackViewer()
        k.Show(WIDGET_TITLE)

    def term(self):
        pass

def PLUGIN_ENTRY():
    return VarViewer()


repository: BroilFish/LazyIDA_mem
file: LazyIDA2.py
description: LazyIDA脚本添加 内存指针调整的快捷键功能
url: https://github.com/BroilFish/LazyIDA_mem/blob/25e6d1de0aaf6c6247c3d40c5eb4fab11af7993b/LazyIDA2.py
wanted_name: LazyIDA
comment: LazyIDA
created_at: 2022-06-08 03:36:44+00:00
pushed_at: 2022-06-22 03:49:54+00:00
content:
from __future__ import division
from __future__ import print_function
import binascii
from struct import unpack
import idaapi
import idautils
import idc
import base64

from PyQt5 import QtCore
from PyQt5.Qt import QApplication
from PyQt5.QtWidgets import QDialog, QHBoxLayout, QVBoxLayout, QLabel, QRadioButton, QTextEdit, QPushButton, QLineEdit, \
    QMessageBox, QFileDialog, QComboBox

ACTION_CONVERT = ["lazyida:convert%d" % i for i in range(10)]
ACTION_SCANVUL = "lazyida:scanvul"
ACTION_COPYEA = "lazyida:copyea"
ACTION_GOTOCLIP = "lazyida:gotoclip"
ACTION_XORDATA = "lazyida:xordata"
ACTION_FILLNOP = "lazyida:fillnop"
ACTION_PASTE = "lazyida:paste"
ACTION_DUMPER = "lazyida:dumper"
ACTION_JMP = "lazyida:jmper"
ACTION_GOTOEA = "huskygg:gotoea"

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

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

ARCH = 0
BITS = 0

history_jmp_base = []


def dump_bytes(addr, size):
    buffer = bytearray()
    for i in range(size):
        buffer.append(idc.Byte(addr + i))
    return buffer


def toHex(x):
    x = hex(x)
    x = x.upper()
    x = x.replace("0X", "")
    x = x.replace("L", "")
    return x


def hex_cleaner(s):
    s = s.strip()
    s = s.replace("0x", "")
    s = s.replace("h", "")
    s = s.replace("L", "")
    return s


class jmper_windows(QDialog):
    def __init__(self, inital_target_addr=None):
        global history_jmp_base
        super(jmper_windows, self).__init__()

        self.setWindowTitle("Lazy Jumper")
        self.cur_addr = idc.get_screen_ea()
        self.cur_image_base = idaapi.get_imagebase(self.cur_addr)

        layout_main = QVBoxLayout()
        layout_main.addWidget(QLabel("Jump without rebase the idb."))

        # current image base
        layout_cur_base = QHBoxLayout()
        layout_cur_base.addWidget(QLabel("Current Base:"))
        self.edit_cur_base = QLineEdit()
        self.edit_cur_base.setText(toHex(self.cur_image_base))
        self.edit_cur_base.setEnabled(False)
        layout_cur_base.addWidget(self.edit_cur_base)

        # new image base
        layout_new_base = QHBoxLayout()
        self.combobox_new_base = QComboBox()
        layout_new_base.addWidget(QLabel("New Base:"))
        layout_new_base.addWidget(self.combobox_new_base)

        # target address
        layout_target = QHBoxLayout()
        self.edit_target_addr = QLineEdit()
        self.edit_target_addr.returnPressed.connect(self.jmp_clicked)

        if inital_target_addr is not None:
            self.edit_target_addr.setText(inital_target_addr)

        layout_target.addWidget(QLabel("Target Addr:"))
        layout_target.addWidget(self.edit_target_addr)

        # initize combobox with history jmp image base.
        for addr in history_jmp_base:
            self.combobox_new_base.addItem(addr)

        if len(history_jmp_base) > 0:
            self.combobox_new_base.setCurrentIndex(len(history_jmp_base) - 1)
        else:
            self.combobox_new_base.addItem(toHex(self.cur_image_base))
            self.combobox_new_base.setCurrentIndex(0)

        self.combobox_new_base.setEditable(True)
        self.btn_jmp = QPushButton("Jump")

        layout_main.addLayout(layout_cur_base)
        layout_main.addLayout(layout_new_base)
        layout_main.addLayout(layout_target)
        layout_main.addWidget(self.btn_jmp)
        self.btn_jmp.clicked.connect(self.jmp_clicked)

        self.setLayout(layout_main)
        self.show()
        self.exec_()

    def keyPressEvent(self, event):
        key_code = event.key()
        if key_code == QtCore.Qt.Key_Escape:
             self.close()
        elif key_code == QtCore.Qt.Key_Enter:
            self.jmp_clicked()

    def jmp_clicked(self):
        target_base_hex = hex_cleaner(self.combobox_new_base.currentText())
        target = int(hex_cleaner(self.edit_target_addr.text()), 16)
        target_base = int(target_base_hex, 16)
        offset = target - target_base
        real_offset = offset + self.cur_image_base
        if target_base_hex not in history_jmp_base:
            history_jmp_base.append(target_base_hex)
        print("original base: %x new base: %x offset:%x" % (self.cur_image_base, target_base, offset))
        idc.Jump(real_offset)
        self.close()


class dumper_windows(QDialog):
    def __init__(self):
        super(dumper_windows, self).__init__()
        self.addr = idc.get_screen_ea()
        self.setWindowTitle("Lazy dumper.")
        layout_main = QVBoxLayout()
        layout_base = QHBoxLayout()
        layout_base.addWidget(QLabel("Base(HEX):"))
        self.edit_base = QLineEdit()
        layout_base.addWidget(self.edit_base)
        layout_size = QHBoxLayout()
        layout_size.addWidget(QLabel("Size(HEX):"))
        self.edit_size = QLineEdit()
        layout_size.addWidget(self.edit_size)

        self.btn_cancel = QPushButton("Cancel")
        self.btn_dump = QPushButton("Dump")

        layout_main.addLayout(layout_base)
        layout_main.addLayout(layout_size)
        layout_main.addWidget(self.btn_dump)
        self.edit_base.setText(hex(self.addr))

        self.setLayout(layout_main)

        self.btn_dump.clicked.connect(self.click_dump)

        self.show()
        self.exec_()

    def click_cancel(self):
        self.close()

    def click_dump(self):
        addr = self.edit_base.text()
        size = self.edit_size.text()
        try:
            addr = int(hex_cleaner(addr), 16)
            size = int(hex_cleaner(size), 16)
        except ValueError as e:
            QMessageBox.warning(self, " Error ", "Wrong numbers! please check!")
            return

        print("dump from %x size:%x" % (addr, size))
        data = dump_bytes(addr, size)
        fileName, filetype = QFileDialog.getSaveFileName(self,
                                                         "File Saving",
                                                         "",
                                                         "All Files (*)")
        if fileName != u'':
            fp = open(fileName, 'wb')
            fp.write(data)
            fp.close()
            print("saved to : " + fileName)
            self.close()

class paste_data_window(QDialog):
    def __init__(self, target_addr):
        super(paste_data_window, self).__init__()
        self.addr = target_addr
        self.setWindowTitle('Paste data')
        layout_main = QVBoxLayout()
        layout_option = QHBoxLayout()
        layout_option.addWidget(QLabel("Input Type: "))
        self.option_types = [QRadioButton("HEX"), QRadioButton("BASE64"), QRadioButton("ASCII")]
        self.option_types[0].setChecked(True)
        for qcheck in self.option_types:
            layout_option.addWidget(qcheck)
        self.edit = QTextEdit()
        self.btn_apply = QPushButton("Apply")
        layout_main.addWidget(QLabel("Target Addr: %s " % hex(target_addr)[2:].upper()))
        layout_main.addLayout(layout_option)
        layout_main.addWidget(self.edit)
        layout_main.addWidget(self.btn_apply)
        self.btn_apply.clicked.connect(self.event_apply_onclicked)
        self.setLayout(layout_main)
        self.show()
        self.exec_()

    def event_apply_onclicked(self):
        text = self.edit.toPlainText()
        if self.option_types[0].isChecked():
            text = text.strip()
            stopWords = [",", "0x", "{", "}", "H", "h", "[", "]", " ", "\n", ";"]
            for ch in stopWords:
                text = text.replace(ch, "")
            print("HEX:" + text)
            hex_bytes = bytearray(binascii.a2b_hex(text))
            for i in range(len(hex_bytes)):
                idc.PatchByte(self.addr + i, hex_bytes[i])
            self.close()
        elif self.option_types[1].isChecked():
            text = text.strip()
            hex_bytes = bytearray(base64.b64decode(text))
            for i in range(len(hex_bytes)):
                idc.PatchByte(self.addr + i, hex_bytes[i])
            self.close()
        elif self.option_types[2].isChecked():
            hex_bytes = bytearray(text.encode('utf-8'))
            for i in range(len(hex_bytes)):
                idc.PatchByte(self.addr + i, hex_bytes[i])
            self.close()


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


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

def jump_ea():
    ea_point = idc.here()
    ea = idc.Dword(ea_point)
    if idc.jumpto(ea):
        print("jump to 0x%x success" % ea)
    else:
        print("0x%x jump err" % ea)

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


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

    def __init__(self, title, items, icon, embedded=False):
        idaapi.Choose.__init__(self, title, [["Address", 20], ["Function", 30], ["Format", 30]], embedded=embedded)
        self.items = items
        self.icon = 45

    def GetItems(self):
        return self.items

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

    def OnClose(self):
        pass

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

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

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


class hotkey_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for hotkey actions
    """

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

    def activate(self, ctx):
        if self.action == ACTION_COPYEA:
            ea = idc.get_screen_ea()
            if ea != idaapi.BADADDR:
                copy_to_clip("0x%X" % ea)
                print("Address 0x%X has been copied to clipboard" % ea)
        elif self.action == ACTION_GOTOCLIP:
            # loc = parse_location(clip_text())
            # if loc != idaapi.BADADDR:
            #   print("Goto location 0x%x" % loc)
            #   idc.jumpto(loc)
            jmper_windows(hex_cleaner(clip_text()))
        elif self.action == ACTION_GOTOEA:
            jump_ea()
        return 1

    def update(self, ctx):
        if ctx.form_type in (idaapi.BWN_DISASM, idaapi.BWN_DUMP):
            return idaapi.AST_ENABLE_FOR_WIDGET
        else:
            return idaapi.AST_DISABLE_FOR_WIDGET


class menu_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for menu actions
    """

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

    def activate(self, ctx):
        if self.action in ACTION_CONVERT:
            # convert
            t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(), idaapi.get_current_viewer()
            if idaapi.read_selection(view, t0, t1):
                start, end = t0.place(view).toea(), t1.place(view).toea()
                size = end - start+1
            elif idc.get_item_size(idc.get_screen_ea()) > 1:
                start = idc.get_screen_ea()
                size = idc.get_item_size(start)+1
                end = start + size
            else:
                return False

            data = idc.get_bytes(start, size)
            if isinstance(data, str):  # python2 compatibility
                data = bytearray(data)
            name = idc.get_name(start, idc.GN_VISIBLE)
            if not name:
                name = "data"
            if data:
                print("\n[+] Dump 0x%X - 0x%X (%u bytes) :" % (start, end, size))
                if self.action == ACTION_CONVERT[0]:
                    # escaped string
                    print('"%s"' % "".join("\\x%02X" % b for b in data))
                elif self.action == ACTION_CONVERT[1]:
                    # hex string
                    print("".join("%02X" % b for b in data))
                elif self.action == ACTION_CONVERT[2]:
                    # C array
                    output = "unsigned char %s[%d] = {" % (name, size)
                    for i in range(size):
                        if i % 16 == 0:
                            output += "\n    "
                        output += "0x%02X, " % data[i]
                    output = output[:-2] + "\n};"
                    print(output)
                elif self.action == ACTION_CONVERT[3]:
                    # C array word
                    data += b"\x00"
                    array_size = (size + 1) // 2
                    output = "unsigned short %s[%d] = {" % (name, array_size)
                    for i in range(0, size, 2):
                        if i % 16 == 0:
                            output += "\n    "
                        output += "0x%04X, " % u16(data[i:i + 2])
                    output = output[:-2] + "\n};"
                    print(output)
                elif self.action == ACTION_CONVERT[4]:
                    # C array dword
                    data += b"\x00" * 3
                    array_size = (size + 3) // 4
                    output = "unsigned int %s[%d] = {" % (name, array_size)
                    for i in range(0, size, 4):
                        if i % 32 == 0:
                            output += "\n    "
                        output += "0x%08X, " % u32(data[i:i + 4])
                    output = output[:-2] + "\n};"
                    print(output)
                elif self.action == ACTION_CONVERT[5]:
                    # C array qword
                    data += b"\x00" * 7
                    array_size = (size + 7) // 8
                    output = "unsigned long %s[%d] = {" % (name, array_size)
                    for i in range(0, size, 8):
                        if i % 32 == 0:
                            output += "\n    "
                        output += "%#018X, " % u64(data[i:i + 8])
                    output = output[:-2] + "\n};"
                    print(output.replace("0X", "0x"))
                elif self.action == ACTION_CONVERT[6]:
                    # python list
                    print("[%s]" % ", ".join("0x%02X" % b for b in data))
                elif self.action == ACTION_CONVERT[7]:
                    # python list word
                    data += b"\x00"
                    print("[%s]" % ", ".join("0x%04X" % u16(data[i:i + 2]) for i in range(0, size, 2)))
                elif self.action == ACTION_CONVERT[8]:
                    # python list dword
                    data += b"\x00" * 3
                    print("[%s]" % ", ".join("0x%08X" % u32(data[i:i + 4]) for i in range(0, size, 4)))
                elif self.action == ACTION_CONVERT[9]:
                    # python list qword
                    data += b"\x00" * 7
                    print("[%s]" % ", ".join("%#018X" % u64(data[i:i + 8]) for i in range(0, size, 8)).replace("0X",
                                                                                                               "0x"))
        elif self.action == ACTION_XORDATA:
            t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(), idaapi.get_current_viewer()
            if idaapi.read_selection(view, t0, t1):
                start, end = t0.place(view).toea(), t1.place(view).toea()
            else:
                if idc.get_item_size(idc.get_screen_ea()) > 1:
                    start = idc.get_screen_ea()
                    end = start + idc.get_item_size(start)
                else:
                    return False

            data = idc.get_bytes(start, end - start)
            if isinstance(data, str):  # python2 compatibility
                data = bytearray(data)
            x = idaapi.ask_long(0, "Xor with...")
            if x:
                x &= 0xFF
                print("\n[+] Xor 0x%X - 0x%X (%u bytes) with 0x%02X:" % (start, end, end - start, x))
                print(repr("".join(chr(b ^ x) for b in data)))
        elif self.action == ACTION_FILLNOP:
            t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(), idaapi.get_current_viewer()
            if idaapi.read_selection(view, t0, t1):
                start, end = t0.place(view).toea(), t1.place(view).toea()
                idaapi.patch_bytes(start, b"\x90" * (end - start))
                print("\n[+] Fill 0x%X - 0x%X (%u bytes) with NOPs" % (start, end, end - start))
        elif self.action == ACTION_SCANVUL:
            print("\n[+] Finding Format String Vulnerability...")
            found = []
            for addr in idautils.Functions():
                name = idc.get_func_name(addr)
                if "printf" in name and "v" not in name and idc.get_segm_name(addr) in (".text", ".plt", ".idata"):
                    xrefs = idautils.CodeRefsTo(addr, False)
                    for xref in xrefs:
                        vul = self.check_fmt_function(name, xref)
                        if vul:
                            found.append(vul)
            if found:
                print("[!] Done! %d possible vulnerabilities found." % len(found))
                ch = VulnChoose("Vulnerability", found, None, False)
                ch.Show()
            else:
                print("[-] No format string vulnerabilities found.")
        elif self.action == ACTION_PASTE:
            print("paste data.")
            paste_data_window(idc.get_screen_ea())
        elif self.action == ACTION_DUMPER:
            print("dump data.")
            dumper_windows()
        elif self.action == ACTION_JMP:
            print("jmper")
            jmper_windows()
        else:
            return 0

        return 1

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

    @staticmethod
    def check_fmt_function(name, addr):
        """
        Check if the format string argument is not valid
        """
        function_head = idc.get_func_attr(addr, idc.FUNCATTR_START)

        while True:
            addr = idc.prev_head(addr)
            op = idc.print_insn_mnem(addr).lower()
            dst = idc.print_operand(addr, 0)

            if op in ("ret", "retn", "jmp", "b") or addr < function_head:
                return

            c = idc.get_cmt(addr, 0)
            if c and c.lower() == "format":
                break
            elif name.endswith(("snprintf_chk",)):
                if op in ("mov", "lea") and dst.endswith(("r8", "r8d", "[esp+10h]")):
                    break
            elif name.endswith(("sprintf_chk",)):
                if op in ("mov", "lea") and (dst.endswith(("rcx", "[esp+0Ch]", "R3")) or
                                             dst.endswith("ecx") and BITS == 64):
                    break
            elif name.endswith(("snprintf", "fnprintf")):
                if op in ("mov", "lea") and (dst.endswith(("rdx", "[esp+8]", "R2")) or
                                             dst.endswith("edx") and BITS == 64):
                    break
            elif name.endswith(("sprintf", "fprintf", "dprintf", "printf_chk")):
                if op in ("mov", "lea") and (dst.endswith(("rsi", "[esp+4]", "R1")) or
                                             dst.endswith("esi") and BITS == 64):
                    break
            elif name.endswith("printf"):
                if op in ("mov", "lea") and (dst.endswith(("rdi", "[esp]", "R0")) or
                                             dst.endswith("edi") and BITS == 64):
                    break

        # format arg found, check its type and value
        # get last oprend
        op_index = idc.generate_disasm_line(addr, 0).count(",")
        op_type = idc.get_operand_type(addr, op_index)
        opnd = idc.print_operand(addr, op_index)

        if op_type == idc.o_reg:
            # format is in register, try to track back and get the source
            _addr = addr
            while True:
                _addr = idc.prev_head(_addr)
                _op = idc.print_insn_mnem(_addr).lower()
                if _op in ("ret", "retn", "jmp", "b") or _addr < function_head:
                    break
                elif _op in ("mov", "lea", "ldr") and idc.print_operand(_addr, 0) == opnd:
                    op_type = idc.get_operand_type(_addr, 1)
                    opnd = idc.print_operand(_addr, 1)
                    addr = _addr
                    break

        if op_type == idc.o_imm or op_type == idc.o_mem:
            # format is a memory address, check if it's in writable segment
            op_addr = idc.get_operand_value(addr, op_index)
            seg = idaapi.getseg(op_addr)
            if seg:
                if not seg.perm & idaapi.SEGPERM_WRITE:
                    # format is in read-only segment
                    return

        print("0x%X: Possible Vulnerability: %s, format = %s" % (addr, name, opnd))
        return ["0x%X" % addr, name, opnd]


class hexrays_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for hexrays actions
    """

    def __init__(self, action):
        idaapi.action_handler_t.__init__(self)
        self.action = action
        self.ret_type = {}

    def activate(self, ctx):
        if self.action == ACTION_HX_REMOVERETTYPE:
            vdui = idaapi.get_widget_vdui(ctx.widget)
            self.remove_rettype(vdui)
            vdui.refresh_ctext()
        elif self.action == ACTION_HX_COPYEA:
            ea = idaapi.get_screen_ea()
            if ea != idaapi.BADADDR:
                copy_to_clip("0x%X" % ea)
                print("Address 0x%X has been copied to clipboard" % ea)
        elif self.action == ACTION_HX_COPYNAME:
            name = idaapi.get_highlight(idaapi.get_current_viewer())[0]
            if name:
                copy_to_clip(name)
                print("%s has been copied to clipboard" % name)
        elif self.action == ACTION_HX_GOTOCLIP:
            loc = parse_location(clip_text())
            print("Goto location 0x%x" % loc)
            idc.jumpto(loc)
        else:
            return 0

        return 1

    def update(self, ctx):
        vdui = idaapi.get_widget_vdui(ctx.widget)
        return idaapi.AST_ENABLE_FOR_WIDGET if vdui else idaapi.AST_DISABLE_FOR_WIDGET

    def remove_rettype(self, vu):
        if vu.item.citype == idaapi.VDI_FUNC:
            # current function
            ea = vu.cfunc.entry_ea
            old_func_type = idaapi.tinfo_t()
            if not vu.cfunc.get_func_type(old_func_type):
                return False
        elif vu.item.citype == idaapi.VDI_EXPR and vu.item.e.is_expr() and vu.item.e.type.is_funcptr():
            # call xxx
            ea = vu.item.get_ea()
            old_func_type = idaapi.tinfo_t()

            func = idaapi.get_func(ea)
            if func:
                try:
                    cfunc = idaapi.decompile(func)
                except idaapi.DecompilationFailure:
                    return False

                if not cfunc.get_func_type(old_func_type):
                    return False
            else:
                return False
        else:
            return False

        fi = idaapi.func_type_data_t()
        if ea != idaapi.BADADDR and old_func_type.get_func_details(fi):
            # Return type is already void
            if fi.rettype.is_decl_void():
                # Restore ret type
                if ea not in self.ret_type:
                    return True
                ret = self.ret_type[ea]
            else:
                # Save ret type and change it to void
                self.ret_type[ea] = fi.rettype
                ret = idaapi.BT_VOID

            # Create new function info with new rettype
            fi.rettype = idaapi.tinfo_t(ret)

            # Create new function type with function info
            new_func_type = idaapi.tinfo_t()
            new_func_type.create_func(fi)

            # Apply new function type
            if idaapi.apply_tinfo(ea, new_func_type, idaapi.TINFO_DEFINITE):
                return vu.refresh_view(True)

        return False


class UI_Hook(idaapi.UI_Hooks):
    def __init__(self):
        idaapi.UI_Hooks.__init__(self)

    def finish_populating_widget_popup(self, form, popup):
        form_type = idaapi.get_widget_type(form)
        if form_type == idaapi.BWN_DISASM or form_type == idaapi.BWN_DUMP:
            idaapi.attach_action_to_popup(form, popup, ACTION_PASTE, None)
            idaapi.attach_action_to_popup(form, popup, ACTION_DUMPER, None)
            idaapi.attach_action_to_popup(form, popup, ACTION_JMP, None)
            t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(), idaapi.get_current_viewer()
            if idaapi.read_selection(view, t0, t1) or idc.get_item_size(idc.get_screen_ea()) > 1:
                idaapi.attach_action_to_popup(form, popup, ACTION_XORDATA, None)
                idaapi.attach_action_to_popup(form, popup, ACTION_FILLNOP, None)
                for action in ACTION_CONVERT:
                    idaapi.attach_action_to_popup(form, popup, action, "Convert/")

        if form_type == idaapi.BWN_DISASM and (ARCH, BITS) in [(idaapi.PLFM_386, 32),
                                                               (idaapi.PLFM_386, 64),
                                                               (idaapi.PLFM_ARM, 32), ]:
            idaapi.attach_action_to_popup(form, popup, ACTION_SCANVUL, None)


class HexRays_Hook(object):
    def callback(self, event, *args):
        if event == idaapi.hxe_populating_popup:
            form, phandle, vu = args
            if vu.item.citype == idaapi.VDI_FUNC or (
                    vu.item.citype == idaapi.VDI_EXPR and vu.item.e.is_expr() and vu.item.e.type.is_funcptr()):
                idaapi.attach_action_to_popup(form, phandle, ACTION_HX_REMOVERETTYPE, None)
        elif event == idaapi.hxe_double_click:
            vu, shift_state = args
            # auto jump to target if clicked item is xxx->func();
            if vu.item.citype == idaapi.VDI_EXPR and vu.item.e.is_expr():
                expr = idaapi.tag_remove(vu.item.e.print1(None))
                if "->" in expr:
                    # find target function
                    name = expr.split("->")[-1]
                    addr = idc.get_name_ea_simple(name)
                    if addr == idaapi.BADADDR:
                        # try class::function
                        e = vu.item.e
                        while e.x:
                            e = e.x
                        addr = idc.get_name_ea_simple("%s::%s" % (str(e.type).split()[0], name))

                    if addr != idaapi.BADADDR:
                        idc.jumpto(addr)
                        return 1
        return 0


class LazyIDA_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = "LazyIDA"
    help = ""
    wanted_name = "LazyIDA"
    wanted_hotkey = ""

    def init(self):
        self.hexrays_inited = False
        self.registered_actions = []
        self.registered_hx_actions = []

        global ARCH
        global BITS
        ARCH = idaapi.ph_get_id()
        info = idaapi.get_inf_structure()
        if info.is_64bit():
            BITS = 64
        elif info.is_32bit():
            BITS = 32
        else:
            BITS = 16

        print("LazyIDA (v1.0.0.3) plugin has been loaded.")

        # Register menu actions
        menu_actions = (
            idaapi.action_desc_t(ACTION_CONVERT[0], "Convert to string", menu_action_handler_t(ACTION_CONVERT[0]), None,
                                 None, 80),
            idaapi.action_desc_t(ACTION_CONVERT[1], "Convert to hex string", menu_action_handler_t(ACTION_CONVERT[1]),
                                 None, None, 8),
            idaapi.action_desc_t(ACTION_CONVERT[2], "Convert to C/C++ array (BYTE)",
                                 menu_action_handler_t(ACTION_CONVERT[2]), None, None, 38),
            idaapi.action_desc_t(ACTION_CONVERT[3], "Convert to C/C++ array (WORD)",
                                 menu_action_handler_t(ACTION_CONVERT[3]), None, None, 38),
            idaapi.action_desc_t(ACTION_CONVERT[4], "Convert to C/C++ array (DWORD)",
                                 menu_action_handler_t(ACTION_CONVERT[4]), None, None, 38),
            idaapi.action_desc_t(ACTION_CONVERT[5], "Convert to C/C++ array (QWORD)",
                                 menu_action_handler_t(ACTION_CONVERT[5]), None, None, 38),
            idaapi.action_desc_t(ACTION_CONVERT[6], "Convert to python list (BYTE)",
                                 menu_action_handler_t(ACTION_CONVERT[6]), None, None, 201),
            idaapi.action_desc_t(ACTION_CONVERT[7], "Convert to python list (WORD)",
                                 menu_action_handler_t(ACTION_CONVERT[7]), None, None, 201),
            idaapi.action_desc_t(ACTION_CONVERT[8], "Convert to python list (DWORD)",
                                 menu_action_handler_t(ACTION_CONVERT[8]), None, None, 201),
            idaapi.action_desc_t(ACTION_CONVERT[9], "Convert to python list (QWORD)",
                                 menu_action_handler_t(ACTION_CONVERT[9]), None, None, 201),
            idaapi.action_desc_t(ACTION_XORDATA, "Get xored data", menu_action_handler_t(ACTION_XORDATA), None, None,
                                 9),
            idaapi.action_desc_t(ACTION_FILLNOP, "Fill with NOPs", menu_action_handler_t(ACTION_FILLNOP), None, None,
                                 9),
            idaapi.action_desc_t(ACTION_PASTE, "Paste Data", menu_action_handler_t(ACTION_PASTE), None, None, 9),
            idaapi.action_desc_t(ACTION_DUMPER, "Lazy Dumper", menu_action_handler_t(ACTION_DUMPER), None, None, 9),
            idaapi.action_desc_t(ACTION_JMP, "Lazy Jumper [Shift + G]", menu_action_handler_t(ACTION_JMP), None, None,
                                 9),
            idaapi.action_desc_t(ACTION_SCANVUL, "Scan format string vulnerabilities",
                                 menu_action_handler_t(ACTION_SCANVUL), None, None, 160),
        )
        for action in menu_actions:
            idaapi.register_action(action)
            self.registered_actions.append(action.name)

        # Register hotkey actions
        hotkey_actions = (
            idaapi.action_desc_t(ACTION_COPYEA, "Copy EA", hotkey_action_handler_t(ACTION_COPYEA), "w",
                                 "Copy current EA", 0),
            idaapi.action_desc_t(ACTION_GOTOCLIP, "Goto clip EA", hotkey_action_handler_t(ACTION_GOTOCLIP), "Shift-G",
                                 "Goto clipboard EA", 0),
            idaapi.action_desc_t(ACTION_GOTOEA, "Goto EA", hotkey_action_handler_t(ACTION_GOTOEA), "Shift-Q",
                                 "Goto EA", 0),

        )
        for action in hotkey_actions:
            idaapi.register_action(action)
            self.registered_actions.append(action.name)

        # Add ui hook
        self.ui_hook = UI_Hook()
        self.ui_hook.hook()

        # Add hexrays ui callback
        if idaapi.init_hexrays_plugin():
            addon = idaapi.addon_info_t()
            addon.id = "tw.l4ys.lazyida"
            addon.name = "LazyIDA"
            addon.producer = "Lays"
            addon.url = "https://github.com/L4ys/LazyIDA"
            addon.version = "1.0.0.3"
            idaapi.register_addon(addon)

            hx_actions = (
                idaapi.action_desc_t(ACTION_HX_REMOVERETTYPE, "Remove return type",
                                     hexrays_action_handler_t(ACTION_HX_REMOVERETTYPE), "v"),
                idaapi.action_desc_t(ACTION_HX_COPYEA, "Copy ea", hexrays_action_handler_t(ACTION_HX_COPYEA), "w"),
                idaapi.action_desc_t(ACTION_HX_COPYNAME, "Copy name", hexrays_action_handler_t(ACTION_HX_COPYNAME),
                                     "c"),
                idaapi.action_desc_t(ACTION_HX_GOTOCLIP, "Goto clipboard ea",
                                     hexrays_action_handler_t(ACTION_HX_GOTOCLIP), "Shift-G"),
            )
            for action in hx_actions:
                idaapi.register_action(action)
                self.registered_hx_actions.append(action.name)

            self.hx_hook = HexRays_Hook()
            idaapi.install_hexrays_callback(self.hx_hook.callback)
            self.hexrays_inited = True

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        if hasattr(self, "ui_hook"):
            self.ui_hook.unhook()

        # Unregister actions
        for action in self.registered_actions:
            idaapi.unregister_action(action)

        if self.hexrays_inited:
            # Unregister hexrays actions
            for action in self.registered_hx_actions:
                idaapi.unregister_action(action)
            if self.hx_hook:
                idaapi.remove_hexrays_callback(self.hx_hook.callback)
            idaapi.term_hexrays_plugin()


def PLUGIN_ENTRY():
    print("init huskygg-IDA")
    return LazyIDA_t()


repository: CKCat/IDAPlugInContest
file: 2015/PlugIn/ifl.py
description: IDAPlugIn Contest 备份
comment: Interactive Functions List
created_at: 2021-09-23 15:11:52+00:00
pushed_at: 2024-01-29 08:05:30+00:00
url:
https://github.com/CKCat/IDAPlugInContest/blob/7ee5e83dee340fdbce81e86d310287f0a34d7250/2015/PlugIn/ifl.py

content:
#!/usr/bin/env python
#
# IFL - Interactive Functions List
#
# how to install: copy the script into plugins directory, i.e: C:\Program Files\IDA 6.8\plugins
# then:
# run from IDA menu: View -> PLUGIN_NAME
# or press: PLUGIN_HOTKEY
#
"""
(c) hasherezade, 2015 run via IDA Pro 6.8
"""
__VERSION__ = '1.0'
__AUTHOR__ = 'hasherezade'

PLUGIN_NAME = "IFL - Interactive Functions List"
PLUGIN_HOTKEY = "Alt-F"

import idautils
from idaapi import *
from idc import *

from idaapi import PluginForm
from PySide import QtGui, QtCore
from PySide.QtCore import QObject, Signal, Slot

# --------------------------------------------------------------------------
# custom functions:
# --------------------------------------------------------------------------

def rva_to_va(rva):
    base = idaapi.get_imagebase()
    return rva + base
    
def va_to_rva(va):
    base = idaapi.get_imagebase()
    return va - base

def function_at(ea):
    start = ea
    functions = Functions(start)
    for func in Functions():
        return func
    return None
    
def parse_function_args(ea):
  local_variables = [ ]
  arguments = [ ]
  current = local_variables

  frame = idc.GetFrame(ea)
  arg_string = ""
  if frame == None:
    return ""
        
  start = idc.GetFirstMember(frame)
  end = idc.GetLastMember(frame)
  count = 0
  max_count = 10000
  args_str = ""     
  while start <= end and count <= max_count:
    size = idc.GetMemberSize(frame, start)
    count = count + 1
    if size == None:
      start = start + 1
      continue

    name = idc.GetMemberName(frame, start)  
    start += size
            
    if name in [" r", " s"]:
      # Skip return address and base pointer
      current = arguments
      continue
    arg_string += " " + name
    current.append(name)
  args_str = ", ".join(arguments)
  if len(args_str) == 0:
    args_str = "void"
  return "(" + args_str + ")"

def parse_function_type(ea, end=None):
  frame = idc.GetFrame(ea)
  if frame == None:
    return ""
  if end == None: #try to find end
      func = function_at(ea)
      if not func :
        return "?"
      end = PrevAddr(GetFunctionAttr(func, FUNCATTR_END)) 
  end_addr = end
  mnem = GetDisasm(end_addr)
  
  if not "ret" in mnem:
    #it's not a real end, get instruction before...
    end_addr = PrevAddr(end)
    if end_addr == BADADDR:
      #cannot get the real end
      return ""
    mnem = GetDisasm(end_addr)

  if not "ret" in mnem:
    #cannot get the real end
    return ""

  op = GetOpType(end_addr, 0)
  if op == o_void:
    #retn has NO parameters
    return "__cdecl"
  #retn has parameters
  return "__stdcall"

def _getFunctionType(start, end=None):
    type = GetType(start)
    if type == None:
        return parse_function_type(start, end)
    args_start = type.find('(')
    if not args_start == None:
        type = type[:args_start]
    return type
    
def _isFunctionMangled(ea):
    name = GetFunctionName(ea)
    disable_mask = GetLongPrm(INF_SHORT_DN)
    if Demangle(name, disable_mask) == None:
        return False
    return True
    
def _getFunctionNameAt(ea):
    name = GetFunctionName(ea)
    disable_mask = GetLongPrm(INF_SHORT_DN)
    demangled_name = Demangle(name, disable_mask)
    if demangled_name == None:
        return name
    args_start = demangled_name.find('(')
    if args_start == None:
        return demangled_name
    return demangled_name[:args_start]

def _getArgsDescription(ea):
    name = Demangle(GetFunctionName(ea), GetLongPrm(INF_SHORT_DN)) #get from mangled name
    if not name:
        name = GetType(ea) #get from type
        if not name:
            return parse_function_args(ea) #cannot get params from the mangled name
    args_start = name.find('(')
    if args_start != None and args_start != (-1):
        return name[args_start:]
    return ""
    
def _getArgsNum(ea):
    args = _getArgsDescription(ea)
    if not args:
        return 0
    delimiter = ','
    args_list = args.split(delimiter)
    args_num = 0
    for arg in args_list:
        if arg == "()" or arg == "(void)":
            continue
        args_num += 1
    return args_num
    
# --------------------------------------------------------------------------
# custom data types:
# --------------------------------------------------------------------------

#Global DataManager

class DataManager(QObject):
    updateSignal = Signal()
    
    def __init__(self, parent=None):
        QtCore.QObject.__init__(self, parent=parent)
        self.currentRva = long(BADADDR)
        
    def setFunctionName(self, start, func_name):
        flags = idaapi.SN_NOWARN | idaapi.SN_NOCHECK
        if idc.MakeNameEx(start, func_name, flags):
            self.updateSignal.emit()
            return True
        return False
 
    def setCurrentRva(self, rva):
        if self.currentRva == rva:
            return # nothing changed
            
        if rva is None:
            rva = long(BADADDR) 
        self.currentRva = long(rva)
        self.updateSignal.emit()
 
# --------------------------------------------------------------------------

class FunctionInfo_t():
    def __init__(self, start, end, refs_list, called_list, is_import=False):
        self.start = start
        self.end = end
        self.args_num = _getArgsNum(start)
        self.type = _getFunctionType(start, end)
        self.is_import = is_import
        self.refs_list = refs_list
        self.called_list = called_list
    
    def contains(self, addr):
        bng = self.start
        end = self.end
        #swap if order is opposite:
        if self.start > self.end:
            end = self.start
            start = self.end
        if addr >= bgn and  addr < end:
            return True
        return False
# --------------------------------------------------------------------------
# custom models:
# --------------------------------------------------------------------------
class TableModel_t(QtCore.QAbstractTableModel):
    """Model for the table """
    COL_START = 0
    COL_END = 1
    COL_NAME = 2
    COL_TYPE = 3
    COL_ARGS = 4
    COL_REFS = 5
    COL_CALLED = 6
    COL_IMPORT = 7
    COL_COUNT = 8
    header_names = ['Start', 'End', 'Name', 'Type', 'Args', 'Is refered by', 'Refers to', 'Imported?']
    
#private:

    def _displayHeader(self, orientation, col):
      if orientation == QtCore.Qt.Vertical:
        return None
      if col == self.COL_START:
        return self.header_names[self.COL_START]
      if col == self.COL_END:
        return self.header_names[self.COL_END]
      if col == self.COL_TYPE:
        return self.header_names[self.COL_TYPE]
      if col == self.COL_ARGS:
        return self.header_names[self.COL_ARGS]
      if col == self.COL_NAME:
        return self.header_names[self.COL_NAME]
      if col == self.COL_REFS:
        return self.header_names[self.COL_REFS]
      if col == self.COL_CALLED:
        return self.header_names[self.COL_CALLED]
      if col == self.COL_IMPORT:
        return self.header_names[self.COL_IMPORT]
      return None
      
    def _displayData(self, row, col):
      func_info = self.function_info_list[row]
      if col == self.COL_START:
        return "%08x" % func_info.start
      if col == self.COL_END:
        return "%08x" % func_info.end
      if col == self.COL_TYPE:
        return func_info.type
      if col == self.COL_ARGS:
        return _getArgsDescription(func_info.start)
      if col == self.COL_NAME:
        return _getFunctionNameAt(func_info.start)
      if col == self.COL_REFS:
        return len(func_info.refs_list)
      if col == self.COL_CALLED:
        return len(func_info.called_list)
      if col == self.COL_IMPORT:
        if func_info.is_import:
            return "+"
        return "-"
      return None
      
    def _displayToolTip(self, row, col):
      func_info = self.function_info_list[row]
      if col == self.COL_START or col == self.COL_END:
        return "Double Click to follow"
      if col == self.COL_NAME:
        return "Double Click to edit"
      if col == self.COL_REFS:
        return self._listRefs(func_info.refs_list)
      if col == self.COL_CALLED:
        return self._listRefs(func_info.called_list)
      return ""
      
    def _displayBackground(self, row, col):
      func_info = self.function_info_list[row]
      if col == self.COL_START or col == self.COL_END:
        return QtGui.QColor("lightblue")
        
      if col == self.COL_NAME:
        if func_info.is_import :
            return QtGui.QColor("orange")
        return QtGui.QColor("khaki")
      return None
     
    def _listRefs(self, refs_list):
      str_list = []
      for ea, ea_to in refs_list:
        str = "%08x @ %s" % (ea, _getFunctionNameAt(ea_to))
        str_list.append(str)
      return '\n'.join(str_list)  
      
#public:
    def __init__(self, function_info_list, parent=None, *args):
        super(TableModel_t, self).__init__()
        self.function_info_list = function_info_list
        
    def isFollowable(self, col):
        if col == self.COL_START:
            return True
        if col == self.COL_END:
            return True
        return False
        
#Qt4 API
    def rowCount(self, parent):
        return len(self.function_info_list)

    def columnCount(self, parent):
        return self.COL_COUNT
    
    def setData(self, index, content, role):
      if not index.isValid():
        return False
      func_info = self.function_info_list[index.row()]
      if index.column() == self.COL_NAME:
        MakeNameEx(func_info.start, str(content), SN_NOWARN) 
      return True
      
    def data(self, index, role):
        if not index.isValid():
            return None
        col = index.column()
        row = index.row()
        if len(self.function_info_list) <= row:
          return None
          
        func_info = self.function_info_list[row]
        
        if role == QtCore.Qt.UserRole:
          if col == self.COL_START:
            return func_info.start
          elif col == self.COL_END:
            return func_info.end
          else:
            return func_info.start
        elif role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
          return self._displayData(row, col)
        elif role == QtCore.Qt.ToolTipRole:
          return self._displayToolTip(row, col)
        elif role == QtCore.Qt.BackgroundColorRole:
          return self._displayBackground(row, col)
        else:
            return None

    def flags(self, index):
        if not index.isValid():
            return None
        flags = QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
        if index.column() == self.COL_NAME:
          return flags | QtCore.Qt.ItemIsEditable
        return flags
        
    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
        if role == QtCore.Qt.DisplayRole:
            return self._displayHeader(orientation, section)
        else:
            return None
# --------------------------------------------------------------------------
class RefsTableModel_t(QtCore.QAbstractTableModel):
    """Model for the table """
    COL_NAME = 0
    COL_ADDR = 1
    COL_TOADDR = 2
    COL_COUNT = 3
    
#private:
    def _displayHeader(self, orientation, col):
        if orientation == QtCore.Qt.Vertical:
            return None
        if col == self.COL_ADDR:
            return "From Address"
        if col == self.COL_TOADDR:
            return "To Address"
        if col == self.COL_NAME:
            return "Foreign Val."
        return None
    
    def _getTargetAddr(self, row):
        curr_ref_fromaddr = self.refs_list[row][0] #fromaddr
        curr_ref_addr = self.refs_list[row][1] #toaddr
        target_addr = BADADDR
        if self.is_refs_to :
            target_addr = curr_ref_fromaddr
        else:
            target_addr = curr_ref_addr
        return target_addr
            
    def _getForeignFuncName(self, row):
        curr_ref_fromaddr = self.refs_list[row][0] #fromaddr
        curr_ref_addr = self.refs_list[row][1] #toaddr
        
        target_addr = self._getTargetAddr(row)
        if GetMnem(target_addr) != "":
            func_name = _getFunctionNameAt(target_addr)
            if func_name:
                return func_name

        addr_str = "[%08lx]" % target_addr
        target_name = GetDisasm(target_addr)
        return addr_str+ " : " + GetDisasm(target_addr)
        
    def _displayData(self, row, col):
    
      if len(self.refs_list) <= row:
          return None    
      curr_ref_fromaddr = self.refs_list[row][0] #fromaddr
      curr_ref_addr = self.refs_list[row][1] #toaddr
      if col == self.COL_ADDR:
        return "%08x" % curr_ref_fromaddr
      if col == self.COL_TOADDR:
        return "%08x" % curr_ref_addr
      if col == self.COL_NAME:
        return self._getForeignFuncName(row)
      return None
    
    def _getAddrToFollow(self, row, col):
        if col == self.COL_ADDR:
            return self.refs_list[row][0]
        if col == self.COL_TOADDR:
            return self.refs_list[row][1]
        return BADADDR
        
    def _displayBackground(self, row, col):
        if self.isFollowable(col):
            return QtGui.QColor("lightblue")
        return None
      
#public:
    def __init__(self, function_info_list, is_refs_to=True, parent=None, *args):
        super(RefsTableModel_t, self).__init__()
        self.function_info_list = function_info_list
        self.curr_index = (-1)
        self.refs_list = []
        self.is_refs_to = is_refs_to
        
    def isFollowable(self, col):
        if col == self.COL_ADDR:
            return True
        if col == self.COL_TOADDR:
            return True
        return False
        
    def findOffsetIndex(self, data):
        index = 0
        for func_info in self.function_info_list:
            if data >= func_info.start and data <= func_info.end:
                return index
            index += 1
        return (-1)

    def setCurrentIndex(self, curr_index):
        self.curr_index = curr_index
        if self.curr_index == (-1) or self.curr_index >= len(self.function_info_list):
          #reset list
          self.refs_list = []
        else:
            if self.is_refs_to :
                self.refs_list = self.function_info_list[self.curr_index].refs_list
            else:
                self.refs_list = self.function_info_list[self.curr_index].called_list
        self.reset()
    
#Qt4 API
    def rowCount(self, parent=None):
        return len(self.refs_list)

    def columnCount(self, parent):
        return self.COL_COUNT
    
    def data(self, index, role):
        if not index.isValid():
            return None
        col = index.column()
        row = index.row()
        if len(self.refs_list) <= row:
          return None
          
        curr_ref_addr = self.refs_list[row][0]
        
        if role == QtCore.Qt.UserRole:
          if self.isFollowable(col):
            return self._getAddrToFollow(row, col)
        elif role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
          return self._displayData(row, col)
        elif role == QtCore.Qt.BackgroundColorRole:
          return self._displayBackground(row, col)
        else:
            return None

    def flags(self, index):
        if not index.isValid():
            return None
        flags = QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
        return flags
        

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
        if role == QtCore.Qt.DisplayRole:
            return self._displayHeader(orientation, section)
        else:
            return None

# --------------------------------------------------------------------------
# custom views:
  
COLOR_NORMAL = 0xFFFFFF

class FunctionsView_t(QtGui.QTableView):

    # private    
    def _set_segment_color(self, ea, color):
        seg = idaapi.getseg(ea)
        seg.color = COLOR_NORMAL
        seg.update()
   
    # public
    def __init__(self, dataManager, color_hilight, func_model, parent=None):
        super(FunctionsView_t, self).__init__(parent=parent)
        self.setSelectionMode(QtGui.QAbstractItemView.SingleSelection)
        #
        self.prev_addr = BADADDR
        self.color_hilight = color_hilight
        self.func_model = func_model
        self.dataManager = dataManager
        #
        self.setMouseTracking(True)
        self.setAutoFillBackground(True)
        
    #Qt API
    def currentChanged(self, current, previous):
        index_data = self.get_index_data(current)
        self.dataManager.setCurrentRva(index_data)
        
    def hilight_addr(self, addr):
        if self.prev_addr != BADADDR:
            ea = self.prev_addr
            self._set_segment_color(ea, COLOR_NORMAL)
            SetColor(ea, CIC_ITEM, COLOR_NORMAL)
        if addr != BADADDR:
            ea = addr
            self._set_segment_color(ea, COLOR_NORMAL)
            SetColor(addr, CIC_ITEM, self.color_hilight)
        self.prev_addr = addr
        
    def get_index_data(self, index):
        if not index.isValid():
            return None
            
        index_data = index.data(QtCore.Qt.UserRole)

        if not type(index_data) is long:
            return None
        return index_data
        
    def mousePressEvent(self, event):
        event.accept()
        index = self.indexAt(event.pos())
        data = self.get_index_data(index)
        super(QtGui.QTableView, self).mousePressEvent(event)
      
    def mouseDoubleClickEvent(self, event):
        event.accept()
        index = self.indexAt(event.pos())
        if not index.isValid():
            return
        data = self.get_index_data(index)  
        if not data:
            super(QtGui.QTableView, self).mouseDoubleClickEvent(event)
            return
        col = index.column()
        if self.func_model.isFollowable(col):
            self.hilight_addr(data)
            Jump(data)
        super(QtGui.QTableView, self).mouseDoubleClickEvent(event)
        
    def mouseMoveEvent(self, event):
        index = self.indexAt(event.pos())
        if not index.isValid():
            return
        col = index.column()
        if self.func_model.isFollowable(col):
            self.setCursor(QtCore.Qt.PointingHandCursor)
        else:
            self.setCursor(QtCore.Qt.ArrowCursor)
        
    def leaveEvent(self, event):
        self.setCursor(QtCore.Qt.ArrowCursor)
        
    def OnDestroy(self):
        self.hilight_addr(BADADDR)
        
# --------------------------------------------------------------------------

class FunctionsListForm_t(PluginForm):
#private
    _COLOR_HILIGHT_FUNC = 0xFFDDBB # BBGGRR
    _COLOR_HILIGHT_REFTO = 0xBBFFBB
    _COLOR_HILIGHT_REFFROM = 0xDDBBFF
    
    def _getCallingOffset(self, func, called_list):
        start = GetFunctionAttr(func, FUNCATTR_START)
        end = PrevAddr(GetFunctionAttr(func, FUNCATTR_END))
        func_name = _getFunctionNameAt(start)
        curr = start
        calling_list = []
        while (True):
            if curr >= end:
                break
            op = GetOperandValue(curr, 0)
            if op in called_list:
                calling_list.append((curr, op))
            curr = NextAddr(curr)
        return calling_list
    
    def _listFunctionsAddr(self):
        fn_list = list()
        for func in Functions():
            start = GetFunctionAttr(func, FUNCATTR_START)
            fn_list.append(start)
        return fn_list
        
    def _saveFunctionsNames(self, file_name):
        if file_name is None or len(file_name) == 0:
            return False
        fn_list = list()
        for func in Functions():
            start = GetFunctionAttr(func, FUNCATTR_START)
            func_name = _getFunctionNameAt(start)
            line = "%lx : %s" %(start, func_name)
            fn_list.append(line)   
        idaapi.msg(str(file_name))
        with open(file_name, 'w') as f:
            for item in fn_list:
                f.write("%s\n" % item)
            return True     
        return False

    def _loadFunctionsNames(self,file_name):
        if file_name is None or len(file_name) == 0:
            return False
        curr_functions = self._listFunctionsAddr()
        delim = ":"
        loaded = 0
        with open(file_name, 'r') as f:
            for line in f.readlines():
                line = line.strip()
                fn = line.split(delim)
                if len(fn) != 2:
                    continue
                start = int(fn[0].strip(), 16)
                func_name = fn[1].strip()
                if start in curr_functions:
                    if self.subDataManager.setFunctionName(start, func_name) == True:
                        loaded += 1
        return loaded

    def imports_names_callback(self, ea, name, ord):
        self.importsSet.add(ea)
        self.importNamesSet.add(name)
        # True -> Continue enumeration
        return True
        
    def _loadImports(self):
        self.importsSet = set()
        self.importNamesSet = set()
        nimps = idaapi.get_import_module_qty()
        for i in xrange(0, nimps):
            idaapi.enum_import_names(i, self.imports_names_callback)

    def _isImportName(self, name):
        if name in self.importNamesSet:
            return True
        return False
    
    def _isImportStart(self, start):
        if start in self.importsSet:
            return True
        if GetMnem(start) == 'call':
            return False
        #print GetMnem(start)
        op = GetOperandValue(start, 0)
        if op in self.importsSet:
            return True
        return False
    
    def _listRefsTo(self, start):
        func_refs_to = XrefsTo(start, 1)
        refs_list = []
        for ref in func_refs_to:
          if idc.GetMnem(ref.frm) == "":
            continue
          refs_list.append((ref.frm, start))
        return refs_list
    
    def _listRefsFrom(self, func, start, end):
        dif = end - start
        called_list = []
        func_name = _getFunctionNameAt(start)
        
        for indx in xrange(0, dif):
          addr = start + indx
          func_refs_from = XrefsFrom(addr, 1)
          for ref in func_refs_from:
            if _getFunctionNameAt(ref.to) == func_name:
              #skip jumps inside self
              continue 
            called_list.append(ref.to)
        calling_list = self._getCallingOffset(func, called_list)
        return calling_list
    
    def _loadLocals(self):
      #loadImports first : 
      self._loadImports()
      
      for func in Functions():
        start = GetFunctionAttr(func, FUNCATTR_START)
        end = PrevAddr(GetFunctionAttr(func, FUNCATTR_END))
        
        is_import = self._isImportStart(start)
        
        refs_list = self._listRefsTo(start)
        calling_list = self._listRefsFrom(func, start, end)

        func_info = FunctionInfo_t(start, end, refs_list, calling_list, is_import)
        self.functionsMap[va_to_rva(start)] = func_info
        self.functionsMap[va_to_rva(end)] = func_info
        self.addr_list.append(func_info)
    
    def _setup_sorted_model(self, view, model):
        sorted_model = QtGui.QSortFilterProxyModel()    
        sorted_model.setDynamicSortFilter(True)
        sorted_model.setSourceModel(model)
        view.setModel(sorted_model)
        view.setSortingEnabled(True)
        #
        sorted_model.setParent(view)        
        model.setParent(sorted_model)
        return sorted_model
        
    def _update_current_offset(self, view, refs_model, offset):
        if offset:
            index = refs_model.findOffsetIndex(offset)
        else:
            index = (-1)
        refs_model.setCurrentIndex(index)
        refs_model.reset()
        view.reset()
        view.repaint()
    
    def _update_function_name(self, ea):
        try:
            func_info = self.functionsMap[va_to_rva(ea)]
        except KeyError:
            return
            
        func_type = func_info.type
        func_args = _getArgsDescription(ea)
        func_name = _getFunctionNameAt(ea)
        func_name = _getFunctionNameAt(ea)
        self.refs_label.setText(func_type + " <b>"+func_name+"</b> " + func_args)

    def _update_ref_tabs(self, ea):
        tocount = 0
        fromcount = 0
        try:
            func_info = self.functionsMap[va_to_rva(ea)]
            tocount = len(func_info.refs_list)
            fromcount = len(func_info.called_list)
        except KeyError:
            pass
        self.refs_tabs.setTabText(0,  "Is refered by %d:" % tocount)
        self.refs_tabs.setTabText(1,  "Refers to %d:" % fromcount)
        
    def adjustColumnsToContents(self):
        self.addr_view.resizeColumnToContents(0)
        self.addr_view.resizeColumnToContents(1)     
        self.addr_view.resizeColumnToContents(2)
        #
        self.addr_view.resizeColumnToContents(5)     
        self.addr_view.resizeColumnToContents(6)
        self.addr_view.resizeColumnToContents(7)
#public
    @QtCore.Slot()
    def longoperationcomplete(self):
        data = g_DataManager.currentRva
        self.setRefOffset(data)
                
    def setRefOffset(self, data):
        if not data:
            return  
        self._update_current_offset(self.refs_view, self.refsto_model, data)
        self._update_current_offset(self.refsfrom_view, self.refsfrom_model, data)
        self._update_ref_tabs(data)
        self._update_function_name(data)
        
    def filterByColumn(self, col_num, str):
        self.addr_sorted_model.setFilterRegExp(QtCore.QRegExp(str, QtCore.Qt.CaseInsensitive, QtCore.QRegExp.FixedString));
        self.addr_sorted_model.setFilterKeyColumn(col_num)
        
    def filterChanged(self):
        self.filterByColumn(self.filter_combo.currentIndex(), self.filter_edit.text() )

    def OnCreate(self, form):
        """
        Called when the plugin form is created
        """
        
        #init data structures:
        self.functionsMap = dict()
        self.addr_list = []
        self._loadLocals()
        
        # Get parent widget
        self.parent = self.FormToPySideWidget(form)
        
        # Create models
        self.table_model = TableModel_t(self.addr_list)
        self.subDataManager = DataManager()
        
        #init
        self.addr_sorted_model = QtGui.QSortFilterProxyModel()    
        self.addr_sorted_model.setDynamicSortFilter(True)
        self.addr_sorted_model.setSourceModel(self.table_model)
        self.addr_view = FunctionsView_t(g_DataManager, self._COLOR_HILIGHT_FUNC, self.table_model)
        self.addr_view.setModel(self.addr_sorted_model)
        self.addr_view.setSortingEnabled(True)
        self.addr_view.setWordWrap(False)
        self.addr_view.setAlternatingRowColors(True)
        self.addr_view.horizontalHeader().setStretchLastSection(False);
        self.addr_view.verticalHeader().show()
    
        self.adjustColumnsToContents()
        #
        self.refsto_model = RefsTableModel_t(self.addr_list, True)
        self.refs_view = FunctionsView_t(self.subDataManager, self._COLOR_HILIGHT_REFTO, self.refsto_model)
        self._setup_sorted_model(self.refs_view, self.refsto_model)
        self.refs_view.setColumnHidden(RefsTableModel_t.COL_TOADDR, True)
        self.refs_view.setWordWrap(False)
        self.refs_view.setAlternatingRowColors(True)
        
        font = self.refs_view.font()
        font.setPointSize(8)
        self.refs_view.setFont(font)
        #
        self.refsfrom_model = RefsTableModel_t(self.addr_list, False)
        self.refsfrom_view = FunctionsView_t(self.subDataManager, self._COLOR_HILIGHT_REFFROM, self.refsfrom_model)
        self._setup_sorted_model(self.refsfrom_view, self.refsfrom_model)
        self.refsfrom_view.setColumnHidden(RefsTableModel_t.COL_TOADDR, True)
        self.refsfrom_view.setWordWrap(False)
        self.refsfrom_view.setAlternatingRowColors(True)
        
        #important for proper order of objects destruction:
        self.table_model.setParent(self.addr_sorted_model)
        self.addr_sorted_model.setParent(self.addr_view)

        # connect SIGNAL
        g_DataManager.updateSignal.connect(self.longoperationcomplete)
        
        # Create a Tab widget for references:
        self.refs_tabs = QtGui.QTabWidget()
        self.refs_tabs.insertTab(0, self.refs_view, "Is refered by")
        self.refs_tabs.insertTab(1, self.refsfrom_view, "Refers to")
        
        # Create filter
        self.filter_edit = QtGui.QLineEdit()
        self.filter_edit.setPlaceholderText("keyword")
        self.filter_combo = QtGui.QComboBox()
        
        self.filter_combo.addItems(TableModel_t.header_names)
        self.filter_combo.setCurrentIndex(TableModel_t.COL_NAME)
        #connect SIGNAL
        self.filter_edit.textChanged.connect(self.filterChanged)
        self.filter_combo.activated.connect(self.filterChanged)
        

        filter_panel = QtGui.QFrame()
        filter_layout = QtGui.QHBoxLayout()
        filter_layout.addWidget(QtGui.QLabel("Where "))
        filter_layout.addWidget(self.filter_combo)
        filter_layout.addWidget(QtGui.QLabel(" contains: "))
        filter_layout.addWidget(self.filter_edit)
        
        filter_panel.setLayout(filter_layout)
        self.filter_edit.setFixedHeight(20)
        filter_panel.setFixedHeight(40)
        filter_panel.setAutoFillBackground(True)
        #
        self.refs_label = QtGui.QLabel("Function")
        self.refs_label.setTextFormat(QtCore.Qt.RichText)
        self.refs_label.setWordWrap(True)
        
        panel1 = QtGui.QFrame()
        layout1 = QtGui.QVBoxLayout()
        panel1.setLayout(layout1)
        
        layout1.addWidget(filter_panel)
        layout1.addWidget(self.addr_view)      
        layout1.setContentsMargins(0,0,0,0)

        panel2 = QtGui.QFrame()
        layout2 = QtGui.QVBoxLayout()
        layout2.addWidget(self.refs_label)
        layout2.addWidget(self.refs_tabs)
        layout2.addWidget(self._makeButtonsPanel())  
        layout2.setContentsMargins(0,10,0,0)
        panel2.setLayout(layout2)
        
        self.main_splitter = QtGui.QSplitter()
        self.main_splitter.setOrientation(QtCore.Qt.Vertical)
        self.main_splitter.addWidget(panel1)
        self.main_splitter.addWidget(panel2)
        
        # Populate PluginForm
        layout = QtGui.QVBoxLayout()
        layout.addWidget(self.main_splitter)
        layout.setSpacing(0)
        layout.setContentsMargins(0,0,0,0)
        self.parent.setLayout(layout)

        idaapi.set_dock_pos(PLUGIN_NAME, "IDA HExview-1", idaapi.DP_RIGHT)
    
    def _makeButtonsPanel(self):
        buttons_panel = QtGui.QFrame()
        buttons_layout = QtGui.QHBoxLayout()
        buttons_panel.setLayout(buttons_layout)
        
        importButton = QtGui.QPushButton("Load names")
        importButton.clicked.connect(self.importNames)
        buttons_layout.addWidget(importButton)
        
        exportButton = QtGui.QPushButton("Save names")
        exportButton.clicked.connect(self.exportNames)
        buttons_layout.addWidget(exportButton)
        return buttons_panel
        
    def importNames(self):
        file_name, ext = QtGui.QFileDialog.getOpenFileName( None, "Export functions names", QtCore.QDir.homePath(), "TXT Files (*.txt);;All files (*)")
        if file_name is not None and len(file_name) > 0 :
            loaded = self._loadFunctionsNames(file_name)
            if loaded == 0:
                idaapi.warning("Failed importing functions names! Not matching offsets!")
            else:
                idaapi.info("Imported %d function names " % (loaded))
                
    def exportNames(self):
        file_name, ext = QtGui.QFileDialog.getSaveFileName( None, "Import functions names", QtCore.QDir.homePath(), "TXT Files (*.txt);;All files (*)")
        if file_name is not None and len(file_name) > 0 :
            if self._saveFunctionsNames(file_name) == False:
                idaapi.warning("Failed exporting functions names!")
            else:
                idaapi.info("Exported to: "+ file_name)
        
    def OnClose(self, form):
        """
        Called when the plugin form is closed
        """
        #clear last selection
        self.addr_view.hilight_addr(BADADDR)
        self.refs_view.hilight_addr(BADADDR)
        self.refsfrom_view.hilight_addr(BADADDR)
        del self
        print "Closed"

    def Show(self):
        """Creates the form if not created or focuses it if it was"""
        return PluginForm.Show(self,
                               PLUGIN_NAME,
                               options = PluginForm.FORM_PERSIST)

# --------------------------------------------------------------------------
class IFLMenuManager():
    """ Manages menu items belonging to IFL"""

    # public
    def __init__(self):
        self.menuItems = None
        
    # private
    def _makeMenuItems(self):
        self.menuItems = list()
        
        if self._addMenuItem("View/", PLUGIN_NAME, PLUGIN_HOTKEY, 0, self._queryItem1, None) == False:
            return False

        return True
        
    def _destroyMenuItems(self):
        if self.menuItems is None:
            return
        for mItem in self.menuItems:
            idaapi.del_menu_item(self.menuItems)
            
    def _addMenuItem(self, menupath, name, hotkey, flags, pyfunc, args):
        menuItem = idaapi.add_menu_item(menupath, name, hotkey, flags, pyfunc, args)
        if menuItem is None:
            return False
        self.menuItems.append(menuItem)
        return True
        
    def _queryItem1(self):
        open_form()
        

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

def open_form():
    global m_functionInfoForm
    global g_DataManager
    #-----
    try:
        g_DataManager
    except:
        g_DataManager = DataManager()
    #-----
    try:
        m_functionInfoForm
    except:
        idaapi.msg("Loading Interactive Function List...")
        m_functionInfoForm = FunctionsListForm_t()

    m_functionInfoForm.Show()

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

#IDA api:

class funclister_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Interactive Functions List"

    help = "Interactive Function List. Comments? Remarks? Mail to: hasherezade@op.pl"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ''
    
    IFLMngr = None
 
    def init(self):
        IFLMngr = IFLMenuManager()
        if IFLMngr._makeMenuItems():
            return idaapi.PLUGIN_OK
        IFLMngr._destroyMenuItems()
        return idaapi.PLUGIN_SKIP
        
    def run(self, arg):
        open_form()
        pass

    def term(self):
        pass

def PLUGIN_ENTRY():
    return funclister_t()
    


repository: 19h/ida-plugins
file: aidapal-main/idapal.py
wanted_name: aiDAPal (Google AI - Multi/Single Func)
comment: aiDAPal: Google AI assistance for Hex-Rays (Multi/Single Function Analysis - google-genai SDK)
created_at: 2025-04-08 22:54:27+00:00
pushed_at: 2025-04-08 23:07:17+00:00
url:
https://github.com/19h/ida-plugins/blob/b045b6f6bc9ae2a53104986fb74f3fe1198e5d2e/aidapal-main/idapal.py

content:
# -*- coding: utf-8 -*-
"""
summary: IDA Pro plugin using Google AI (Gemini) for code analysis assistance.

description:
  This plugin integrates with Google's Generative AI API (Gemini) via the google-genai SDK
  to provide suggestions for function names, comments, and local
  variable renames within IDA Pro's decompiled pseudocode view.

  It adds context menu items to the Pseudocode view to trigger the
  analysis FOR THE CURRENT FUNCTION AND ITS CALLEES (up to depth 2).
  The results are presented in a custom PyQt5 dockable widget with tabs
  for each analyzed function, where the user can review and selectively
  apply the suggestions across multiple functions.

  Includes logic to decompile and include called functions (up to depth 2)
  in the context provided to the AI. Decompilation and other IDA API calls
  are performed synchronously in the main IDA thread.

  Requires:
  - IDA Pro 7.6+ (with Python 3 and PyQt5 support)
  - Hex-Rays Decompiler
  - google-genai library (`pip install google-genai`)
  - pydantic library (`pip install pydantic`)
  - A Google AI API Key set in the environment variable `GOOGLE_API_KEY`.

DEBUG VERSION: Includes extensive logging and multi-function support.
"""

# --- Imports ---
import ida_kernwin
import ida_hexrays
import ida_funcs
import ida_name
import ida_bytes
import idaapi
import idautils
import idc
import ida_xref
import ida_nalt # Although get_cmt moved, keep for potential future use
import ida_ua # Added for instruction analysis in get_callee_decompilation
import ida_idp

import threading
import json
import textwrap
import os
import sys
from functools import partial
import traceback # For detailed error logging

# Third-party libraries
try:
    from PyQt5 import QtCore, QtGui, QtWidgets
    from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QScrollArea,
                                 QLabel, QVBoxLayout, QHBoxLayout, QGridLayout,
                                 QGroupBox, QCheckBox, QPushButton, QFrame,
                                 QTabWidget) # Added QTabWidget
    from PyQt5.QtGui import QPalette, QColor
    print("aiDAPal DEBUG: PyQt5 imported successfully.")
except ImportError:
    print("aiDAPal Error: PyQt5 not found. Please ensure it's installed in IDA's Python environment.")
    # Optionally, prevent the plugin from loading if PyQt5 is missing
    # raise ImportError("PyQt5 is required for the aiDAPal UI.")

try:
    from pydantic import BaseModel, Field
    print("aiDAPal DEBUG: pydantic imported successfully.")
except ImportError:
    print("aiDAPal Error: pydantic not found. Please install it: pip install pydantic")
    # raise ImportError("pydantic is required for aiDAPal.")

try:
    # Use the new SDK: google-genai
    from google import genai
    from google.genai import types
    from google.genai import errors as google_genai_errors # Import specific error types
    print("aiDAPal DEBUG: google-genai imported successfully.")
except ImportError:
    print("aiDAPal Error: google-genai not found. Please install it: pip install google-genai")
    # raise ImportError("google-genai is required for aiDAPal.")


# --- Configuration ---

# Google AI Configuration
# Read API Key from environment variable for security
GOOGLE_AI_API_KEY = os.environ.get("GOOGLE_API_KEY") # NEW: Use GOOGLE_API_KEY for google-genai
# Default model to use if multiple are available or for the primary action
# See https://ai.google.dev/models/gemini for available models (check compatibility with google-genai)
DEFAULT_GEMINI_MODEL = "gemini-2.5-pro-exp-03-25"

# List of Google AI models to offer as actions in the context menu.
# You can add more models here if desired, e.g., "gemini-1.5-pro-latest"
# Each model name will create a separate context menu entry.
# Ensure the models listed are compatible with the API usage below.
MODELS_TO_REGISTER = [DEFAULT_GEMINI_MODEL] # Keep it simple for debugging

# Safety settings for Google AI (adjust thresholds as needed)
# Convert to google.genai.types.SafetySetting format
try:
    DEFAULT_SAFETY_SETTINGS = [
        types.SafetySetting(category='HARM_CATEGORY_HATE_SPEECH', threshold='BLOCK_NONE'),
        types.SafetySetting(category='HARM_CATEGORY_DANGEROUS_CONTENT', threshold='BLOCK_NONE'),
        types.SafetySetting(category='HARM_CATEGORY_HARASSMENT', threshold='BLOCK_NONE'),
        types.SafetySetting(category='HARM_CATEGORY_SEXUALLY_EXPLICIT', threshold='BLOCK_NONE'),
        types.SafetySetting(category='HARM_CATEGORY_CIVIC_INTEGRITY', threshold='BLOCK_NONE'),
    ]
    print("aiDAPal DEBUG: Default safety settings converted to google.genai.types format.")
except NameError:
    print("aiDAPal Warning: google.genai.types not fully imported. Safety settings might not work correctly.")
    DEFAULT_SAFETY_SETTINGS = [] # Fallback to empty list
except AttributeError:
    print("aiDAPal Warning: google.genai.types.SafetySetting not found. Safety settings might not work correctly.")
    DEFAULT_SAFETY_SETTINGS = [] # Fallback


# Timeout for API requests in seconds (Note: Timeout for generate_content might not be directly supported yet in google-genai's config)
API_TIMEOUT = 180

# --- Concurrency Control ---
g_analysis_in_progress = set() # Set of primary func_ea currently being analyzed
g_analysis_lock = threading.Lock() # To protect the set
print("aiDAPal DEBUG: Concurrency control variables initialized.")


# --- Pydantic Models for Conceptual Structure ---
# Although we use an explicit dict schema for the API call,
# these models help document the expected structure.

class VariableRename(BaseModel):
    """Represents a suggested rename for a single variable."""
    original_name: str = Field(description="The original name of the local variable or argument as seen in the pseudocode.")
    new_name: str = Field(description="The suggested new, more descriptive name for the variable.")
    rename_reason: str | None = Field(None, description="The reason for the suggested rename.")
    rename_reason_findings: str | None = Field(None, description="Any specific findings that led to the suggested rename.")

class SingleFunctionAnalysis(BaseModel):
    """Defines the expected structure for the analysis of ONE function."""
    # function_ea: int = Field(description="The entry address (effective address) of the function being analyzed.") # REMOVED
    original_function_name: str = Field(description="The original name of the function (including sub_...) as provided in the context header.") # ADDED
    function_name: str = Field(description="A concise and descriptive suggested name for the function, suitable for IDA Pro naming conventions.")
    comment: str = Field(description="A multi-line comment explaining the function's purpose, parameters (if discernible), return value (if discernible), and overall logic. Formatted for C-style block comments without the enclosing /* */.")
    variables: list[VariableRename] = Field(description="A list of suggested renames for local variables and function arguments within this specific function.")
    observations: list[dict] | None = Field(None, description="Observations about this function.")
    function_name_reason: str | None = Field(None, description="Reason for the suggested function name.")
    function_name_reason_findings: str | None = Field(None, description="Findings for the suggested function name.")
    comment_reason: str | None = Field(None, description="Reason for the suggested comment.")
    comment_reason_findings: str | None = Field(None, description="Findings for the suggested comment.")


class MultiFunctionAnalysisResult(BaseModel):
    """Defines the expected structure of the JSON analysis result from the AI for multiple functions."""
    function_analyses: list[SingleFunctionAnalysis] = Field(description="A list containing the analysis results for each function provided in the context.")


# --- Explicit Schema Definition for Google AI API (MULTI-FUNCTION) ---
# This dictionary schema remains the same as it's passed to the API.
explicit_multi_function_analysis_schema = {
    "type": "object",
    "properties": {
        "function_analyses": {
            "type": "array",
            "description": "A list containing the analysis results for each function provided in the context.",
            "items": {
                "type": "object",
                "properties": {
                    "original_function_name": {
                        "type": "string",
                        "description": "The original name of the function (e.g., 'sub_140001000' or 'MyFunction') exactly as it appears in the '// --- Function: NAME (EA) ---' header provided in the context."
                    },
                    "observations": {
                        "type": "array",
                        "description": "Any observations about this function that are important to consider when trying to interpret its purpose, parameters, return value, and overall logic.",
                        "items": {
                            "type": "object",
                            "properties": {
                                "observation": {
                                    "type": "string",
                                    "description": "A single observation about this function."
                                },
                                "observation_impact": {
                                    "type": "string",
                                    "description": "The impact of the observation on this function's purpose, parameters, return value, and overall logic."
                                }
                            },
                            "required": ["observation", "observation_impact"],
                            "propertyOrdering": ["observation", "observation_impact"],
                        }
                    },
                    "comment_reason": {
                        "type": "string",
                        "description": "The reason for the suggested comment for this function. Reason must precede the comment."
                    },
                    "comment_reason_findings": {
                        "type": "string",
                        "description": "Any specific findings that led to the suggested comment for this function. Reason must precede the comment."
                    },
                    "function_name_reason": {
                        "type": "string",
                        "description": "The reason for the suggested function name for this function. Reason must precede the function name."
                    },
                    "function_name_reason_findings": {
                        "type": "string",
                        "description": "Any specific findings that led to the suggested function name for this function. Reason must precede the function name."
                    },
                    "variables": {
                        "type": "array",
                        "description": "A list of suggested renames for local variables and function arguments *within this specific function only*. If unsure, leave the variable out or suggest the original name. Reason must precede the new name. Variables must precede new function name and comment.",
                        "items": {
                            "type": "object",
                            "properties": {
                                "rename_reason": {
                                    "type": "string",
                                    "description": "The reason for the suggested rename."
                                },
                                "rename_reason_findings": {
                                    "type": "string",
                                    "description": "Any specific findings that led to the suggested rename."
                                },
                                "original_name": {
                                    "type": "string",
                                    "description": "The original name of the local variable or argument as seen in this function's pseudocode."
                                },
                                "new_name": {
                                    "type": "string",
                                    "description": "The suggested new, more descriptive name for the variable."
                                }
                            },
                            "required": ["rename_reason", "rename_reason_findings", "original_name", "new_name"],
                            "propertyOrdering": ["rename_reason", "rename_reason_findings", "original_name", "new_name"],
                        }
                    },
                    "comment": {
                        "type": "string",
                        "description": "A multi-line comment explaining this function's purpose, parameters, return value, and logic. Formatted for C-style block comments without the enclosing /* */. Use plain text only."
                    },
                    "function_name": {
                        "type": "string",
                        "description": "A concise and descriptive suggested name for this function, suitable for IDA Pro naming conventions."
                    },
                },
                "required": [
                    "original_function_name",
                    "observations",
                    "function_name_reason",
                    "function_name_reason_findings",
                    "function_name",
                    "comment_reason",
                    "comment_reason_findings",
                    "comment",
                    "variables"
                ],
                "propertyOrdering": [
                    "original_function_name",
                    "observations",
                    "function_name_reason",
                    "function_name_reason_findings",
                    "function_name",
                    "comment_reason",
                    "comment_reason_findings",
                    "comment",
                    "variables"
                ],
            }
        }
    },
    # Require the top-level list
    "required": ["function_analyses"]
}
print("aiDAPal DEBUG: Explicit MULTI-FUNCTION JSON schema defined (using original_function_name).")


# --- UI Widget Classes (Mostly Unchanged, used within tabs) ---

class FunctionNameWidget(QWidget):
    """Widget to display and accept/reject the suggested function name."""
    accepted = True # Default state

    def __init__(self, function_name):
        super(FunctionNameWidget, self).__init__()
        # print(f"aiDAPal DEBUG: [FunctionNameWidget.__init__] Initializing with name: '{function_name}'") # Less noisy
        layout = QVBoxLayout()
        layout.setAlignment(QtCore.Qt.AlignLeft | QtCore.Qt.AlignVCenter)

        group_box = QGroupBox("Suggested Function Name")
        group_layout = QHBoxLayout()
        group_layout.setAlignment(QtCore.Qt.AlignLeft | QtCore.Qt.AlignVCenter)
        group_layout.setSpacing(10)

        self.checkbox = QCheckBox()
        self.checkbox.setCheckState(QtCore.Qt.Checked)
        self.checkbox.stateChanged.connect(self.accepted_state_change)

        self.name_label = QLabel(function_name)
        self.name_label.setWordWrap(True) # Allow wrapping if name is long

        group_layout.addWidget(self.checkbox)
        group_layout.addWidget(self.name_label)

        group_box.setLayout(group_layout)
        layout.addWidget(group_box)
        self.setLayout(layout)
        # print("aiDAPal DEBUG: [FunctionNameWidget.__init__] Initialization complete.")

    def accepted_state_change(self, state):
        self.accepted = (state == QtCore.Qt.Checked)
        # print(f'aiDAPal DEBUG: [FunctionNameWidget.accepted_state_change] Function Name Accepted state changed to: {self.accepted}')


class CommentWidget(QWidget):
    """Widget to display and accept/reject the suggested comment."""
    accepted = True # Default state

    def __init__(self, comment):
        super(CommentWidget, self).__init__()
        # print(f"aiDAPal DEBUG: [CommentWidget.__init__] Initializing with comment (length: {len(comment)}).")
        layout = QVBoxLayout()
        layout.setAlignment(QtCore.Qt.AlignLeft | QtCore.Qt.AlignVCenter)

        group_box = QGroupBox("Suggested Comment")
        group_layout = QHBoxLayout()
        # Align checkbox to top-left of the comment area
        group_layout.setAlignment(QtCore.Qt.AlignLeft | QtCore.Qt.AlignTop)
        group_layout.setSpacing(10)

        self.checkbox = QCheckBox()
        self.checkbox.setCheckState(QtCore.Qt.Checked)
        self.checkbox.stateChanged.connect(self.accepted_state_change)

        self.comment_area = QLabel(comment)
        self.comment_area.setWordWrap(True)
        # Set a reasonable minimum width, height will adjust
        self.comment_area.setMinimumWidth(400)
        self.comment_area.setTextInteractionFlags(QtCore.Qt.TextSelectableByMouse) # Allow text selection

        group_layout.addWidget(self.checkbox)
        group_layout.addWidget(self.comment_area)
        group_box.setLayout(group_layout)
        layout.addWidget(group_box)
        self.setLayout(layout)
        # print("aiDAPal DEBUG: [CommentWidget.__init__] Initialization complete.")

    def accepted_state_change(self, state):
        self.accepted = (state == QtCore.Qt.Checked)
        # print(f'aiDAPal DEBUG: [CommentWidget.accepted_state_change] Comment Accepted state changed to: {self.accepted}')


class VariableWidget(QWidget):
    """Widget to display and accept/reject suggested variable renames in a scrollable grid."""
    def __init__(self, variables):
        super(VariableWidget, self).__init__()
        # print(f"aiDAPal DEBUG: [VariableWidget.__init__] Initializing with {len(variables)} variable suggestions.")
        main_layout = QVBoxLayout()
        main_layout.setAlignment(QtCore.Qt.AlignTop | QtCore.Qt.AlignLeft)

        group_box = QGroupBox("Suggested Variable Renames")
        group_layout = QGridLayout()
        group_layout.setAlignment(QtCore.Qt.AlignTop | QtCore.Qt.AlignLeft)
        group_layout.setSpacing(10)
        group_layout.setColumnStretch(1, 1) # Allow labels to expand
        group_layout.setColumnStretch(3, 1) # Allow labels to expand

        self.checkboxes = []
        self.variable_data = variables # Store the original data
        columns = 2 # Adjust number of columns for variable display
        # print(f"aiDAPal DEBUG: [VariableWidget.__init__] Creating grid with {columns} columns.")
        for i, var_data in enumerate(variables):
            row = i // columns
            col_base = (i % columns) * 4 # 4 widgets per entry: checkbox, old, arrow, new

            original_name = var_data.get('original_name', 'N/A')
            new_name = var_data.get('new_name', 'N/A')

            checkbox = QCheckBox()
            checkbox.setCheckState(QtCore.Qt.Checked)
            checkbox.stateChanged.connect(lambda state, index=i: self.accepted_state_change(state, index))
            self.checkboxes.append(checkbox)

            lbl_original = QLabel(original_name)
            lbl_arrow = QLabel("→") # Use arrow symbol
            lbl_new = QLabel(new_name)
            lbl_original.setTextInteractionFlags(QtCore.Qt.TextSelectableByMouse)
            lbl_new.setTextInteractionFlags(QtCore.Qt.TextSelectableByMouse)

            # Add widgets to the grid layout
            group_layout.addWidget(checkbox, row, col_base + 0)
            group_layout.addWidget(lbl_original, row, col_base + 1)
            group_layout.addWidget(lbl_arrow, row, col_base + 2, alignment=QtCore.Qt.AlignCenter)
            group_layout.addWidget(lbl_new, row, col_base + 3)

        group_box.setLayout(group_layout)

        # Create ScrollArea
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded) # Show scrollbar only when needed
        scroll_area.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded)
        scroll_area.setWidget(group_box)
        # Set a minimum height for the scroll area to prevent it collapsing
        scroll_area.setMinimumHeight(100) # Reduced min height slightly
        # print("aiDAPal DEBUG: [VariableWidget.__init__] ScrollArea created.")

        main_layout.addWidget(scroll_area)
        self.setLayout(main_layout)
        # print("aiDAPal DEBUG: [VariableWidget.__init__] Initialization complete.")

    def accepted_state_change(self, state, index):
        # This callback currently doesn't modify overall state,
        # relies on get_selected_variables() for final check.
        # print(f'aiDAPal DEBUG: [VariableWidget.accepted_state_change] Variable {index} Accepted State changed to: {state == QtCore.Qt.Checked}')
        pass

    def get_selected_variables(self):
        """Returns a list of variable data dicts for the checked items."""
        selected = []
        for i, checkbox in enumerate(self.checkboxes):
            if checkbox.isChecked():
                selected.append(self.variable_data[i])
        # print(f"aiDAPal DEBUG: [VariableWidget.get_selected_variables] Returning {len(selected)} selected variables.")
        return selected


# --- Main UI Form Class (Modified for Tabs) ---

class aiDAPalUIForm(ida_kernwin.PluginForm):
    """The main dockable widget form for displaying AI suggestions using tabs for multiple functions."""

    def __init__(self, analysis_results_list, primary_func_ea):
        super(aiDAPalUIForm, self).__init__()
        primary_func_ea_str = f"0x{primary_func_ea:X}" if primary_func_ea else "None"
        print(f"aiDAPal DEBUG: [aiDAPalUIForm.__init__] Initializing multi-function form for analysis triggered by {primary_func_ea_str}.")
        self.analysis_results = analysis_results_list if analysis_results_list else []
        self.primary_func_ea = primary_func_ea # Store the EA of the function that triggered the analysis
        # Store references to widgets, keyed by function EA
        self.widgets_by_ea = {} # { func_ea: {'name': FunctionNameWidget, 'comment': CommentWidget, 'vars': VariableWidget}, ... }
        self.parent_widget = None # Store the PyQt widget reference
        print(f"aiDAPal DEBUG: [aiDAPalUIForm.__init__] Received {len(self.analysis_results)} function analysis results.")
        print(f"aiDAPal DEBUG: [aiDAPalUIForm.__init__] Initialization complete for {primary_func_ea_str}.")

    def OnCreate(self, form):
        """Called when the IDA form is created."""
        print("aiDAPal DEBUG: [aiDAPalUIForm.OnCreate] Form creation started.")
        try:
            self.parent_widget = self.FormToPyQtWidget(form)
            print("aiDAPal DEBUG: [aiDAPalUIForm.OnCreate] Parent widget obtained.")
            self.PopulateForm()
            print("aiDAPal DEBUG: [aiDAPalUIForm.OnCreate] Form population finished.")
        except Exception as e:
            print(f"aiDAPal Error: [aiDAPalUIForm.OnCreate] Exception: {e}")
            traceback.print_exc()

    def PopulateForm(self):
        """Creates and arranges the PyQt widgets within the form using tabs."""
        print("aiDAPal DEBUG: [aiDAPalUIForm.PopulateForm] Starting form population with tabs.")
        if not self.parent_widget:
            print("aiDAPal Error: [aiDAPalUIForm.PopulateForm] Parent widget not available.")
            return

        # Main layout for the entire form
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(0, 0, 0, 0) # Use full space

        # Create Tab Widget
        tab_widget = QTabWidget()
        self.widgets_by_ea = {} # Reset widget storage

        print(f"aiDAPal DEBUG: [aiDAPalUIForm.PopulateForm] Processing {len(self.analysis_results)} results to create tabs.")
        if not self.analysis_results:
             main_layout.addWidget(QLabel("No analysis results received from AI."))

        # Sort results to potentially show primary function first (optional)
        # sorted_results = sorted(self.analysis_results, key=lambda r: 0 if r.get('function_ea') == self.primary_func_ea else 1)
        sorted_results = self.analysis_results # Or keep original order

        for result_data in sorted_results:
            # IMPORTANT: Use the function_ea added back during the mapping phase
            func_ea = result_data.get('function_ea')
            if func_ea is None:
                print(f"aiDAPal Warning: [aiDAPalUIForm.PopulateForm] Skipping result item missing 'function_ea' (mapping likely failed): {result_data.get('original_function_name', 'N/A')}")
                continue

            func_ea_str = f"0x{func_ea:X}"
            # Use execute_sync to safely get the function name from the main thread during UI setup
            func_name_ida = None
            name_container = [None]
            def get_name_main(ea, container):
                try:
                    container[0] = ida_funcs.get_func_name(ea) or f"sub_{ea:X}"
                    return 1
                except:
                    container[0] = f"sub_{ea:X}" # Fallback
                    return 0
            name_sync_status = ida_kernwin.execute_sync(lambda: get_name_main(func_ea, name_container), ida_kernwin.MFF_READ)
            if name_sync_status == 1:
                func_name_ida = name_container[0]
            else:
                func_name_ida = f"sub_{func_ea_str}" # Fallback if sync failed
                print(f"aiDAPal Warning: [PopulateForm] Failed to get function name for {func_ea_str} via execute_sync.")


            tab_title = f"{func_name_ida} ({func_ea_str})"
            print(f"aiDAPal DEBUG: [aiDAPalUIForm.PopulateForm] Creating tab for: {tab_title}")

            # Create a container widget for the tab content
            tab_content_widget = QWidget()
            tab_layout = QVBoxLayout()
            tab_layout.setAlignment(QtCore.Qt.AlignTop | QtCore.Qt.AlignLeft)

            # --- Function Name Section ---
            func_name_sugg = result_data.get('function_name', 'N/A')
            name_widget = FunctionNameWidget(func_name_sugg)
            tab_layout.addWidget(name_widget)

            # --- Comment Section ---
            comment_sugg = result_data.get('comment', 'No comment suggested.')
            comment_widget = CommentWidget(comment_sugg)
            tab_layout.addWidget(comment_widget)

            # --- Variables Section ---
            variables_sugg = result_data.get('variables', [])
            variable_widget = None # Initialize
            if variables_sugg:
                variable_widget = VariableWidget(variables_sugg)
                tab_layout.addWidget(variable_widget)
            else:
                tab_layout.addWidget(QLabel("No variable rename suggestions for this function."))

            # --- (Optional UI Enhancement - Display Observations/Reasons) ---
            # observations = result_data.get('observations', [])
            # if observations:
            #     obs_text = "Observations:\n" + "\n".join([f"- {o.get('observation','?')} (Impact: {o.get('observation_impact','?')})" for o in observations])
            #     obs_label = QLabel(obs_text)
            #     obs_label.setWordWrap(True)
            #     obs_label.setStyleSheet("QLabel { color : gray; margin-top: 10px; }") # Style hint
            #     tab_layout.addWidget(obs_label)
            # # Similar labels could be added for function_name_reason, comment_reason etc.

            tab_layout.addStretch(1) # Push content upwards
            tab_content_widget.setLayout(tab_layout)

            # Add tab to the tab widget
            tab_widget.addTab(tab_content_widget, tab_title)

            # Store widget references for this function
            self.widgets_by_ea[func_ea] = {
                'name': name_widget,
                'comment': comment_widget,
                'vars': variable_widget, # Will be None if no variables suggested
                'data': result_data # Store original data for easy access
            }
            print(f"  aiDAPal DEBUG: [PopulateForm] Widgets stored for {func_ea_str}")

        # Add TabWidget to the main layout
        main_layout.addWidget(tab_widget)

        # --- Buttons ---
        print("aiDAPal DEBUG: [aiDAPalUIForm.PopulateForm] Creating buttons.")
        accept_button = QPushButton("Apply Selected")
        cancel_button = QPushButton("Close")

        # Button layout
        button_layout = QHBoxLayout()
        button_layout.addStretch() # Push buttons to the right
        button_layout.addWidget(accept_button)
        button_layout.addWidget(cancel_button)

        # Add button layout to the main layout (below tabs)
        main_layout.addLayout(button_layout)

        # Connect buttons
        print("aiDAPal DEBUG: [aiDAPalUIForm.PopulateForm] Connecting button signals.")
        accept_button.clicked.connect(self.on_accept_clicked)
        cancel_button.clicked.connect(self.on_cancel_clicked)

        self.parent_widget.setLayout(main_layout)
        # Set a reasonable minimum size for the dockable window
        self.parent_widget.setMinimumSize(600, 500) # Increased size slightly for tabs
        print("aiDAPal DEBUG: [aiDAPalUIForm.PopulateForm] Layout set, minimum size set. Population complete.")

    def on_accept_clicked(self):
        """Handles the 'Apply Selected' button click, gathering selections and running updates in main thread."""
        print("aiDAPal DEBUG: [aiDAPalUIForm.on_accept_clicked] 'Apply Selected' clicked.")

        changes_by_ea = {} # { func_ea: {'function_name': str|None, 'comment': str|None, 'variables': list[dict]}, ... }

        print(f"aiDAPal DEBUG: [on_accept_clicked] Gathering selections from {len(self.widgets_by_ea)} functions...")
        for func_ea, widgets in self.widgets_by_ea.items():
            func_ea_str = f"0x{func_ea:X}"
            # print(f"  aiDAPal DEBUG: [on_accept_clicked] Processing EA: {func_ea_str}") # Noisy
            selected_changes = {}
            original_data = widgets['data']

            # Check Function Name
            if widgets['name'] and widgets['name'].accepted:
                selected_changes['function_name'] = original_data.get('function_name')
            else:
                selected_changes['function_name'] = None

            # Check Comment
            if widgets['comment'] and widgets['comment'].accepted:
                selected_changes['comment'] = original_data.get('comment')
            else:
                selected_changes['comment'] = None

            # Check Variables
            selected_vars = []
            if widgets['vars']: # Check if VariableWidget exists
                selected_vars = widgets['vars'].get_selected_variables()
            selected_changes['variables'] = selected_vars

            # Only add to final dict if there's something to apply for this function
            if selected_changes['function_name'] or selected_changes['comment'] or selected_changes['variables']:
                changes_by_ea[func_ea] = selected_changes
                # print(f"  aiDAPal DEBUG: [on_accept_clicked] Changes queued for {func_ea_str}") # Noisy
            # else: print(f"  aiDAPal DEBUG: [on_accept_clicked] No changes selected for {func_ea_str}") # Noisy

        if not changes_by_ea:
            print("aiDAPal DEBUG: [aiDAPalUIForm.on_accept_clicked] No changes selected across all functions. Closing.")
            ida_kernwin.info("aiDAPal: No changes were selected.")
            self.Close(0)
            return

        print(f"aiDAPal DEBUG: [aiDAPalUIForm.on_accept_clicked] Final changes gathered for {len(changes_by_ea)} functions. Preparing main thread update.")

        # --- Execute Update Logic in Main Thread ---
        update_success = False
        update_result_container = [False] # Container for success status

        def run_update_in_main_thread(changes_dict, result_container):
            print(f"aiDAPal DEBUG: [run_update_in_main_thread] Executing in main thread...")
            try:
                # Call the actual update logic function
                result_container[0] = self._perform_ida_updates(changes_dict)
                print(f"aiDAPal DEBUG: [run_update_in_main_thread] _perform_ida_updates returned: {result_container[0]}")
                return 1 # Indicate success to execute_sync
            except Exception as e_update_main:
                print(f"aiDAPal Error: [run_update_in_main_thread] Exception during main thread update: {e_update_main}")
                traceback.print_exc()
                result_container[0] = False # Ensure failure state
                return 0 # Indicate failure to execute_sync

        # Execute synchronously
        sync_status = ida_kernwin.execute_sync(
            lambda: run_update_in_main_thread(changes_by_ea, update_result_container),
            ida_kernwin.MFF_WRITE # Essential for database modifications
        )

        print(f"aiDAPal DEBUG: [on_accept_clicked] execute_sync for updates returned status: {sync_status}")

        if sync_status == 1:
            update_success = update_result_container[0]
            print(f"aiDAPal DEBUG: [on_accept_clicked] Update success status from main thread: {update_success}")
            if update_success:
                 ida_kernwin.info(f"aiDAPal: Applied selected changes to {len(changes_by_ea)} function(s).")
            else:
                 ida_kernwin.warning(f"aiDAPal: Failed to apply some or all changes. Check Output window.")
        else:
            print(f"aiDAPal Error: [on_accept_clicked] execute_sync for updates failed or was cancelled (status: {sync_status}).")
            ida_kernwin.warning(f"aiDAPal: Failed to execute update operation in main thread.")

        # Close the form regardless of update success/failure after attempting
        print("aiDAPal DEBUG: [aiDAPalUIForm.on_accept_clicked] Closing form.")
        self.Close(0)
        print("aiDAPal DEBUG: [aiDAPalUIForm.on_accept_clicked] Exiting.")


    def _perform_ida_updates(self, changes_by_ea):
        """
        Internal function containing the logic to apply changes to IDA.
        MUST be called from the main IDA thread (e.g., via execute_sync).
        """
        print(f"aiDAPal DEBUG: [_perform_ida_updates] Starting update process for {len(changes_by_ea)} functions (MAIN THREAD).")

        overall_changes_made = False
        overall_refresh_needed = False # Track if Hex-Rays view needs refresh for *any* function
        all_updates_succeeded = True # Track if any individual update fails

        for func_ea, changes in changes_by_ea.items():
            func_ea_str = f"0x{func_ea:X}"
            print(f"--- Updating function {func_ea_str} ---")
            func_changes_made = False
            func_refresh_needed = False

            if func_ea == idaapi.BADADDR:
                 print(f"  aiDAPal Error: [_perform_ida_updates] Invalid function entry address {func_ea_str} in changes dict. Skipping.")
                 all_updates_succeeded = False # Mark as failed if we skip invalid data
                 continue

            # --- Pre-Update Checks ---
            func_t_obj = ida_funcs.get_func(func_ea) # Safe now in main thread
            if not func_t_obj:
                print(f"  aiDAPal Error: [_perform_ida_updates] Pre-check failed: ida_funcs.get_func(0x{func_ea:X}) returned None. Cannot apply comment or reliably rename.")
                # Don't mark all_updates_succeeded as False yet, name setting might still work
            else:
                 print(f"  aiDAPal DEBUG: [_perform_ida_updates] Pre-check OK: Found func_t object for 0x{func_ea:X}.")

            can_decompile_for_rename = False
            if changes.get("variables"): # Only check decompilation if variable renames are requested
                try:
                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Pre-check: Attempting ida_hexrays.decompile(0x{func_ea:X}) for rename validation...")
                    cfunc = ida_hexrays.decompile(func_ea)
                    if cfunc:
                        print(f"  aiDAPal DEBUG: [_perform_ida_updates] Pre-check OK: Decompilation successful for 0x{func_ea:X}.")
                        can_decompile_for_rename = True
                    else:
                        print(f"  aiDAPal Warning: [_perform_ida_updates] Pre-check failed: ida_hexrays.decompile(0x{func_ea:X}) returned None. Variable rename will likely fail.")
                        all_updates_succeeded = False # Mark failure if rename is expected but pre-check fails
                except ida_hexrays.DecompilationFailure as e:
                    print(f"  aiDAPal Warning: [_perform_ida_updates] Pre-check failed: ida_hexrays.decompile(0x{func_ea:X}) raised DecompilationFailure: {e}. Variable rename will fail.")
                    all_updates_succeeded = False # Mark failure
                except Exception as e_decomp:
                     print(f"  aiDAPal Error: [_perform_ida_updates] Pre-check failed: Unexpected error during ida_hexrays.decompile(0x{func_ea:X}): {e_decomp}")
                     all_updates_succeeded = False # Mark failure


            # --- Apply Comment ---
            new_cmt_text = changes.get("comment")
            if new_cmt_text:
                if func_t_obj: # Check if we got the object earlier
                    wrapped_cmt = '\n'.join(textwrap.wrap(new_cmt_text, width=80))
                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Attempting to set comment (length: {len(wrapped_cmt)}) for {func_ea_str}...")
                    if ida_funcs.set_func_cmt(func_t_obj, wrapped_cmt, False): # Safe now
                        print(f"  aiDAPal DEBUG: [_perform_ida_updates] Successfully set function comment for {func_ea_str}.")
                        func_changes_made = True
                    else:
                        print(f"  aiDAPal DEBUG: [_perform_ida_updates] Failed to set function comment for {func_ea_str} (set_func_cmt returned False).")
                        all_updates_succeeded = False # Mark failure
                else:
                    # Already logged the error above
                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Skipping comment set for {func_ea_str} due to missing func_t object.")
                    all_updates_succeeded = False # Cannot apply comment without func_t
            # else: print(f"  aiDAPal DEBUG: [_perform_ida_updates] No comment to apply for {func_ea_str}.")

            # --- Apply Function Name ---
            new_name = changes.get("function_name")
            if new_name:
                print(f'  aiDAPal DEBUG: [_perform_ida_updates] Attempting function name update to: {new_name} for {func_ea_str}')
                # Check current name only if func_t_obj exists, otherwise assume different
                current_name = ida_funcs.get_func_name(func_ea) if func_t_obj else None # Safe now
                if func_t_obj and current_name == new_name:
                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Function name '{new_name}' is already set for {func_ea_str}. Skipping.")
                else:
                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Calling set_name for {func_ea_str} with name '{new_name}'...")
                    # set_name might work even if get_func failed, try it
                    if ida_name.set_name(func_ea, new_name, ida_name.SN_CHECK | ida_name.SN_FORCE | ida_name.SN_NOWARN): # Safe now
                        print(f'  aiDAPal DEBUG: [_perform_ida_updates] Successfully updated function name for {func_ea_str}.')
                        func_changes_made = True
                    else:
                        print(f'  aiDAPal DEBUG: [_perform_ida_updates] Failed to update function name for {func_ea_str} (set_name returned False).')
                        all_updates_succeeded = False # Mark failure
            # else: print(f"  aiDAPal DEBUG: [_perform_ida_updates] No function name to apply for {func_ea_str}.")

            # --- Apply Variable Renames ---
            variables_to_apply = changes.get("variables", [])
            if variables_to_apply:
                if can_decompile_for_rename: # Check pre-check result
                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Attempting {len(variables_to_apply)} variable renames for {func_ea_str}...")
                    renamed_count = 0
                    failed_count = 0
                    skipped_count = 0

                    for var_data in variables_to_apply:
                        original_name = var_data.get('original_name')
                        new_name_suggested = var_data.get('new_name')

                        if not original_name or not new_name_suggested:
                            print("    aiDAPal DEBUG: [_perform_ida_updates] Skipping rename: Missing original or new name in data.")
                            failed_count += 1
                            continue

                        if original_name == new_name_suggested:
                            print(f"    aiDAPal DEBUG: [_perform_ida_updates] Skipping rename: Suggested name '{new_name_suggested}' is same as original.")
                            skipped_count += 1
                            continue

                        print(f"    aiDAPal DEBUG: [_perform_ida_updates] Renaming '{original_name}' -> '{new_name_suggested}' using ida_hexrays.rename_lvar for func {func_ea_str}")
                        if ida_hexrays.rename_lvar(func_ea, original_name, new_name_suggested): # Safe now
                            print(f"      aiDAPal DEBUG: [_perform_ida_updates] Successfully renamed '{original_name}' to '{new_name_suggested}'.")
                            renamed_count += 1
                            func_changes_made = True
                            func_refresh_needed = True # Variable renames require refresh
                        else:
                            print(f"      aiDAPal DEBUG: [_perform_ida_updates] rename_lvar failed for '{original_name}' -> '{new_name_suggested}' in {func_ea_str}.")
                            failed_count += 1
                            all_updates_succeeded = False # Mark failure

                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Variable renaming for {func_ea_str} finished. Success: {renamed_count}, Failed: {failed_count}, Skipped (same name): {skipped_count}")
                else:
                    # Already logged the decompilation failure reason
                     print(f"  aiDAPal DEBUG: [_perform_ida_updates] Skipping variable renames for {func_ea_str} due to failed decompilation pre-check.")
                     # If variables were requested but couldn't be applied, mark as failure
                     all_updates_succeeded = False
            # else: print(f"  aiDAPal DEBUG: [_perform_ida_updates] No variable renames to apply for {func_ea_str}.")

            # --- Mark for Refresh ---
            if func_refresh_needed:
                print(f"  aiDAPal DEBUG: [_perform_ida_updates] Refresh needed for {func_ea_str}. Marking cfunc dirty.")
                ida_hexrays.mark_cfunc_dirty(func_ea) # Safe now
                overall_refresh_needed = True
            elif func_changes_made:
                 print(f"  aiDAPal DEBUG: [_perform_ida_updates] Changes made (name/comment) for {func_ea_str}, marking dirty.")
                 # Only mark dirty if the function is known to the decompiler
                 if can_decompile_for_rename or func_t_obj: # Check if it's a known/decompilable func
                    ida_hexrays.mark_cfunc_dirty(func_ea) # Safe now
                    overall_refresh_needed = True # Mark dirty might trigger refresh anyway
                 else:
                    print(f"  aiDAPal DEBUG: [_perform_ida_updates] Skipping mark_cfunc_dirty for {func_ea_str} as it might not be a valid Hex-Rays function.")


            if func_changes_made:
                overall_changes_made = True

            print(f"--- Finished update for {func_ea_str} ---")

        # --- Final Status ---
        if overall_changes_made:
            print(f"aiDAPal DEBUG: [_perform_ida_updates] Overall changes were made.")
            if overall_refresh_needed:
                 print(f"aiDAPal DEBUG: [_perform_ida_updates] Refresh needed for one or more functions.")
        else:
            print(f"aiDAPal DEBUG: [_perform_ida_updates] No changes were applied.")

        print(f"aiDAPal DEBUG: [_perform_ida_updates] Multi-function update finished. Success status: {all_updates_succeeded}")
        return all_updates_succeeded # Return True if all individual updates succeeded


    def on_cancel_clicked(self):
        """Handles the 'Close' button click."""
        print("aiDAPal DEBUG: [aiDAPalUIForm.on_cancel_clicked] 'Close' button clicked. Closing form.")
        self.Close(0) # Indicate cancellation

    def OnClose(self, form):
        """Called when the IDA form is being closed."""
        primary_func_ea_str = f"0x{self.primary_func_ea:X}" if self.primary_func_ea else "None"
        print(f"aiDAPal DEBUG: [aiDAPalUIForm.OnClose] Form closing for analysis triggered by {primary_func_ea_str}.")
        # Clean up references if necessary (handled by UI wrapper now)
        pass


# --- UI Wrapper Class ---

class aiDAPalUI:
    """Helper class to instantiate and show the UI form, preventing duplicates for the *same analysis trigger*."""
    open_forms = {} # Class variable: Key: primary_func_ea, Value: aiDAPalUIForm instance

    def __init__(self, analysis_results_list=None, primary_func_ea=None):
        print("aiDAPal DEBUG: [aiDAPalUI.__init__] Initializing UI wrapper.")
        if primary_func_ea is None or primary_func_ea == idaapi.BADADDR:
            print("aiDAPal Error: [aiDAPalUI.__init__] Cannot show UI without a valid primary function EA.")
            return

        func_ea_str = f"0x{primary_func_ea:X}"
        print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Target primary function EA: {func_ea_str}")

        # Check if a form for an analysis *triggered by this primary function* is already open
        if primary_func_ea in aiDAPalUI.open_forms:
            print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Form for analysis triggered by {func_ea_str} already open. Attempting to activate.")
            existing_form = aiDAPalUI.open_forms[primary_func_ea]
            try:
                widget = existing_form.GetWidget()
                if widget:
                    print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Activating existing widget for {func_ea_str}.")
                    ida_kernwin.activate_widget(widget, True)
                else:
                    print(f"aiDAPal Warning: [aiDAPalUI.__init__] Could not get widget for existing form {func_ea_str}.")
            except Exception as e:
                print(f"aiDAPal Warning: [aiDAPalUI.__init__] Error activating existing widget for {func_ea_str}: {e}")
            return # Don't create a new one

        if analysis_results_list is None:
            print("aiDAPal Warning: [aiDAPalUI.__init__] Initialized with no results list.")
            self.analysis_results = []
        else:
            self.analysis_results = analysis_results_list
            print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] UI wrapper initialized for analysis triggered by {func_ea_str} with {len(self.analysis_results)} results.")

        # Create and show the form
        print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Creating new aiDAPalUIForm instance for {func_ea_str}.")
        self.plg = aiDAPalUIForm(self.analysis_results, primary_func_ea) # Pass results list and primary EA
        aiDAPalUI.open_forms[primary_func_ea] = self.plg # Store instance before showing
        print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Form instance stored in open_forms for {func_ea_str}.")

        # Set a unique caption based on the primary function
        func_name_str = ida_funcs.get_func_name(primary_func_ea) or f"sub_{primary_func_ea:X}"
        form_caption = f"aiDAPal Suggestions: {func_name_str} & Callees" # Updated caption
        print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Setting form caption: '{form_caption}'")
        # Show as a dockable, persistent window
        show_flags = ida_kernwin.WOPN_PERSIST | ida_kernwin.WOPN_RESTORE
        print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Calling Show() for {func_ea_str} with flags {show_flags}.")
        self.plg.Show(form_caption, show_flags)

        # Override OnClose to remove from tracking dict
        print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] Wrapping OnClose for form triggered by {func_ea_str}.")
        original_on_close = self.plg.OnClose
        # Use a lambda that captures the primary_func_ea for removal
        wrapped_on_close = lambda form, ea=primary_func_ea: self.on_close_handler(form, ea, original_on_close)
        self.plg.OnClose = wrapped_on_close # Assign the wrapper

        print(f"aiDAPal DEBUG: [aiDAPalUI.__init__] UI wrapper initialization complete for {func_ea_str}.")

    # Separate handler function for OnClose logic
    def on_close_handler(self, form, func_ea, original_on_close_func):
        func_ea_str = f"0x{func_ea:X}"
        print(f"aiDAPal DEBUG: [aiDAPalUI.on_close_handler] Wrapper called for form triggered by {func_ea_str}.")
        print(f"aiDAPal DEBUG: [aiDAPalUI.on_close_handler] Removing form for {func_ea_str} from tracking.")
        aiDAPalUI.open_forms.pop(func_ea, None) # Remove safely using the captured EA
        print(f"aiDAPal DEBUG: [aiDAPalUI.on_close_handler] Calling original OnClose for {func_ea_str}.")
        original_on_close_func(form) # Call original OnClose
        print(f"aiDAPal DEBUG: [aiDAPalUI.on_close_handler] Original OnClose returned for {func_ea_str}.")


# --- Helper Functions (get_references_from_function, get_function_data_ref_comments, get_callee_decompilation are unchanged) ---
# NOTE: get_callee_decompilation MUST be called from the main thread via execute_sync

def get_references_from_function(func_ea):
    """
    Walks a function's items and collects unique data reference target addresses.
    MUST be called from the main IDA thread or via execute_sync.
    """
    # print(f"aiDAPal DEBUG: [get_references_from_function] Starting for EA: 0x{func_ea:X}") # Noisy
    refs = set()
    func = ida_funcs.get_func(func_ea) # Main thread only
    if not func:
        print(f"aiDAPal Error: [get_references_from_function] Could not get function object for EA 0x{func_ea:X}")
        return refs

    fii = ida_funcs.func_item_iterator_t()
    ok = fii.set_range(func.start_ea, func.end_ea) # Main thread only
    if not ok:
        print(f"aiDAPal Error: [get_references_from_function] Failed to initialize function item iterator for 0x{func_ea:X}")
        return refs

    item_count = 0
    dref_count = 0
    while fii.next_code(): # Main thread only
        insn_ea = fii.current()
        item_count += 1
        dref_ea = ida_xref.get_first_dref_from(insn_ea) # Main thread only
        while dref_ea != idaapi.BADADDR:
            refs.add(dref_ea)
            dref_count += 1
            dref_ea = ida_xref.get_next_dref_from(insn_ea, dref_ea) # Main thread only

    # print(f"aiDAPal DEBUG: [get_references_from_function] Finished for 0x{func_ea:X}. Scanned: {item_count}. Refs: {dref_count}. Unique: {len(refs)}.") # Noisy
    return refs

def get_function_data_ref_comments(current_func_ea):
    """
    Extracts comments from data locations referenced by the given function.
    MUST be called from the main IDA thread or via execute_sync.
    """
    ea_str = f"0x{current_func_ea:X}" if current_func_ea is not None else "None"
    print(f"aiDAPal DEBUG: [get_function_data_ref_comments] Starting for EA: {ea_str}")

    if current_func_ea is None:
        print("aiDAPal DEBUG: [get_function_data_ref_comments] No function EA provided.")
        return None

    # This function itself calls main-thread-only functions
    references = get_references_from_function(current_func_ea)
    if not references:
         print(f"aiDAPal DEBUG: [get_function_data_ref_comments] No data references found for {ea_str}.")
         return None

    print(f"aiDAPal DEBUG: [get_function_data_ref_comments] Found {len(references)} unique data references for {ea_str}. Checking for comments...")

    comment_lines = []
    for ref in sorted(list(references)):
        cmt_text = ""
        cmt_r = ida_bytes.get_cmt(ref, True)  # Main thread only
        cmt_n = ida_bytes.get_cmt(ref, False) # Main thread only

        if cmt_r: cmt_text = cmt_r.strip()
        if cmt_n:
            cmt_n_stripped = cmt_n.strip()
            if cmt_text and cmt_text != cmt_n_stripped: cmt_text += f" // {cmt_n_stripped}"
            elif not cmt_text: cmt_text = cmt_n_stripped

        if cmt_text:
            name = ida_name.get_name(ref) # Main thread only
            display_ref = name if name else f"0x{ref:X}"
            comment_lines.append(f" * {display_ref}: {cmt_text}")

    if not comment_lines:
        print(f"aiDAPal DEBUG: [get_function_data_ref_comments] No comments found for any references in {ea_str}.")
        return None
    else:
        header = "/* Referenced Data Comments:"
        footer = "*/"
        full_comment = header + "\n" + "\n".join(comment_lines) + "\n" + footer
        print(f"aiDAPal DEBUG: [get_function_data_ref_comments] Found {len(comment_lines)} comments for {ea_str}. Returning comment block (length: {len(full_comment)}).")
        return full_comment

def get_callee_decompilation(start_ea, max_depth):
    """
    Performs a depth-limited traversal of the call graph starting from start_ea,
    decompiles functions up to max_depth, and returns their code.
    MUST be called from the main IDA thread or via execute_sync.
    """
    print(f"aiDAPal DEBUG: [get_callee_decompilation] Starting traversal from 0x{start_ea:X} with max_depth={max_depth} (MAIN THREAD)")
    decompiled_funcs = {}
    funcs_to_visit = set()
    visited_ea = set()

    # Check start_ea validity (already requires main thread)
    if ida_funcs.get_func(start_ea) is None:
        print(f"aiDAPal Error: [get_callee_decompilation] Invalid start_ea: 0x{start_ea:X}")
        return decompiled_funcs

    funcs_to_visit.add(start_ea)
    visited_ea.add(start_ea)

    for current_depth in range(max_depth + 1):
        print(f"aiDAPal DEBUG: [get_callee_decompilation] Processing depth {current_depth}")
        next_level_funcs = set()
        if not funcs_to_visit:
            print(f"aiDAPal DEBUG: [get_callee_decompilation] No more functions to visit at depth {current_depth}.")
            break

        for func_ea in funcs_to_visit:
            # Decompile (requires main thread)
            if func_ea not in decompiled_funcs:
                print(f"  aiDAPal DEBUG: [get_callee_decompilation] Decompiling 0x{func_ea:X} at depth {current_depth}")
                func_name = ida_name.get_name(func_ea) or f"sub_{func_ea:X}" # Main thread
                try:
                    cfunc = ida_hexrays.decompile(func_ea) # Main thread
                    if cfunc:
                        decompiled_code = str(cfunc)
                        decompiled_funcs[func_ea] = f"// --- Function: {func_name} (0x{func_ea:X}) ---\n{decompiled_code}\n// --- End Function: {func_name} (0x{func_ea:X}) ---\n"
                        # print(f"    aiDAPal DEBUG: Decompilation successful for 0x{func_ea:X} (length: {len(decompiled_code)})") # Noisy
                    else:
                        decompiled_funcs[func_ea] = f"// --- Decompilation FAILED for {func_name} (0x{func_ea:X}) ---\n"
                        print(f"    aiDAPal Warning: Decompilation returned None for 0x{func_ea:X}")
                except ida_hexrays.DecompilationFailure as e:
                    decompiled_funcs[func_ea] = f"// --- Decompilation ERROR for {func_name} (0x{func_ea:X}): {e} ---\n"
                    print(f"    aiDAPal Error: Decompilation failed for 0x{func_ea:X}: {e}")
                except Exception as e:
                    decompiled_funcs[func_ea] = f"// --- Decompilation UNEXPECTED ERROR for {func_name} (0x{func_ea:X}): {e} ---\n"
                    print(f"    aiDAPal Error: Unexpected error decompiling 0x{func_ea:X}: {e}")
                    traceback.print_exc()

            # Find Callees (requires main thread) - Using operand iteration method
            if current_depth < max_depth:
                func = ida_funcs.get_func(func_ea) # Main thread
                if not func:
                    print(f"  aiDAPal Warning: [get_callee_decompilation] Could not get func object for 0x{func_ea:X} to find callees.")
                    continue

                current_item_ea = func.start_ea
                while current_item_ea < func.end_ea:
                    # All API calls here require main thread
                    flags = idaapi.get_flags(current_item_ea)
                    if idaapi.is_code(flags):
                        # *** CORRECTED is_call_insn usage ***
                        if idaapi.is_call_insn(current_item_ea):
                            # *** SWITCHED to operand iteration ***
                            for i in range(idaapi.UA_MAXOP): # Iterate through operands
                                op_type = idc.get_operand_type(current_item_ea, i)
                                # Check if the operand is a code address (near or far)
                                if op_type == idc.o_near or op_type == idc.o_far:
                                    ref = idc.get_operand_value(current_item_ea, i)
                                    if ref != idaapi.BADADDR: # Ensure we got a valid address
                                        called_func = ida_funcs.get_func(ref)
                                        # Ensure the reference is the start of a function
                                        if called_func and called_func.start_ea == ref:
                                            callee_ea = ref
                                            if callee_ea not in visited_ea:
                                                # print(f"    aiDAPal DEBUG: Found callee 0x{callee_ea:X} from operand {i} at 0x{current_item_ea:X}.") # Debugging noise
                                                next_level_funcs.add(callee_ea)
                                                visited_ea.add(callee_ea)
                                            # Assuming a direct call instruction targets only one function via operands
                                            break # Move to next instruction once a valid call target is found in operands
                                # *** END operand iteration block ***

                    next_ea = idc.next_head(current_item_ea, func.end_ea) # Main thread
                    if next_ea <= current_item_ea:
                         print(f"  aiDAPal Warning: [get_callee_decompilation] next_head did not advance from 0x{current_item_ea:X}. Breaking loop for 0x{func_ea:X}.")
                         break
                    current_item_ea = next_ea
                    if current_item_ea == idaapi.BADADDR:
                        break

        funcs_to_visit = next_level_funcs

    print(f"aiDAPal DEBUG: [get_callee_decompilation] Traversal finished. Found {len(decompiled_funcs)} functions up to depth {max_depth}.")
    return decompiled_funcs


# --- Google AI Interaction (UPDATED for google-genai SDK) ---

def do_google_ai_analysis(code_prompt, model_name):
    """
    Sends the code prompt to Google AI (Gemini) via google-genai SDK and expects
    structured JSON output containing analysis for multiple functions.
    Uses original_function_name for mapping.

    Args:
        code_prompt (str): The C code for multiple functions and any extra context.
        model_name (str): The Gemini model name (e.g., "gemini-1.5-flash-latest").

    Returns:
        list: A list of dictionaries, where each dictionary contains the analysis
              results for one function (matching SingleFunctionAnalysis schema, but with
              original_function_name instead of function_ea), or None on failure.
    """
    print(f"aiDAPal DEBUG: [do_google_ai_analysis] Starting MULTI-FUNCTION analysis request to Google AI model: {model_name} using google-genai SDK")

    # API Key Check
    if not GOOGLE_AI_API_KEY or GOOGLE_AI_API_KEY == "YOUR_API_KEY_HERE":
        print("aiDAPal Error: [do_google_ai_analysis] GOOGLE_API_KEY not set or is placeholder.")
        ida_kernwin.warning("aiDAPal Error: Google AI API Key not configured.\nSet the GOOGLE_API_KEY variable in the script or environment variable.")
        return None
    print("aiDAPal DEBUG: [do_google_ai_analysis] API Key seems present.")

    try:
        # Create Client
        print("aiDAPal DEBUG: [do_google_ai_analysis] Creating genai.Client...")
        # Note: Ensure GOOGLE_AI_API_KEY is the correct key for the Gemini Developer API
        client = genai.Client(api_key=GOOGLE_AI_API_KEY)
        print("aiDAPal DEBUG: [do_google_ai_analysis] genai.Client created.")

        # Prepare Generation Config using types.GenerateContentConfig
        print("aiDAPal DEBUG: [do_google_ai_analysis] Setting up types.GenerateContentConfig with MULTI-FUNCTION schema...")
        try:
            generation_config = types.GenerateContentConfig(
                response_mime_type="application/json",
                response_schema=explicit_multi_function_analysis_schema,
                temperature=0.0,
                # Add safety settings here
                safety_settings=DEFAULT_SAFETY_SETTINGS
                # candidate_count=1 # Explicitly ask for one candidate if needed
            )
            print("aiDAPal DEBUG: [do_google_ai_analysis] types.GenerateContentConfig set.")
        except NameError:
             print("aiDAPal Error: [do_google_ai_analysis] google.genai.types not available. Cannot create GenerateContentConfig.")
             ida_kernwin.warning("aiDAPal Error: google.genai SDK types could not be used. Check installation.")
             return None
        except AttributeError:
            print("aiDAPal Error: [do_google_ai_analysis] GenerateContentConfig or SafetySetting attribute error. Check google.genai installation/version.")
            ida_kernwin.warning("aiDAPal Error: google.genai SDK types seem incompatible. Check installation/version.")
            return None


        # Prepare prompt content (just the string)
        contents = code_prompt
        print(f"aiDAPal DEBUG: [do_google_ai_analysis] Sending prompt (length: {len(contents)} chars) to {model_name} via client.models.generate_content...")

        # Call the API using the client
        print("aiDAPal DEBUG: [do_google_ai_analysis] Calling client.models.generate_content()...")
        response = client.models.generate_content(
            model=f'models/{model_name}', # Model name usually needs 'models/' prefix
            contents=contents,
            config=generation_config
            # Note: safety_settings are now part of the config object
        )
        print(f"aiDAPal DEBUG: [do_google_ai_analysis] Received response object from {model_name}.")

        # --- Process Response ---
        # Check if response.text exists and is not empty
        if not hasattr(response, 'text') or not response.text:
             print("aiDAPal Error: [do_google_ai_analysis] Response has no text content or is empty/blocked.")
             # Try to get more info if available (less standardized in google-genai response object for blocks)
             # The prompt_feedback might contain block reasons
             block_reason = "Unknown"
             safety_ratings = "N/A"
             try:
                 if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
                     if hasattr(response.prompt_feedback, 'block_reason'):
                         block_reason = response.prompt_feedback.block_reason.name if response.prompt_feedback.block_reason else "Not Specified"
                     if hasattr(response.prompt_feedback, 'safety_ratings'):
                        safety_ratings = str([f"{sr.category.name}: {sr.probability.name}" for sr in response.prompt_feedback.safety_ratings])

                 # Check candidates as a fallback, structure might differ
                 elif hasattr(response, 'candidates') and response.candidates:
                     # google-genai candidate structure might be different
                     candidate = response.candidates[0]
                     if hasattr(candidate, 'finish_reason') and candidate.finish_reason:
                         # Map finish reason if possible, e.g., SAFETY
                         if candidate.finish_reason.name == 'SAFETY':
                             block_reason = 'SAFETY'
                         else:
                             block_reason = candidate.finish_reason.name
                     if hasattr(candidate, 'safety_ratings') and candidate.safety_ratings:
                         safety_ratings = str([f"{sr.category.name}: {sr.probability.name}" for sr in candidate.safety_ratings])

                 print(f"  aiDAPal DEBUG: Finish/Block Reason: {block_reason}"); print(f"  aiDAPal DEBUG: Safety Ratings: {safety_ratings}")
                 ida_kernwin.warning(f"aiDAPal: Google AI response was empty or blocked.\nReason: {block_reason}\nSafety: {safety_ratings}\nCheck safety settings or prompt.")

             except Exception as e_info:
                 print(f"aiDAPal Error: [do_google_ai_analysis] Error getting finish reason/safety details: {e_info}")
                 ida_kernwin.warning("aiDAPal: Google AI response was empty or blocked.")
             return None

        print("aiDAPal DEBUG: [do_google_ai_analysis] Response has text content. Attempting to parse response.text...")
        try:
            raw_text = response.text
            print(f"aiDAPal DEBUG: [do_google_ai_analysis] Raw response text (first 100 chars): {raw_text[:100]}...")

            # --- Robust JSON Extraction ---
            json_start = -1
            json_end = -1
            cleaned_text = ""
            search_start_index = 0

            # 1. Find where to start searching for '{'
            # Prefer '```json' then '```' as start markers
            start_marker_json = raw_text.find("```json")
            start_marker_plain = raw_text.find("```")

            if start_marker_json != -1:
                search_start_index = start_marker_json + 7
                print(f"aiDAPal DEBUG: [do_google_ai_analysis] Found '```json' marker. Starting search for '{{' at index {search_start_index}.")
            elif start_marker_plain != -1:
                # Check if there's a '{' likely after this '```'
                potential_json_start = raw_text.find('{', start_marker_plain + 3)
                if potential_json_start != -1:
                     # Assume this ``` is the start fence
                     search_start_index = start_marker_plain + 3
                     print(f"aiDAPal DEBUG: [do_google_ai_analysis] Found '```' marker likely before JSON. Starting search for '{{' at index {search_start_index}.")
                # else: Keep search_start_index = 0, maybe this ``` is the end fence or irrelevant text

            # 2. Find the first '{' at or after search_start_index
            json_start = raw_text.find('{', search_start_index)

            if json_start != -1:
                print(f"aiDAPal DEBUG: [do_google_ai_analysis] Found potential JSON start '{{' at index {json_start}.")
                # 3. Found the starting '{'. Now find the end.
                # Look for the closing fence '```' after the starting '{'
                end_marker = raw_text.find("```", json_start)
                print(f"aiDAPal DEBUG: [do_google_ai_analysis] Searching for closing '```' after index {json_start}. Found at: {end_marker}")

                if end_marker != -1:
                    # Found a closing fence. Find the last '}' before it.
                    json_end = raw_text.rfind('}', json_start, end_marker)
                    print(f"aiDAPal DEBUG: [do_google_ai_analysis] Found closing '```'. Searching for last '}}' between {json_start} and {end_marker}. Found at: {json_end}")
                else:
                    # No closing fence found after '{'. Find the last '}' in the rest of the string.
                    json_end = raw_text.rfind('}', json_start)
                    print(f"aiDAPal DEBUG: [do_google_ai_analysis] No closing '```' found. Searching for last '}}' after index {json_start}. Found at: {json_end}")


                # 4. Validate indices and extract
                if json_end != -1 and json_end > json_start:
                    cleaned_text = raw_text[json_start : json_end + 1].strip()
                    print(f"aiDAPal DEBUG: [do_google_ai_analysis] Extracted potential JSON (length: {len(cleaned_text)}). Preview (first/last 50 chars): '{cleaned_text[:50]}...{cleaned_text[-50:]}'")
                else:
                    print("aiDAPal Warning: [do_google_ai_analysis] Found '{' but could not find a corresponding '}' afterwards or indices were invalid.")
                    # Keep cleaned_text empty
            else:
                print(f"aiDAPal Warning: [do_google_ai_analysis] Could not find the start of a JSON object ('{{') after searching from index {search_start_index}.")
                # Keep cleaned_text empty

            # 5. Check if extraction was successful before attempting parse
            if not cleaned_text:
                 print(f"aiDAPal Error: [do_google_ai_analysis] Failed to extract JSON block from the response text.")
                 print(f"aiDAPal DEBUG: Raw response text that failed extraction was:\n---\n{raw_text}\n---")
                 ida_kernwin.warning(f"aiDAPal failed to extract a JSON block from the Google AI response.\nCheck Output window for raw response.")
                 # Returning None here prevents the JSONDecodeError later, provides clearer context
                 return None

            print("aiDAPal DEBUG: [do_google_ai_analysis] Calling json.loads()...")
            parsed_response = json.loads(cleaned_text)
            print("aiDAPal DEBUG: [do_google_ai_analysis] JSON parsing successful.")

            # --- Log Parsed JSON Payload (YAML-like) ---
            def format_value(value, max_len=100):
                """Format value for logging, handling strings specifically."""
                if isinstance(value, str):
                    # Remove newlines and truncate
                    cleaned_str = value.replace('\n', '\\n').replace('\r', '')
                    if len(cleaned_str) > max_len:
                        return f'"{cleaned_str[:max_len]}..."'
                    else:
                        return f'"{cleaned_str}"'
                elif isinstance(value, (int, float, bool)) or value is None:
                    return repr(value) # Use standard repr for others
                else:
                    # For lists/dicts, just indicate type, recursion handles content
                    return f"<{type(value).__name__}>"

            def log_dict_yaml_style(data, indent=0, max_depth=10):
                """Recursively log dictionary/list in YAML-like format."""
                if indent >= max_depth * 2: # Prevent excessively deep recursion
                    print("  " * indent + "... (max depth reached)")
                    return

                prefix = "  " * indent
                if isinstance(data, dict):
                    for key, value in data.items():
                        if isinstance(value, (dict, list)):
                            print(f"{prefix}- {key}:")
                            log_dict_yaml_style(value, indent + 1, max_depth)
                        else:
                            print(f"{prefix}- {key}: {format_value(value)}")
                elif isinstance(data, list):
                    for i, item in enumerate(data):
                        if isinstance(item, (dict, list)):
                            print(f"{prefix}- [{i}]") # Indicate list item index
                            log_dict_yaml_style(item, indent + 1, max_depth)
                        else:
                            print(f"{prefix}- {format_value(item)}")
                else: # Should not happen if called with dict/list initially
                    print(f"{prefix}{format_value(data)}")

            print("aiDAPal DEBUG: [do_google_ai_analysis] Logging parsed JSON payload:")
            print("--- START PARSED JSON ---")
            log_dict_yaml_style(parsed_response)
            print("--- END PARSED JSON ---")
            # --- End Logging ---

            # --- Validation for Multi-Function Schema (remains the same) ---
            print("aiDAPal DEBUG: [do_google_ai_analysis] Validating parsed multi-function structure...")
            if not isinstance(parsed_response, dict):
                raise ValueError("Top level response is not a dictionary.")
            if "function_analyses" not in parsed_response:
                raise ValueError("Missing required top-level key: 'function_analyses'")
            if not isinstance(parsed_response["function_analyses"], list):
                raise ValueError("'function_analyses' field is not a list.")

            validated_analyses = []
            # Define the required keys based on the updated schema
            required_inner_keys = [
                "original_function_name", # ADDED
                "observations", "function_name_reason",
                "function_name_reason_findings", "function_name", "comment_reason",
                "comment_reason_findings", "comment", "variables"
            ]
            required_variable_keys = ["rename_reason", "rename_reason_findings", "original_name", "new_name"]
            required_observation_keys = ["observation", "observation_impact"]

            print(f"aiDAPal DEBUG: [do_google_ai_analysis] Validating {len(parsed_response['function_analyses'])} function analysis items...")
            for i, func_analysis in enumerate(parsed_response["function_analyses"]):
                if not isinstance(func_analysis, dict):
                    print(f"aiDAPal Warning: [do_google_ai_analysis] Item at index {i} in 'function_analyses' is not a dictionary. Skipping item: {func_analysis}")
                    continue

                # Check required keys for the function analysis itself
                valid_item = True
                for key in required_inner_keys:
                    if key not in func_analysis:
                        print(f"aiDAPal Warning: [do_google_ai_analysis] Missing required key '{key}' in function analysis item at index {i}. Providing default/skipping.")
                        if key == 'original_function_name': # Cannot proceed without name for mapping
                            valid_item = False; break
                        elif key == 'variables': func_analysis[key] = []
                        elif key == 'observations': func_analysis[key] = []
                        elif key == 'function_name': func_analysis[key] = f"suggested_name_{i}"
                        elif key == 'comment': func_analysis[key] = "(No comment suggested)"
                        else: func_analysis[key] = "" # Default for reason/findings strings

                if not valid_item: continue # Skip if original_function_name was missing

                # Validate 'observations' list structure (remains same)
                if not isinstance(func_analysis.get('observations'), list):
                     print(f"aiDAPal Warning: [do_google_ai_analysis] 'observations' field is not a list in item {i}. Setting to empty list.")
                     func_analysis['observations'] = []
                else:
                    validated_observations = []
                    for obs_idx, obs_item in enumerate(func_analysis['observations']):
                        if not isinstance(obs_item, dict):
                            print(f"aiDAPal Warning: [do_google_ai_analysis] Observation item at index {obs_idx} in function analysis {i} is not a dictionary. Skipping item: {obs_item}")
                            continue
                        valid_obs = True
                        for obs_key in required_observation_keys:
                            if obs_key not in obs_item:
                                print(f"aiDAPal Warning: [do_google_ai_analysis] Missing required key '{obs_key}' in observation item {obs_idx} (func analysis {i}). Skipping item: {obs_item}")
                                valid_obs = False; break
                        if valid_obs:
                            validated_observations.append(obs_item)
                    func_analysis['observations'] = validated_observations


                # Validate 'variables' list structure (remains same)
                if not isinstance(func_analysis.get('variables'), list):
                     print(f"aiDAPal Warning: [do_google_ai_analysis] 'variables' field is not a list in item {i}. Setting to empty list.")
                     func_analysis['variables'] = []
                else:
                    validated_variables = []
                    for j, var_item in enumerate(func_analysis.get('variables', [])):
                        if not isinstance(var_item, dict):
                            print(f"aiDAPal Warning: [do_google_ai_analysis] Variable item at index {j} in function analysis {i} is not a dictionary. Skipping item: {var_item}")
                            continue
                        valid_var = True
                        for v_key in required_variable_keys: # Use updated required keys
                            if v_key not in var_item:
                                print(f"aiDAPal Warning: [do_google_ai_analysis] Missing required key '{v_key}' in variable item {j} (func analysis {i}). Skipping item: {var_item}")
                                valid_var = False; break
                        if valid_var:
                            validated_variables.append(var_item)
                    func_analysis['variables'] = validated_variables

                validated_analyses.append(func_analysis) # Add the validated function analysis

            print(f"aiDAPal DEBUG: [do_google_ai_analysis] Validation complete. {len(validated_analyses)} valid function analyses found.")

            # Return the list of validated function analysis dictionaries
            print("aiDAPal DEBUG: [do_google_ai_analysis] Successfully parsed and validated MULTI-FUNCTION JSON. Returning list of analyses.")
            return validated_analyses # Return the list directly

        # Error handling for JSON parsing/validation
        except (json.JSONDecodeError, ValueError) as e:
            print(f"aiDAPal Error: [do_google_ai_analysis] Failed to parse or validate MULTI-FUNCTION JSON: {e}")
            # Check if raw_text was assigned before trying to print it
            if 'raw_text' in locals():
                print(f"aiDAPal DEBUG: Response text that failed parsing was:\n---\n{raw_text}\n---")
            else:
                print("aiDAPal DEBUG: Raw response text was not available for logging.")
            ida_kernwin.warning(f"aiDAPal failed to parse or validate JSON response from Google AI.\nError: {e}\n\nCheck Output window for response text.")
            return None
        except Exception as e_parse:
             print(f"aiDAPal Error: [do_google_ai_analysis] Unexpected error processing response.text: {e_parse}")
             traceback.print_exc()
             ida_kernwin.warning("aiDAPal: Unexpected error processing Google AI response text.")
             return None


    # Error handling for google-genai API calls
    except google_genai_errors.APIError as e_api:
        print(f"aiDAPal Error: [do_google_ai_analysis] Google AI APIError occurred: {e_api}")
        print(f"  Code: {e_api.code}")
        print(f"  Message: {e_api.message}")
        traceback.print_exc()
        ida_kernwin.warning(f"aiDAPal: Google AI API Error ({e_api.code}): {e_api.message}")
        return None
    # Error handling for general exceptions during API interaction
    except ImportError:
        # This case should ideally be caught earlier, but included for completeness
        print("aiDAPal Error: [do_google_ai_analysis] Required Python libraries not found (ImportError).")
        ida_kernwin.warning("aiDAPal Error: Required Python libraries ('google-genai', 'pydantic') not found.")
        return None
    except Exception as e:
        print(f"aiDAPal Error: [do_google_ai_analysis] An unexpected error occurred during API interaction: {e}")
        traceback.print_exc()
        ida_kernwin.warning(f"aiDAPal: An unexpected error occurred during Google AI interaction: {e}")
        return None


# --- Asynchronous Task Handling ---

def do_show_ui(results_list, primary_func_ea):
    """
    Callback function to show the UI form (with tabs) in the main IDA thread.

    Args:
        results_list (list): The list of analysis result dicts from the AI.
        primary_func_ea (int): The EA of the function that triggered the analysis.

    Returns:
        bool: Always False (as required by execute_ui_requests).
    """
    func_ea_str = f"0x{primary_func_ea:X}" if primary_func_ea else "None"
    print(f"aiDAPal DEBUG: [do_show_ui] Entered. Will show multi-function UI triggered by {func_ea_str}.")

    if primary_func_ea is None or primary_func_ea == idaapi.BADADDR:
        print("aiDAPal Error: [do_show_ui] No valid primary function EA provided.")
        return False # Must return bool

    print(f"aiDAPal DEBUG: [do_show_ui] Preparing to show UI for analysis triggered by {func_ea_str}")
    try:
        print(f"aiDAPal DEBUG: [do_show_ui] Instantiating aiDAPalUI for {func_ea_str}...")
        # Pass the list of results and the primary EA to the UI wrapper
        aiDAPalUI(results_list, primary_func_ea)
        print(f"aiDAPal DEBUG: [do_show_ui] aiDAPalUI instantiation complete. UI display initiated.")
    except Exception as e:
        print(f"aiDAPal Error: [do_show_ui] Error creating or showing UI: {e}")
        traceback.print_exc()

    print(f"aiDAPal DEBUG: [do_show_ui] Exiting for {func_ea_str}. Returning False.")
    return False # Required for execute_ui_requests

# --- Analysis Task Wrapper (for concurrency control) ---
# Wrapper remains the same, but async_call now handles multiple functions
def analysis_task_wrapper(primary_cfunc, model, analyze_only_current, context): # MODIFIED: Added analyze_only_current
    """Wrapper for async_call to handle concurrency flag."""
    func_ea_local = primary_cfunc.entry_ea if primary_cfunc else idaapi.BADADDR
    func_ea_str = f"0x{func_ea_local:X}" if func_ea_local != idaapi.BADADDR else "BADADDR"
    analysis_type_str = "SINGLE" if analyze_only_current else "MULTI" # Added for logging
    print(f"aiDAPal DEBUG: [analysis_task_wrapper] Thread started for {analysis_type_str} analysis triggered by {func_ea_str}, model {model}.")
    try:
        print(f"aiDAPal DEBUG: [analysis_task_wrapper] Calling async_call for {func_ea_str}...")
        # Pass the primary cfunc object AND the flags/context
        async_call(primary_cfunc, model, analyze_only_current, context) # MODIFIED: Pass flags/context
        print(f"aiDAPal DEBUG: [analysis_task_wrapper] async_call finished for {func_ea_str}.")
    except Exception as e:
         print(f"aiDAPal Error: [analysis_task_wrapper] Exception during async_call for {func_ea_str}: {e}")
         traceback.print_exc()
    finally:
        # Ensure the flag is cleared even if async_call fails or UI doesn't show
        if func_ea_local != idaapi.BADADDR:
            print(f"aiDAPal DEBUG: [analysis_task_wrapper] Acquiring lock to clear analysis flag for {func_ea_str}...")
            with g_analysis_lock:
                g_analysis_in_progress.discard(func_ea_local) # Use primary EA for the flag
                print(f"aiDAPal DEBUG: [analysis_task_wrapper] Analysis flag cleared for {func_ea_str}. Current in progress: {g_analysis_in_progress}")
        else:
            print("aiDAPal DEBUG: [analysis_task_wrapper] Cannot clear flag, invalid primary func_ea.")
        print(f"aiDAPal DEBUG: [analysis_task_wrapper] Thread finished for analysis triggered by {func_ea_str}.")


# --- MODIFIED ASYNC CALL (Using Name Mapping & google-genai) ---
def async_call(primary_cfunc, model_name, analyze_only_current=False, extra_context=None): # Signature is correct
    """
    Handles the asynchronous call to the AI analysis function for the primary
    function and its callees, expecting multi-function results (using names),
    maps results back to EAs, and schedules the UI update. Uses google-genai SDK.

    Args:
        primary_cfunc (ida_hexrays.cfunc_t): The decompiled object of the primary function.
        model_name (str): The Google AI model to use.
        analyze_only_current (bool): If True, instruct the AI to only analyze the primary function.
        extra_context (str | None): Optional string context (e.g., data references).
    """
    primary_func_ea = idaapi.BADADDR
    primary_func_name_for_prompt = "UNKNOWN_FUNCTION" # Fallback
    if primary_cfunc:
        primary_func_ea = primary_cfunc.entry_ea
        # Get the name for the prompt here, before potentially going async
        # Use execute_sync to safely get the name
        name_container = [None]
        def get_name_main(ea, container):
            try:
                container[0] = ida_funcs.get_func_name(ea) or f"sub_{ea:X}"
                return 1
            except:
                container[0] = f"sub_{ea:X}" # Fallback
                return 0
        name_sync_status = ida_kernwin.execute_sync(lambda: get_name_main(primary_func_ea, name_container), ida_kernwin.MFF_READ)
        if name_sync_status == 1:
            primary_func_name_for_prompt = name_container[0]
        else:
            primary_func_name_for_prompt = f"sub_{primary_func_ea:X}"


    primary_func_ea_str = f"0x{primary_func_ea:X}" if primary_func_ea != idaapi.BADADDR else "BADADDR"
    analysis_type_str = "SINGLE-FUNCTION (with context)" if analyze_only_current else "MULTI-FUNCTION"

    print(f"aiDAPal DEBUG: [async_call] Starting {analysis_type_str} analysis in background thread triggered by {primary_func_ea_str} ({primary_func_name_for_prompt}) using Google AI model: {model_name}")

    if not primary_cfunc:
        print("aiDAPal Error: [async_call] No primary function context (cfunc object) provided.")
        return

    # --- Get Decompiled Code (RUN IN MAIN THREAD) ---
    all_decompiled_codes = None # {ea: code_string}
    decomp_result_container = [None]
    # ... (run_decompilation_in_main_thread and execute_sync call remain the same) ...
    try:
        max_callee_depth = 2
        print(f"aiDAPal DEBUG: [async_call] Preparing to run decompilation in main thread for {primary_func_ea_str}...")
        def run_decompilation_in_main_thread(container):
            print(f"aiDAPal DEBUG: [run_decompilation_in_main_thread] Executing in main thread for {primary_func_ea_str}...")
            try:
                ida_kernwin.show_wait_box(f"Decompiling primary & callees (depth {max_callee_depth}) for {model_name}...")
                container[0] = get_callee_decompilation(primary_func_ea, max_callee_depth)
                ida_kernwin.hide_wait_box()
                print(f"aiDAPal DEBUG: [run_decompilation_in_main_thread] Decompilation finished. Stored result type: {type(container[0])}")
                return 1
            except Exception as e_main:
                 print(f"aiDAPal Error: [run_decompilation_in_main_thread] Exception: {e_main}")
                 traceback.print_exc()
                 ida_kernwin.hide_wait_box()
                 container[0] = None
                 return 0
        print(f"aiDAPal DEBUG: [async_call] Calling execute_sync for decompilation...")
        sync_status = ida_kernwin.execute_sync(lambda: run_decompilation_in_main_thread(decomp_result_container), ida_kernwin.MFF_WRITE)
        print(f"aiDAPal DEBUG: [async_call] execute_sync for decompilation returned status: {sync_status}")
        if sync_status == 1:
            all_decompiled_codes = decomp_result_container[0]
            if isinstance(all_decompiled_codes, dict):
                 print(f"aiDAPal DEBUG: [async_call] Retrieved {len(all_decompiled_codes)} decompiled functions via main thread.")
            elif all_decompiled_codes is None: raise RuntimeError("Decompilation returned None from main thread.")
            else: raise RuntimeError(f"Unexpected type returned from decompilation: {type(all_decompiled_codes)}")
        else: raise RuntimeError(f"Decompilation failed or was cancelled in main thread (status: {sync_status}).")
        if not all_decompiled_codes:
             print(f"aiDAPal Warning: [async_call] No decompiled code returned. Cannot proceed.")
             return
    except Exception as e:
        print(f"aiDAPal Error: [async_call] Failed during decompilation process orchestration for {primary_func_ea_str}: {e}")
        traceback.print_exc()
        return

    # --- Create EA <-> Name Mapping (RUN IN MAIN THREAD) ---
    ea_to_original_name = None # {ea: name_string}
    name_map_container = [None]
    try:
        print(f"aiDAPal DEBUG: [async_call] Preparing to create EA->Name map in main thread...")
        # Get EAs from the decompiled codes dictionary
        target_eas = list(all_decompiled_codes.keys())

        def create_name_map_main(eas, container):
            print(f"aiDAPal DEBUG: [create_name_map_main] Executing in main thread for {len(eas)} EAs...")
            mapping = {}
            try:
                for ea in eas:
                    name = ida_funcs.get_func_name(ea) or f"sub_{ea:X}"
                    mapping[ea] = name
                container[0] = mapping
                print(f"aiDAPal DEBUG: [create_name_map_main] Created map with {len(mapping)} entries.")
                return 1
            except Exception as e_map:
                print(f"aiDAPal Error: [create_name_map_main] Exception: {e_map}")
                traceback.print_exc()
                container[0] = None
                return 0

        print(f"aiDAPal DEBUG: [async_call] Calling execute_sync for name mapping...")
        sync_status_map = ida_kernwin.execute_sync(lambda: create_name_map_main(target_eas, name_map_container), ida_kernwin.MFF_READ)
        print(f"aiDAPal DEBUG: [async_call] execute_sync for name mapping returned status: {sync_status_map}")

        if sync_status_map == 1:
            ea_to_original_name = name_map_container[0]
            if not ea_to_original_name: # Check if map is empty or None
                 raise RuntimeError("Failed to create EA to Name map in main thread.")
            print(f"aiDAPal DEBUG: [async_call] Successfully created EA->Name map with {len(ea_to_original_name)} entries.")
        else:
            raise RuntimeError(f"Name mapping failed or was cancelled in main thread (status: {sync_status_map}).")

    except Exception as e_map_orch:
        print(f"aiDAPal Error: [async_call] Failed during name mapping orchestration: {e_map_orch}")
        traceback.print_exc()
        return # Cannot proceed without the map

    # --- Prepare Prompt ---
    full_code_context_str = ""
    if all_decompiled_codes:
        print(f"aiDAPal DEBUG: [async_call] Formatting {len(all_decompiled_codes)} functions for prompt...")
        sorted_eas = sorted(all_decompiled_codes.keys())
        # Use the generated map to get the correct name for the header
        code_context_list = []
        for ea in sorted_eas:
            code_str = all_decompiled_codes[ea]
            # Find the original header and replace the name if needed (more robust)
            original_name_in_map = ea_to_original_name.get(ea, f"sub_{ea:X}") # Get name from map
            # Reconstruct the header line for clarity in the prompt
            header = f"// --- Function: {original_name_in_map} (0x{ea:X}) ---"
            # Find the existing header in the code string (assuming it's the first line)
            lines = code_str.split('\n', 1)
            if len(lines) > 0 and lines[0].startswith("// --- Function:"):
                code_context_list.append(header + "\n" + lines[1]) # Replace header
            else:
                code_context_list.append(header + "\n" + code_str) # Prepend header if missing

        full_code_context_str = "\n\n---\n\n".join(code_context_list)
        print(f"aiDAPal DEBUG: [async_call] Full code context string length: {len(full_code_context_str)}")
    else:
        print(f"aiDAPal Error: [async_call] No code context generated, cannot proceed.")
        return

    # Build the final prompt content (persona and context remain the same)
    prompt_parts = []
    prompt_parts.append("""You are an expert reverse engineer and software analyst tasked with analyzing decompiled C-like pseudocode from IDA Pro's Hex-Rays.
Your goal is to provide insightful analysis to aid a human analyst in understanding the code's functionality and improving its readability.

Analyze the provided pseudocode, which may contain typical decompilation artifacts such as:
- Abstracted or generic names (e.g., sub_140001000, v1, a2, dword_140020000)
- Unclear variable types or usage patterns
- Complex or flattened control flow
- Compiler optimizations or artifacts
- Missing context (e.g., incomplete type information, external dependencies)

For each function provided, perform a detailed analysis including:
1.  **Function Purpose:** Briefly describe the overall goal or main operation.
2.  **Key Logic/Algorithm:** Identify and explain the core logic, algorithms, or significant operations.
3.  **Parameters:** Describe the inferred purpose and usage of each parameter.
4.  **Return Value:** Describe the inferred meaning or purpose of the return value.
5.  **Suggested Function Name:** Propose a concise, descriptive name based on its purpose, following C/C++ naming conventions (e.g., `calculate_checksum`, `parse_network_packet`), including reasoning and specific code findings.
6.  **Suggested Comment:** Generate a comprehensive C-style comment block explaining the function's purpose, parameters, return value, and key logic, including reasoning and code findings.
7.  **Variable Renames:** Suggest meaningful names for unclear variables (e.g., `v5` to `user_id`), including reasoning and code findings. Only suggest renames that significantly enhance readability.
8.  **Observations:** List notable observations, potential issues (e.g., buffer overflows), complex sections, standard library functions used, or interesting patterns, explaining the impact/significance.

Present your analysis clearly and concisely. Ensure that the `original_function_name` field in your output EXACTLY matches the function name provided in the input context (including the address if present, e.g., 'sub_140001000' or 'ProcessData_1400020A0'). This is critical for mapping the results correctly.

Focus on providing practical insights that would help a reverse engineer understand and potentially refactor the code. Adhere strictly to the requested JSON output format defined by the API schema.
""")

    # --- CHOOSE PROMPT INSTRUCTIONS BASED ON analysis_only_current FLAG ---
    if analyze_only_current:
        print(f"aiDAPal DEBUG: [async_call] Using SINGLE-FUNCTION analysis instructions for primary function {primary_func_name_for_prompt} ({primary_func_ea_str}).")
        prompt_parts.append(f"\nAnalyze ONLY the primary function '{primary_func_name_for_prompt}' (at address {primary_func_ea_str}) from the following C-like pseudocode functions obtained from IDA Pro. Use the other functions provided SOLELY as context for understanding the primary function's interactions.")

        if extra_context:
            print(f"aiDAPal DEBUG: [async_call] Adding data reference context (for primary func {primary_func_ea_str}, length {len(extra_context)}).")
            prompt_parts.append(f"\n/* Context: Data locations referenced by the *primary* function ({primary_func_name_for_prompt} at {primary_func_ea_str}) and their comments (if any):\n{extra_context}\n*/")

        prompt_parts.append(f"\n--- START FUNCTION CODE CONTEXT ---\n{full_code_context_str}\n--- END FUNCTION CODE CONTEXT ---")

        # --- SIMPLIFIED SINGLE-FUNCTION INSTRUCTIONS ---
        prompt_parts.append(f"""
Instructions for analyzing ONLY the primary function '{primary_func_name_for_prompt}' ({primary_func_ea_str}):
Focus your analysis exclusively on the function '{primary_func_name_for_prompt}' identified by its name and address {primary_func_ea_str}.
Ensure the `original_function_name` field in your output EXACTLY matches the function name provided in the input context ('{primary_func_name_for_prompt}'). This is CRITICAL.

Return the results as a single JSON object adhering to the expected schema. The `function_analyses` array MUST contain EXACTLY ONE entry, corresponding ONLY to the primary function '{primary_func_name_for_prompt}' ({primary_func_ea_str}). Do NOT include analysis for any other functions, even though their code was provided for context.

JSON Response:""")
        # --- END SIMPLIFIED SINGLE-FUNCTION INSTRUCTIONS ---

    else: # Original multi-function analysis
        print(f"aiDAPal DEBUG: [async_call] Using MULTI-FUNCTION analysis instructions.")
        prompt_parts.append(f"\nAnalyze ALL of the following C-like pseudocode functions obtained from IDA Pro. The analysis was triggered for the function '{primary_func_name_for_prompt}' at {primary_func_ea_str}.")

        if extra_context:
            print(f"aiDAPal DEBUG: [async_call] Adding data reference context (for primary func {primary_func_ea_str}, length {len(extra_context)}).")
            prompt_parts.append(f"\n/* Context: Data locations referenced by the *triggering* function ({primary_func_name_for_prompt} at {primary_func_ea_str}) and their comments (if any):\n{extra_context}\n*/")

        prompt_parts.append(f"\n--- START FUNCTION CODE CONTEXT ---\n{full_code_context_str}\n--- END FUNCTION CODE CONTEXT ---")

        # --- SIMPLIFIED MULTI-FUNCTION INSTRUCTIONS ---
        prompt_parts.append("""
Instructions:
Analyze EACH function provided in the 'FUNCTION CODE CONTEXT' section above.
For EACH function, ensure the `original_function_name` field in your output EXACTLY matches the function name provided in the input context (e.g., 'sub_140001000' or 'MyFunction'). This is CRITICAL for mapping the results back.

Return the results as a single JSON object adhering to the expected schema, containing a `function_analyses` array with an entry for EACH function presented in the context.

JSON Response:""")
        # --- END SIMPLIFIED MULTI-FUNCTION INSTRUCTIONS ---
    # --- END INSTRUCTION SELECTION ---

    prompt_content = "\n".join(prompt_parts)

    print(f"aiDAPal DEBUG: [async_call] Final {analysis_type_str} prompt content length: {len(prompt_content)} chars for analysis triggered by {primary_func_ea_str}")

    # --- Call AI Analysis using google-genai ---
    results_list_with_names = None # Expecting list with names
    analysis_exception = None
    try:
        wait_box_msg = f"Asking Google AI ({model_name})... ({analysis_type_str})"
        print(f"aiDAPal DEBUG: [async_call] Showing AI wait box for {primary_func_ea_str}: '{wait_box_msg}'")
        ida_kernwin.execute_ui_requests([lambda: ida_kernwin.show_wait_box(wait_box_msg)])

        print(f"aiDAPal DEBUG: [async_call] Calling do_google_ai_analysis for {primary_func_ea_str}...")
        # This function now uses the google-genai SDK
        results_list_with_names = do_google_ai_analysis(prompt_content, model_name) # Expects list or None
        print(f"aiDAPal DEBUG: [async_call] do_google_ai_analysis returned. Result is None: {results_list_with_names is None}. Count: {len(results_list_with_names) if results_list_with_names else 0}")

    except Exception as e: # Catch any unexpected error during the call itself
        print(f"aiDAPal Error: [async_call] Exception during AI analysis call orchestration for {primary_func_ea_str}: {e}")
        analysis_exception = e
        traceback.print_exc()
        results_list_with_names = None # Ensure it's None on error
    finally:
        print(f"aiDAPal DEBUG: [async_call] Scheduling hide_wait_box for {primary_func_ea_str} via execute_ui_requests.")
        ida_kernwin.execute_ui_requests([ida_kernwin.hide_wait_box])
        print(f"aiDAPal DEBUG: [async_call] hide_wait_box request submitted for {primary_func_ea_str}.")

    # --- Map Results Back to EA (Logic remains the same) ---
    results_list_with_ea = []
    if results_list_with_names is not None:
        print(f"aiDAPal DEBUG: [async_call] Mapping {len(results_list_with_names)} results back to EAs...")
        # Create reverse map (Name -> EA). Handle potential duplicates crudely by taking the first encountered.
        name_to_ea = {}
        duplicate_names = set()
        # Ensure ea_to_original_name is a dictionary before iterating
        if isinstance(ea_to_original_name, dict):
            for ea, name in ea_to_original_name.items():
                if name in name_to_ea:
                    if name not in duplicate_names:
                        print(f"aiDAPal Warning: [async_call] Duplicate function name '{name}' encountered in context. Mapping might be ambiguous.")
                        duplicate_names.add(name)
                else:
                    name_to_ea[name] = ea
        else:
            print(f"aiDAPal Error: [async_call] Cannot map results, ea_to_original_name is not a dictionary (type: {type(ea_to_original_name)}).")
            results_list_with_names = None # Prevent further processing

        mapped_count = 0
        unmapped_count = 0
        # Check again if results list is still valid after potential name map failure
        if results_list_with_names is not None:
            for analysis in results_list_with_names:
                original_name = analysis.get("original_function_name")
                if not original_name:
                    print(f"aiDAPal Warning: [async_call] Skipping analysis result missing 'original_function_name': {analysis.get('function_name', 'N/A')}")
                    unmapped_count += 1
                    continue

                found_ea = name_to_ea.get(original_name)
                if found_ea is not None:
                    # --- Sanity Check for Single-Function Mode ---
                    # If we requested single-function analysis, ensure the result matches the primary EA
                    if analyze_only_current and found_ea != primary_func_ea:
                        print(f"aiDAPal Warning: [async_call] Requested single-function analysis for {primary_func_ea_str} ('{primary_func_name_for_prompt}'), but AI returned analysis for '{original_name}' (0x{found_ea:X}). Discarding.")
                        unmapped_count += 1
                        continue
                    # --- End Sanity Check ---

                    analysis['function_ea'] = found_ea # Add the EA back to the dictionary
                    results_list_with_ea.append(analysis)
                    mapped_count += 1
                    # print(f"  Mapped '{original_name}' back to EA 0x{found_ea:X}") # Noisy
                else:
                    print(f"aiDAPal Warning: [async_call] Could not map analysis result for name '{original_name}' back to an EA. Discarding result.")
                    unmapped_count += 1

        print(f"aiDAPal DEBUG: [async_call] Mapping complete. Mapped: {mapped_count}, Unmapped/Discarded: {unmapped_count}")

    # --- Schedule UI Update ---
    # Pass the list containing dictionaries *with EAs*
    if results_list_with_ea: # Only show UI if we successfully mapped at least one result
        print(f"aiDAPal DEBUG: [async_call] Analysis successful (mapped {len(results_list_with_ea)} items) for {primary_func_ea_str}. Preparing UI update.")
        call_do_show_ui = partial(do_show_ui, results_list_with_ea, primary_func_ea)
        print(f"aiDAPal DEBUG: [async_call] Scheduling UI update via execute_ui_requests for {primary_func_ea_str}.")
        ida_kernwin.execute_ui_requests([call_do_show_ui])
        print(f"aiDAPal DEBUG: [async_call] UI update request submitted for {primary_func_ea_str}.")
    else:
        # Handle cases where AI failed, returned nothing, or mapping failed for all items
        print(f"aiDAPal DEBUG: [async_call] Analysis failed, returned no results, or failed to map results for {primary_func_ea_str}.")
        if analysis_exception:
             print(f"aiDAPal Error: [async_call] Analysis failed due to exception: {analysis_exception}")
        # Optionally show an info message if mapping failed but AI call succeeded
        elif results_list_with_names is not None and not results_list_with_ea:
             ida_kernwin.warning("aiDAPal: Received analysis from AI, but could not map results back to functions based on names, or the AI analyzed the wrong function. Check AI response format or function names.")


    print(f"aiDAPal DEBUG: [async_call] Exiting background thread analysis function for {primary_func_ea_str}.")
# --- END MODIFIED ASYNC CALL ---


# --- IDA Plugin Integration (Action Handler, Hooks, Plugin Class) ---
# These parts remain largely the same, as the trigger point (right-click on primary function)
# and the mechanism for adding menu items don't change.
# The Action Handler now passes the primary cfunc to the wrapper.

class MyActionHandler(ida_kernwin.action_handler_t):
    """Handles the activation of the context menu actions."""
    def __init__(self, model_name, analyze_only_current=False): # ADDED analyze_only_current flag
        self.model = model_name # Store the specific model for this handler instance
        self.analyze_only_current = analyze_only_current # Store the analysis type
        ida_kernwin.action_handler_t.__init__(self)
        analysis_type_str = "SINGLE" if analyze_only_current else "MULTI"
        print(f"aiDAPal DEBUG: [MyActionHandler.__init__] Action handler created for model '{self.model}', type: {analysis_type_str}")

    def activate(self, ctx):
        """Called when the context menu item is clicked."""
        analysis_type_str = "SINGLE" if self.analyze_only_current else "MULTI"
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Action '{self.model}' ({analysis_type_str}) activated.")
        widget = ctx.widget
        widget_type = ida_kernwin.get_widget_type(widget)
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Context widget type: {widget_type}")
        if widget_type != ida_kernwin.BWN_PSEUDOCODE:
            print("aiDAPal DEBUG: [MyActionHandler.activate] Action activated outside pseudocode view. Ignoring.")
            return 1 # Indicate action not handled

        print("aiDAPal DEBUG: [MyActionHandler.activate] Getting vdui from pseudocode widget...")
        vu = ida_hexrays.get_widget_vdui(widget)
        if not vu:
            print("aiDAPal Error: [MyActionHandler.activate] Could not get vdui for the pseudocode widget.")
            ida_kernwin.warning("aiDAPal: Could not get pseudocode view details (vdui).")
            return 1
        if not vu.cfunc: # Check if cfunc exists (decompilation successful)
            print("aiDAPal Error: [MyActionHandler.activate] vdui exists, but vu.cfunc is None. Decompilation might have failed.")
            ida_kernwin.warning("aiDAPal: Decompilation failed or not available for this function.")
            return 1
        print("aiDAPal DEBUG: [MyActionHandler.activate] Got vdui and cfunc.")

        # This is the PRIMARY function the user clicked on
        primary_cfunc = vu.cfunc
        primary_func_ea = primary_cfunc.entry_ea
        primary_func_ea_str = f"0x{primary_func_ea:X}"
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Target primary function EA: {primary_func_ea_str}")

        # --- Concurrency Check (using primary function EA) ---
        concurrency_key = primary_func_ea
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Acquiring lock to check analysis status for {primary_func_ea_str}...")
        with g_analysis_lock:
            print(f"aiDAPal DEBUG: [MyActionHandler.activate] Lock acquired. Checking if {primary_func_ea_str} is in {g_analysis_in_progress}")
            if concurrency_key in g_analysis_in_progress:
                print(f"aiDAPal DEBUG: [MyActionHandler.activate] Analysis already in progress for {primary_func_ea_str}. Showing warning.")
                ida_kernwin.warning(f"aiDAPal: Analysis already running for {ida_funcs.get_func_name(primary_func_ea)}. Please wait.")
                print(f"aiDAPal DEBUG: [MyActionHandler.activate] Releasing lock and exiting activate for {primary_func_ea_str}.")
                return 1 # Indicate action handled (by showing warning)
            # Mark as in progress using the primary EA
            g_analysis_in_progress.add(concurrency_key)
            print(f"aiDAPal DEBUG: [MyActionHandler.activate] Marked {primary_func_ea_str} as in progress. Current: {g_analysis_in_progress}")
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Lock released for {primary_func_ea_str}.")

        # --- Get Context (Data References for the PRIMARY function - RUN IN MAIN THREAD) ---
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Preparing to get data reference comments for primary function {primary_func_ea_str} in main thread...")
        dref_comments = None # Initialize actual result variable
        dref_result_container = [None] # Mutable container to hold the result

        try:
            # Define helper to run in main thread, modifying the container
            def get_dref_comments_main(container):
                print(f"aiDAPal DEBUG: [get_dref_comments_main] Running in main thread for {primary_func_ea_str}")
                try:
                    # Store the actual result (string or None) in the container
                    container[0] = get_function_data_ref_comments(primary_func_ea)
                    print(f"aiDAPal DEBUG: [get_dref_comments_main] Stored result in container: Type={type(container[0])}")
                    return 1 # Indicate success to execute_sync
                except Exception as e_inner:
                    print(f"aiDAPal Error: [get_dref_comments_main] Exception inside main thread execution: {e_inner}")
                    traceback.print_exc()
                    container[0] = None # Ensure container holds None on error
                    return 0 # Indicate failure to execute_sync

            # Execute synchronously, passing the container
            sync_status = ida_kernwin.execute_sync(
                lambda: get_dref_comments_main(dref_result_container),
                ida_kernwin.MFF_READ
            )

            print(f"aiDAPal DEBUG: [MyActionHandler.activate] execute_sync for dref comments returned status: {sync_status}")

            # Check the status code returned by execute_sync
            if sync_status == 1:
                # Success: retrieve the actual result from the container
                dref_comments = dref_result_container[0]
                print(f"aiDAPal DEBUG: [MyActionHandler.activate] Successfully retrieved dref comments from main thread. Type: {type(dref_comments)}")
                if dref_comments is not None: # Check if it's not None before using len()
                     print(f"aiDAPal DEBUG: [MyActionHandler.activate] Found data reference comments for {primary_func_ea_str} (length: {len(dref_comments)}).")
                else:
                     print(f"aiDAPal DEBUG: [MyActionHandler.activate] No data reference comments found or returned as None for {primary_func_ea_str}.")
            else:
                # Failure or Cancelled
                print(f"aiDAPal Error: [MyActionHandler.activate] execute_sync for dref comments failed or was cancelled (status: {sync_status}).")
                with g_analysis_lock: g_analysis_in_progress.discard(concurrency_key)
                return 1 # Indicate error occurred

        except Exception as e_dref_sync:
             print(f"aiDAPal Error: [MyActionHandler.activate] Exception occurred during execute_sync call for dref comments: {e_dref_sync}")
             traceback.print_exc()
             with g_analysis_lock: g_analysis_in_progress.discard(concurrency_key)
             return 1 # Indicate error occurred

        # --- Start Analysis Thread ---
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Preparing analysis thread ({analysis_type_str}) for {primary_func_ea_str} using model '{self.model}'.")
        caller = partial(analysis_task_wrapper, primary_cfunc, self.model, self.analyze_only_current, dref_comments)
        analysis_thread = threading.Thread(target=caller)
        analysis_thread.start()
        print(f"aiDAPal DEBUG: [MyActionHandler.activate] Analysis thread started for {primary_func_ea_str}, model '{self.model}'. Exiting activate.")
        return 1 # Indicate action was initiated

    def update(self, ctx):
        """Enables the action only when in a pseudocode view with successful decompilation."""
        if ctx.widget_type == ida_kernwin.BWN_PSEUDOCODE:
            vu = ida_hexrays.get_widget_vdui(ctx.widget)
            if vu and vu.cfunc:
                 return ida_kernwin.AST_ENABLE_FOR_WIDGET
            else:
                 return ida_kernwin.AST_DISABLE_FOR_WIDGET
        else:
            return ida_kernwin.AST_DISABLE_FOR_WIDGET


class MyHooks(ida_kernwin.UI_Hooks):
    """Hooks into IDA's UI to add context menu items."""
    def finish_populating_widget_popup(self, widget, popup_handle, ctx=None):
        """Adds menu items to the widget's context menu."""
        widget_type = ida_kernwin.get_widget_type(widget)
        if widget_type == ida_kernwin.BWN_PSEUDOCODE:
            # print(f"aiDAPal DEBUG: [MyHooks.finish_populating_widget_popup] Adding menu items to Pseudocode popup.") # Noisy
            try:
                submenu_name = "aiDAPal Analysis"
                action_added = False
                for model_name in MODELS_TO_REGISTER:
                    # --- Add MULTI function analysis action ---
                    action_name_multi = f"aidapal:googleai:multi:{model_name}"
                    # print(f"aiDAPal DEBUG: [MyHooks.finish_populating_widget_popup] Attaching action '{action_name_multi}' to submenu '{submenu_name}'...") # Noisy
                    if ida_kernwin.attach_action_to_popup(widget, popup_handle, action_name_multi, f"{submenu_name}/", ida_kernwin.SETMENU_INS):
                        action_added = True
                        # print(f"aiDAPal DEBUG: [MyHooks.finish_populating_widget_popup] Successfully attached action '{action_name_multi}'.") # Noisy
                    else:
                         print(f"aiDAPal Warning: [MyHooks.finish_populating_widget_popup] Failed to attach action '{action_name_multi}' to popup.")

                    # --- Add SINGLE function analysis action ---
                    action_name_single = f"aidapal:googleai:single:{model_name}"
                    # print(f"aiDAPal DEBUG: [MyHooks.finish_populating_widget_popup] Attaching action '{action_name_single}' to submenu '{submenu_name}'...") # Noisy
                    if ida_kernwin.attach_action_to_popup(widget, popup_handle, action_name_single, f"{submenu_name}/", ida_kernwin.SETMENU_INS):
                        action_added = True
                        # print(f"aiDAPal DEBUG: [MyHooks.finish_populating_widget_popup] Successfully attached action '{action_name_single}'.") # Noisy
                    else:
                         print(f"aiDAPal Warning: [MyHooks.finish_populating_widget_popup] Failed to attach action '{action_name_single}' to popup.")


                if action_added:
                    # print(f"aiDAPal DEBUG: [MyHooks.finish_populating_widget_popup] Adding separators to submenu '{submenu_name}'.") # Noisy
                    ida_kernwin.attach_action_to_popup(widget, popup_handle, "-", f"{submenu_name}/", ida_kernwin.SETMENU_INS | ida_kernwin.SETMENU_FIRST)
                # else: print(f"aiDAPal DEBUG: [MyHooks.finish_populating_widget_popup] No actions were added, skipping separators.") # Noisy

            except Exception as e:
                print(f"aiDAPal Error: [MyHooks.finish_populating_widget_popup] Exception attaching actions to popup: {e}")
                traceback.print_exc()


class aidapal_t(idaapi.plugin_t):
    """The main IDA Pro plugin class."""
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "aiDAPal: Google AI assistance for Hex-Rays (Multi/Single Function Analysis - google-genai SDK)" # Updated comment
    help = "Right-click in Pseudocode view to analyze the function (and optionally callees) with Google AI." # Updated help
    wanted_name = "aiDAPal (Google AI - Multi/Single Func)" # Updated name
    wanted_hotkey = ""
    hooks = None # Class attribute to hold hook instance
    registered_actions = [] # Class attribute

    def init(self):
        """Called by IDA when loading the plugin."""
        print("-" * 60)
        print(f"aiDAPal DEBUG: {self.wanted_name} plugin initializing...")

        # Check prerequisites (updated for google-genai)
        print("aiDAPal DEBUG: [init] Checking prerequisites...")
        # Check if GOOGLE_API_KEY environment variable is set or hardcoded key is valid
        # env_key = os.environ.get("GOOGLE_API_KEY")
        # hardcoded_key = GOOGLE_AI_API_KEY
        # if not (env_key or (hardcoded_key and hardcoded_key != "YOUR_API_KEY_HERE")):
        if not GOOGLE_AI_API_KEY or GOOGLE_AI_API_KEY == "YOUR_API_KEY_HERE":
             print("aiDAPal Warning: [init] GOOGLE_API_KEY environment variable not set or GOOGLE_AI_API_KEY is placeholder in script.")
             print("  Plugin actions will not function until the key is set (either via env var GOOGLE_API_KEY or in script) and IDA is restarted.")

        print("aiDAPal DEBUG: [init] Initializing Hex-Rays...")
        if not ida_hexrays.init_hexrays_plugin():
            print("aiDAPal Error: [init] Hex-Rays decompiler is not available. Plugin cannot run.")
            print("-" * 60)
            return idaapi.PLUGIN_SKIP
        print("aiDAPal DEBUG: [init] Hex-Rays initialized successfully.")

        aidapal_t.registered_actions = []
        print("aiDAPal DEBUG: [init] Cleared registered actions list.")

        try:
            # Register actions (tooltips remain similar, ensure model names are valid for google-genai)
            print(f"aiDAPal DEBUG: [init] Registering actions for models: {MODELS_TO_REGISTER}")
            for model_name in MODELS_TO_REGISTER:
                # --- Register MULTI function action ---
                action_name_multi = f"aidapal:googleai:multi:{model_name}"
                action_label_multi = f'Analyze Func & Callees with {model_name}' # Original label
                action_tooltip_multi = f'Send current function and callees (depth 2) to Google AI ({model_name}) for analysis of ALL functions' # Original tooltip
                print(f"aiDAPal DEBUG: [init] Creating action description for '{action_name_multi}'...")
                action_desc_multi = ida_kernwin.action_desc_t(
                    action_name_multi, action_label_multi, MyActionHandler(model_name, analyze_only_current=False), None, action_tooltip_multi, 199) # analyze_only_current=False

                print(f"aiDAPal DEBUG: [init] Registering action '{action_name_multi}'...")
                if ida_kernwin.register_action(action_desc_multi):
                    print(f"  aiDAPal DEBUG: [init] Successfully registered action: '{action_name_multi}'")
                    aidapal_t.registered_actions.append(action_name_multi)
                else:
                    print(f"  aiDAPal Error: [init] Failed to register action: '{action_name_multi}'")

                # --- Register SINGLE function action ---
                action_name_single = f"aidapal:googleai:single:{model_name}"
                action_label_single = f'Analyze Current Func Only with {model_name}' # New label
                action_tooltip_single = f'Send current function and callees (depth 2) to Google AI ({model_name}) for analysis of ONLY the current function' # New tooltip
                print(f"aiDAPal DEBUG: [init] Creating action description for '{action_name_single}'...")
                action_desc_single = ida_kernwin.action_desc_t(
                    action_name_single, action_label_single, MyActionHandler(model_name, analyze_only_current=True), None, action_tooltip_single, 199) # analyze_only_current=True

                print(f"aiDAPal DEBUG: [init] Registering action '{action_name_single}'...")
                if ida_kernwin.register_action(action_desc_single):
                    print(f"  aiDAPal DEBUG: [init] Successfully registered action: '{action_name_single}'")
                    aidapal_t.registered_actions.append(action_name_single)
                else:
                    print(f"  aiDAPal Error: [init] Failed to register action: '{action_name_single}'")


            if not aidapal_t.registered_actions:
                print("aiDAPal Error: [init] No actions were registered successfully. Skipping plugin load.")
                print("-" * 60)
                return idaapi.PLUGIN_SKIP
            print(f"aiDAPal DEBUG: [init] Total actions registered: {len(aidapal_t.registered_actions)}")

            # Install UI hooks (unchanged)
            print("aiDAPal DEBUG: [init] Installing UI hooks...")
            aidapal_t.hooks = MyHooks()
            if aidapal_t.hooks.hook():
                 print("  aiDAPal DEBUG: [init] UI hooks installed successfully.")
            else:
                 print("  aiDAPal Error: [init] Failed to install UI hooks. Unregistering actions and skipping plugin load.")
                 self._unregister_actions()
                 aidapal_t.hooks = None
                 print("-" * 60)
                 return idaapi.PLUGIN_SKIP

        except Exception as e:
            print(f"aiDAPal Error: [init] Exception during initialization: {e}")
            traceback.print_exc()
            print("aiDAPal DEBUG: [init] Cleaning up due to exception...")
            self._unregister_actions()
            if aidapal_t.hooks:
                try: aidapal_t.hooks.unhook()
                except Exception as unhook_e: print(f"aiDAPal Error: [init] Exception during unhooking after error: {unhook_e}")
                aidapal_t.hooks = None
            print("-" * 60)
            return idaapi.PLUGIN_SKIP

        print(f"aiDAPal DEBUG: {self.wanted_name} plugin initialization complete.")
        print("-" * 60)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """Called by IDA when running the plugin from the menu (not used here)."""
        print(f"aiDAPal DEBUG: {self.wanted_name} run() called (arg={arg}). This plugin uses the context menu.")
        ida_kernwin.info(f"{self.wanted_name}: Use the right-click context menu in the Pseudocode view to analyze functions.")

    def term(self):
        """Called by IDA when unloading the plugin."""
        print("-" * 60)
        print(f"aiDAPal DEBUG: {self.wanted_name} plugin terminating...")

        # Unhook UI hooks (unchanged)
        print("aiDAPal DEBUG: [term] Uninstalling UI hooks...")
        try:
            if aidapal_t.hooks:
                aidapal_t.hooks.unhook()
                print("  aiDAPal DEBUG: [term] UI hooks uninstalled successfully.")
                aidapal_t.hooks = None
            else:
                print("  aiDAPal DEBUG: [term] No UI hooks instance to uninstall.")
        except Exception as e:
            print(f"  aiDAPal Error: [term] Exception during unhooking: {e}")
            traceback.print_exc()

        # Unregister actions (unchanged)
        print("aiDAPal DEBUG: [term] Unregistering actions...")
        self._unregister_actions()

        # Clean up any open UI forms (unchanged logic, uses primary EA key)
        print(f"aiDAPal DEBUG: [term] Closing any tracked UI forms ({len(aiDAPalUI.open_forms)} found)...")
        forms_to_close = list(aiDAPalUI.open_forms.items())
        for func_ea, form_instance in forms_to_close:
            func_ea_str = f"0x{func_ea:X}"
            print(f"  aiDAPal DEBUG: [term] Attempting to close tracked UI form triggered by {func_ea_str}")
            try:
                form_instance.Close(0)
                print(f"  aiDAPal DEBUG: [term] Close() called for form {func_ea_str}.")
            except Exception as e:
                print(f"  aiDAPal Error: [term] Exception closing form for {func_ea_str}: {e}")
                aiDAPalUI.open_forms.pop(func_ea, None) # Manual removal on error

        aiDAPalUI.open_forms.clear()
        print("aiDAPal DEBUG: [term] Cleared open_forms tracking dictionary.")

        # Clear concurrency state (unchanged)
        print("aiDAPal DEBUG: [term] Clearing concurrency state...")
        with g_analysis_lock:
            g_analysis_in_progress.clear()
        print("aiDAPal DEBUG: [term] Concurrency state cleared.")


        print(f"aiDAPal DEBUG: {self.wanted_name} plugin termination complete.")
        print("-" * 60)

    def _unregister_actions(self):
        """Helper to unregister all actions this plugin registered."""
        print(f"aiDAPal DEBUG: [_unregister_actions] Unregistering {len(aidapal_t.registered_actions)} actions...")
        if hasattr(aidapal_t, 'registered_actions'):
            actions_to_unregister = list(aidapal_t.registered_actions)
            for action_name in actions_to_unregister:
                try:
                    # print(f"  aiDAPal DEBUG: [_unregister_actions] Unregistering action: '{action_name}'...") # Noisy
                    if ida_kernwin.unregister_action(action_name):
                        # print(f"    aiDAPal DEBUG: [_unregister_actions] Successfully unregistered '{action_name}'.") # Noisy
                        pass
                    else:
                        print(f"    aiDAPal Warning: [_unregister_actions] Failed to unregister action '{action_name}' (API returned False).")
                except Exception as e:
                     print(f"    aiDAPal Error: [_unregister_actions] Exception unregistering action '{action_name}': {e}")
            aidapal_t.registered_actions = []
            print("aiDAPal DEBUG: [_unregister_actions] Cleared registered actions list.")
        else:
            print("aiDAPal DEBUG: [_unregister_actions] No 'registered_actions' attribute found.")


# --- Plugin Entry Point ---

def PLUGIN_ENTRY():
    """Required entry point for IDA Pro plugins."""
    print("aiDAPal DEBUG: PLUGIN_ENTRY() called. Returning aidapal_t instance.")
    return aidapal_t()

# --- End of Script ---
print("aiDAPal DEBUG: Script loaded.")


