# plugins

repository: beerandgin/InitialAnalysis
file: InitialAnalysis.py
description: IDA Script for Initial Analysis from Several Sources Turned Into A Plugin
wanted_name: InitialAnalysis
comment: Wow, such comment.
created_at: 2019-03-08 18:22:36+00:00
pushed_at: 2018-03-07 03:56:23+00:00
url:
https://github.com/beerandgin/InitialAnalysis/blob/5a99bdc8089883954bf740209bd398a81f69a203/InitialAnalysis.py

content:
#!/usr/bin/env python
#
# SetColorSiko by Mike Sikorski
# https://practicalmalwareanalysis.com/setcolorssiko-py/
#
# Func-renamer by Alexander Hanel
# https://bitbucket.org/Alexander_Hanel/func-renamer
#
# ============================================================================
# Copyright (c) 2012, Sebastian Eschweiler <advanced(dot)malware<dot>analyst[at]gmail.com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the <organization> nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# =============================================================================

import idaapi
from idautils import *
from idc import *

DEFCOLOR = 0xffffffff
COLOR_CRYPTO = 0xffd2f8
COLOR_CALL = 0xffffd0


class InitialAnalysis_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Wow, such comment."
    help = "Wow, much help."
    wanted_name = "InitialAnalysis"
    wanted_hotkey = "F3"

    def init(self):
        print "Initial Analysis Plugin Loaded"
        return idaapi.PLUGIN_OK

    def run(self, arg):
        InitialAnalysis_main()

    def term(self):
        pass


def PLUGIN_ENTRY():
    return InitialAnalysis_t()


def findUnidentifiedFunctions():
    # just get all not-function code and convert it to functions
    next = idaapi.cvar.inf.minEA
    while next != idaapi.BADADDR:
        next = idaapi.find_not_func(next, 1)
        flags = idaapi.getFlags(next)
        if idaapi.isCode(flags):
            idc.MakeFunction(next)


def colorize(addr, color):
    idaapi.set_item_color(addr, color)


def revokeAnalysis():
    n = idaapi.netnode("$ initialAnalysis", 0, False)

    if (n == idaapi.BADNODE):
        return

    idx = n.alt1st()
    while idx != idaapi.BADNODE:
        colorize(idx, DEFCOLOR)
        idx = n.altnxt(idx)

    n.kill()

    idaapi.refresh_idaview_anyway()


def setEaInfo(ea, info=""):
    n = idaapi.netnode("$ initialAnalysis", 0, True)
    n.set(info)


def setFunctionInfo(ea, color, info=""):
    f = idaapi.get_func(ea)
    if not f:
        return
    setEaInfo(f.startEA, info)


def setInfo(ea, color, info=""):
    colorize(ea, color)
    setEaInfo(ea, info)
    setFunctionInfo(ea, info)


class CryptoTester(object):
    def __init__(self):
        self.xor_instructions = [
            idaapi.NN_xor, idaapi.NN_pxor, idaapi.NN_xorps, idaapi.NN_xorpd]
        self.other_crypto = [idaapi.NN_ror, idaapi.NN_rol, idaapi.NN_not]

    def instruction(self, cmd):

        colorize = False

        if cmd.itype in self.xor_instructions:
            # check if different operands
            if cmd.Op1.type != cmd.Op2.type or cmd.Op1.reg != cmd.Op2.reg or cmd.Op1.value != cmd.Op2.value:
                colorize = True

        elif cmd.itype in self.other_crypto:
            colorize = True

        if colorize:
            setInfo(cmd.ea, COLOR_CRYPTO, "crypto")


class CallTester(object):
    def __init__(self):
        self.call_instructions = [idaapi.NN_call,
                                  idaapi.NN_callfi, idaapi.NN_callni]

    def instruction(self, cmd):
        if cmd.itype in self.call_instructions:
            setInfo(cmd.ea, COLOR_CALL)


def iterateInstructions():
    next = 0
    while next != idaapi.BADADDR:

        # get next instruction
        next = idc.NextHead(next)

        idaapi.decode_insn(next)
        for handlers in InstructionCallbacks:
            handlers.instruction(idaapi.cmd)

# SetColorSiko 
# Mike Sikorski
# https://practicalmalwareanalysis.com/setcolorssiko-py/

def setcolorsiko():
    '''
    #Color the Calls off-white
    heads = Heads(SegStart(ScreenEA()), SegEnd(ScreenEA()))
    funcCalls = []
    for i in heads:
            if GetMnem(i) == "call":
                    funcCalls.append(i)
    print "Number of calls: %d" % (len(funcCalls))
    for i in funcCalls:
            SetColor(i, CIC_ITEM, 0xc7fdff)
    '''
    # Color Anti-VM instructions Red and print their location
    heads = Heads(SegStart(ScreenEA()), SegEnd(ScreenEA()))
    antiVM = []
    for i in heads:
        if (GetMnem(i) == "sidt" or GetMnem(i) == "sgdt" or GetMnem(i) == "sldt" or GetMnem(i) == "smsw" or GetMnem(i) == "str" or GetMnem(i) == "in" or GetMnem(i) == "cpuid"):
            antiVM.append(i)
    print "Number of potential Anti-VM instructions: %d" % (len(antiVM))
    for i in antiVM:
        print "Anti-VM potential at %x" % i
        SetColor(i, CIC_ITEM, 0x0000ff)
    # Color non-zeroing out xor instructions Orange
    heads = Heads(SegStart(ScreenEA()), SegEnd(ScreenEA()))
    xor = []
    for i in heads:
        if GetMnem(i) == "xor":
            if (GetOpnd(i, 0) != GetOpnd(i, 1)):
                xor.append(i)
    print "Number of xor: %d" % (len(xor))
    for i in xor:
        SetColor(i, CIC_ITEM, 0x00a5ff)

# func-rename
# alexander<dot>hanel<at>gmail<dot>com
# https://bitbucket.org/Alexander_Hanel/func-renamer

class GENAPI:
    def __init__(self):
        self.rename = False
        self.comment = False
        self.reg = [ 'Reg', 'RegCloseKey' ,'RegConnectRegistryA' ,'RegConnectRegistryW' ,'RegCreateKeyA' ,'RegCreateKeyExA' ,'RegCreateKeyExW',\
        'RegCreateKeyW' ,'RegDeleteKeyA' ,'RegDeleteKeyW' ,'RegDeleteValueA' ,'RegDeleteValueW' ,'RegDisablePredefinedCache' ,\
        'RegDisablePredefinedCacheEx' ,'RegEnumKeyA' ,'RegEnumKeyExA' ,'RegEnumKeyExW' ,'RegEnumKeyW' ,'RegEnumValueA' ,\
        'RegEnumValueW' ,'RegFlushKey' ,'RegGetKeySecurity' ,'RegLoadKeyA' ,'RegLoadKeyW' ,'RegNotifyChangeKeyValue' ,\
        'RegOpenCurrentUser' ,'RegOpenKeyA' ,'RegOpenKeyExA' ,'RegOpenKeyExW' ,'RegOpenKeyW' ,'RegOpenUserClassesRoot' ,\
        'RegOverridePredefKey' ,'RegQueryInfoKeyA' ,'RegQueryInfoKeyW' ,'RegQueryMultipleValuesA' ,'RegQueryMultipleValuesW' ,\
        'RegQueryValueA' ,'RegQueryValueExA' ,'RegQueryValueExW' ,'RegQueryValueW' ,'RegReplaceKeyA' ,'RegReplaceKeyW' ,\
        'RegRestoreKeyA' ,'RegRestoreKeyW' ,'RegSaveKeyA' ,'RegSaveKeyExA' ,'RegSaveKeyExW' ,'RegSaveKeyW' ,'RegSetKeySecurity' ,\
        'RegSetValueA' ,'RegSetValueExA' ,'RegSetValueExW' ,'RegSetValueW' ,'RegUnLoadKeyA' ,'RegUnLoadKeyW', 'SHDeleteEmptyKeyA' ,\
        'SHDeleteEmptyKeyW' ,'SHDeleteKeyA' ,'SHDeleteKeyW' ,'SHOpenRegStream2A' ,'SHOpenRegStream2W' ,'SHOpenRegStreamA' ,\
        'SHOpenRegStreamW' ,'SHQueryInfoKeyA' ,'SHQueryInfoKeyW' ,'SHQueryValueExA' ,'SHQueryValueExW' ,'SHRegCloseUSKey' ,\
        'SHRegCreateUSKeyA' ,'SHRegCreateUSKeyW' ,'SHRegDeleteEmptyUSKeyA' ,'SHRegDeleteEmptyUSKeyW' ,'SHRegDeleteUSValueA' ,\
        'SHRegDeleteUSValueW' ,'SHRegDuplicateHKey' ,'SHRegEnumUSKeyA' ,'SHRegEnumUSKeyW' ,'SHRegEnumUSValueA' ,'SHRegEnumUSValueW',\
        'SHRegGetBoolUSValueA' ,'SHRegGetBoolUSValueW' ,'SHRegGetPathA' ,'SHRegGetPathW' ,'SHRegGetUSValueA' ,'SHRegGetUSValueW' ,\
        'SHRegGetValueA' ,'SHRegGetValueW' ,'SHRegOpenUSKeyA' ,'SHRegOpenUSKeyW' ,'SHRegQueryInfoUSKeyA' ,'SHRegQueryInfoUSKeyW' ,\
        'SHRegQueryUSValueA' ,'SHRegQueryUSValueW' ,'SHRegSetPathA' ,'SHRegSetPathW' ,'SHRegSetUSValueA' ,'SHRegSetUSValueW' ,\
        'SHRegWriteUSValueA' ,'SHRegWriteUSValueW' ,'SHDeleteOrphanKeyA' ,'SHDeleteOrphanKeyW' ,'SHDeleteValueA' ,'SHDeleteValueW' ,\
        'SHEnumKeyExA' ,'SHEnumKeyExW' ,'SHEnumValueA' ,'SHEnumValueW' ,'SHGetValueA' ,'SHGetValueW' ,'SHOpenRegStream2A' ,\
        'SHOpenRegStream2W' ,'SHOpenRegStreamA' ,'SHOpenRegStreamW' ,'SHQueryInfoKeyA' ,'SHQueryInfoKeyW' ,'SHQueryValueExA' ,\
        'SHQueryValueExW' ,'SHRegCloseUSKey' ,'SHRegCreateUSKeyA' ,'SHRegCreateUSKeyW' ,'SHRegDeleteEmptyUSKeyA' ,\
        'SHRegDeleteEmptyUSKeyW' ,'SHRegDeleteUSValueA' ,'SHRegDeleteUSValueW' ,'SHRegDuplicateHKey' ,'SHRegEnumUSKeyA' ,\
        'SHRegEnumUSKeyW' ,'SHRegEnumUSValueA' ,'SHRegEnumUSValueW' ,'SHRegGetBoolUSValueA' ,'SHRegGetBoolUSValueW' ,'SHRegGetPathA' ,\
        'SHRegGetPathW' ,'SHRegGetUSValueA' ,'SHRegGetUSValueW' ,'SHRegGetValueA' ,'SHRegGetValueW' ,'SHRegOpenUSKeyA' ,'SHRegOpenUSKeyW' ,\
        'SHRegQueryInfoUSKeyA' ,'SHRegQueryInfoUSKeyW' ,'SHRegQueryUSValueA' ,'SHRegQueryUSValueW' ,'SHRegSetPathA' ,'SHRegSetPathW' ,\
        'SHRegSetUSValueA' ,'SHRegSetUSValueW' ,'SHRegWriteUSValueA' ,'SHRegWriteUSValueW']

        self.winsock = [ 'Ws2', 'FreeAddrInfoW', 'GetAddrInfoW', 'GetNameInfoW', 'WEP', 'WPUCompleteOverlappedRequest', 'WSAAccept', \
        'WSAAddressToStringA', 'WSAAddressToStringW', 'WSAAsyncGetHostByAddr', 'WSAAsyncGetHostByName', 'WSAAsyncGetProtoByName',\
        'WSAAsyncGetProtoByNumber', 'WSAAsyncGetServByName', 'WSAAsyncGetServByPort', 'WSAAsyncSelect', 'WSACancelAsyncRequest',\
        'WSACancelBlockingCall', 'WSACleanup', 'WSACloseEvent', 'WSAConnect', 'WSACreateEvent', 'WSADuplicateSocketA',\
        'WSADuplicateSocketW', 'WSAEnumNameSpaceProvidersA', 'WSAEnumNameSpaceProvidersW', 'WSAEnumNetworkEvents', 'WSAEnumProtocolsA',\
        'WSAEnumProtocolsW', 'WSAEventSelect', 'WSAGetLastError', 'WSAGetOverlappedResult', 'WSAGetQOSByName', \
        'WSAGetServiceClassInfoA', 'WSAGetServiceClassInfoW', 'WSAGetServiceClassNameByClassIdA', 'WSAGetServiceClassNameByClassIdW',\
        'WSAHtonl', 'WSAHtons', 'WSAInstallServiceClassA', 'WSAInstallServiceClassW', 'WSAIoctl', 'WSAIsBlocking', 'WSAJoinLeaf', \
        'WSALookupServiceBeginA', 'WSALookupServiceBeginW', 'WSALookupServiceEnd', 'WSALookupServiceNextA', 'WSALookupServiceNextW', \
        'WSANSPIoctl', 'WSANtohl', 'WSANtohs', 'WSAProviderConfigChange', 'WSARecv', 'WSARecvDisconnect', 'WSARecvFrom', \
        'WSARemoveServiceClass', 'WSAResetEvent', 'WSASend', 'WSASendDisconnect', 'WSASendTo', 'WSASetBlockingHook', 'WSASetEvent',\
        'WSASetLastError', 'WSASetServiceA', 'WSASetServiceW', 'WSASocketA', 'WSASocketW', 'WSAStartup', 'WSAStringToAddressA', \
        'WSAStringToAddressW', 'WSAUnhookBlockingHook', 'WSAWaitForMultipleEvents', 'WSApSetPostRoutine', 'WSCDeinstallProvider', \
        'WSCEnableNSProvider', 'WSCEnumProtocols', 'WSCGetProviderPath', 'WSCInstallNameSpace', 'WSCInstallProvider', 'WSCUnInstallNameSpace',\
        'WSCUpdateProvider', 'WSCWriteNameSpaceOrder', 'WSCWriteProviderOrder', '__WSAFDIsSet', 'accept', 'bind', 'closesocket', 'connect', \
        'freeaddrinfo', 'getaddrinfo', 'gethostbyaddr', 'gethostbyname', 'gethostname', 'getnameinfo', 'getpeername', 'getprotobyname', \
        'getprotobynumber', 'getservbyname', 'getservbyport', 'getsockname', 'getsockopt', 'htonl', 'htons', 'inet_addr', 'inet_ntoa', \
        'ioctlsocket', 'listen', 'ntohl', 'ntohs', 'recv', 'recvfrom', 'select', 'send', 'sendto', 'setsockopt', 'shutdown', 'socket']

        self.WinINet = [ 'WINet', 'CreateMD5SSOHash', 'DetectAutoProxyUrl', 'DllInstall', 'ForceNexusLookup', 'ForceNexusLookupExW', 'InternetAlgIdToStringA',\
        'InternetAlgIdToStringW', 'InternetAttemptConnect', 'InternetAutodial', 'InternetAutodialCallback', 'InternetAutodialHangup',\
        'InternetCanonicalizeUrlA', 'InternetCanonicalizeUrlW', 'InternetCheckConnectionA', 'InternetCheckConnectionW', \
        'InternetClearAllPerSiteCookieDecisions', 'InternetCloseHandle', 'InternetCombineUrlA', 'InternetCombineUrlW', \
        'InternetConfirmZoneCrossing', 'InternetConfirmZoneCrossingA', 'InternetConfirmZoneCrossingW', 'InternetConnectA',\
        'InternetConnectW', 'InternetCrackUrlA', 'InternetCrackUrlW', 'InternetCreateUrlA', 'InternetCreateUrlW', 'InternetDial',\
        'InternetDialA', 'InternetDialW', 'InternetEnumPerSiteCookieDecisionA', 'InternetEnumPerSiteCookieDecisionW', 'InternetErrorDlg',\
        'InternetFindNextFileA', 'InternetFindNextFileW', 'InternetFortezzaCommand', 'InternetGetCertByURL', 'InternetGetCertByURLA',\
        'InternetGetConnectedState', 'InternetGetConnectedStateEx', 'InternetGetConnectedStateExA', 'InternetGetConnectedStateExW',\
        'InternetGetCookieA', 'InternetGetCookieExA', 'InternetGetCookieExW', 'InternetGetCookieW', 'InternetGetLastResponseInfoA', \
        'InternetGetLastResponseInfoW', 'InternetGetPerSiteCookieDecisionA', 'InternetGetPerSiteCookieDecisionW', 'InternetGoOnline',\
        'InternetGoOnlineA', 'InternetGoOnlineW', 'InternetHangUp', 'InternetInitializeAutoProxyDll', 'InternetLockRequestFile',\
        'InternetOpenA', 'InternetOpenUrlA', 'InternetOpenUrlW', 'InternetOpenW', 'InternetQueryDataAvailable', 'InternetQueryFortezzaStatus',\
        'InternetQueryOptionA', 'InternetQueryOptionW', 'InternetReadFile', 'InternetReadFileExA', 'InternetReadFileExW', \
        'InternetSecurityProtocolToStringA', 'InternetSecurityProtocolToStringW', 'InternetSetCookieA', 'InternetSetCookieExA', \
        'InternetSetCookieExW', 'InternetSetCookieW', 'InternetSetDialState', 'InternetSetDialStateA', 'InternetSetDialStateW',\
        'InternetSetFilePointer', 'InternetSetOptionA', 'InternetSetOptionExA', 'InternetSetOptionExW', 'InternetSetOptionW', \
        'InternetSetPerSiteCookieDecisionA', 'InternetSetPerSiteCookieDecisionW', 'InternetSetStatusCallback', 'InternetSetStatusCallbackA',\
        'InternetSetStatusCallbackW', 'InternetShowSecurityInfoByURL', 'InternetShowSecurityInfoByURLA', 'InternetShowSecurityInfoByURLW', \
        'InternetTimeFromSystemTime', 'InternetTimeFromSystemTimeA', 'InternetTimeFromSystemTimeW', 'InternetTimeToSystemTime',\
        'InternetTimeToSystemTimeA', 'InternetTimeToSystemTimeW', 'InternetUnlockRequestFile', 'InternetWriteFile', 'InternetWriteFileExA',\
        'InternetWriteFileExW', 'IsHostInProxyBypassList', 'ParseX509EncodedCertificateForListBoxEntry', 'PrivacyGetZonePreferenceW', \
        'PrivacySetZonePreferenceW', 'ResumeSuspendedDownload', 'ShowCertificate', 'ShowClientAuthCerts', 'ShowSecurityInfo', \
        'ShowX509EncodedCertificate','UrlZonesDetach', '_GetFileExtensionFromUrl']

        self.cache = [ 'Cach','CommitUrlCacheEntryA', 'CommitUrlCacheEntryW', 'CreateUrlCacheContainerA', 'CreateUrlCacheContainerW', 'CreateUrlCacheEntryA',\
        'CreateUrlCacheEntryW', 'CreateUrlCacheGroup', 'DeleteIE3Cache', 'DeleteUrlCacheContainerA', 'DeleteUrlCacheContainerW', \
        'DeleteUrlCacheEntry', 'DeleteUrlCacheEntryA', 'DeleteUrlCacheEntryW', 'DeleteUrlCacheGroup', 'FindCloseUrlCache', 'FindFirstUrlCacheContainerA',\
        'FindFirstUrlCacheContainerW', 'FindFirstUrlCacheEntryA', 'FindFirstUrlCacheEntryExA', 'FindFirstUrlCacheEntryExW', 'FindFirstUrlCacheEntryW', \
        'FindFirstUrlCacheGroup', 'FindNextUrlCacheContainerA', 'FindNextUrlCacheContainerW', 'FindNextUrlCacheEntryA', 'FindNextUrlCacheEntryExA',\
        'FindNextUrlCacheEntryExW', 'FindNextUrlCacheEntryW', 'FindNextUrlCacheGroup', 'FreeUrlCacheSpaceA', 'FreeUrlCacheSpaceW', 'GetUrlCacheConfigInfoA', \
        'GetUrlCacheConfigInfoW', 'GetUrlCacheEntryInfoA', 'GetUrlCacheEntryInfoExA', 'GetUrlCacheEntryInfoExW', 'GetUrlCacheEntryInfoW', \
        'GetUrlCacheGroupAttributeA', 'GetUrlCacheGroupAttributeW', 'GetUrlCacheHeaderData', 'IncrementUrlCacheHeaderData', 'IsUrlCacheEntryExpiredA',\
        'IsUrlCacheEntryExpiredW', 'LoadUrlCacheContent', 'ReadUrlCacheEntryStream', 'RegisterUrlCacheNotification', 'RetrieveUrlCacheEntryFileA', \
        'RetrieveUrlCacheEntryFileW', 'RetrieveUrlCacheEntryStreamA', 'RetrieveUrlCacheEntryStreamW', 'RunOnceUrlCache', 'SetUrlCacheConfigInfoA',\
        'SetUrlCacheConfigInfoW', 'SetUrlCacheEntryGroup', 'SetUrlCacheEntryGroupA', 'SetUrlCacheEntryGroupW', 'SetUrlCacheEntryInfoA', 'SetUrlCacheEntryInfoW',\
        'SetUrlCacheGroupAttributeA', 'SetUrlCacheGroupAttributeW', 'SetUrlCacheHeaderData', 'UnlockUrlCacheEntryFile', 'UnlockUrlCacheEntryFileA', \
        'UnlockUrlCacheEntryFileW', 'UnlockUrlCacheEntryStream', 'UpdateUrlCacheContentPath']

        self.ftp = [ 'Ftp','FtpCommandA' ,'FtpCommandW' ,'FtpCreateDirectoryA' ,'FtpCreateDirectoryW' ,'FtpDeleteFileA' ,'FtpDeleteFileW' ,'FtpFindFirstFileA' ,\
        'FtpFindFirstFileW' ,'FtpGetCurrentDirectoryA' ,'FtpGetCurrentDirectoryW' ,'FtpGetFileA' ,'FtpGetFileEx' ,'FtpGetFileSize' ,'FtpGetFileW' ,\
        'FtpOpenFileA' ,'FtpOpenFileW' ,'FtpPutFileA' ,'FtpPutFileEx' ,'FtpPutFileW' ,'FtpRemoveDirectoryA' ,'FtpRemoveDirectoryW' ,'FtpRenameFileA' ,\
        'FtpRenameFileW' ,'FtpSetCurrentDirectoryA' ,'FtpSetCurrentDirectoryW']

        self.gopher = [ 'Gopher', 'GopherCreateLocatorA', 'GopherCreateLocatorW', 'GopherFindFirstFileA', 'GopherFindFirstFileW', 'GopherGetAttributeA', \
        'GopherGetAttributeW', 'GopherGetLocatorTypeA', 'GopherGetLocatorTypeW', 'GopherOpenFileA', 'GopherOpenFileW']

        self.url = ['Url', 'UrlApplySchemeA' ,'UrlApplySchemeW' ,'UrlCanonicalizeA' ,'UrlCanonicalizeW' ,'UrlCombineA' ,'UrlCombineW' ,'UrlCompareA' ,\
        'UrlCompareW' ,'UrlCreateFromPathA' ,'UrlCreateFromPathW' ,'UrlEscapeA' ,'UrlEscapeW' ,'UrlGetLocationA' ,'UrlGetLocationW' ,'UrlGetPartA'\
        ,'UrlGetPartW' ,'UrlHashA' ,'UrlHashW' ,'UrlIsA' ,'UrlIsNoHistoryA' ,'UrlIsNoHistoryW' ,'UrlIsOpaqueA' ,'UrlIsOpaqueW' ,'UrlIsW' ,'UrlUnescapeA' ,'UrlUnescapeW']

        self.dir = ['Dir','CreateDirectoryA', 'CreateDirectoryExA', 'CreateDirectoryExW', 'CreateDirectoryW', 'GetCurrentDirectoryA', 'GetCurrentDirectoryW',\
        'GetDllDirectoryA', 'GetDllDirectoryW', 'GetSystemDirectoryA', 'GetSystemDirectoryW', 'GetSystemWindowsDirectoryA', 'GetSystemWindowsDirectoryW',\
        'GetSystemWow64DirectoryA', 'GetSystemWow64DirectoryW', 'GetVDMCurrentDirectories', 'GetWindowsDirectoryA', 'GetWindowsDirectoryW', \
        'ReadDirectoryChangesW', 'RemoveDirectoryA', 'RemoveDirectoryW', 'SetCurrentDirectoryA', 'SetCurrentDirectoryW', 'SetDllDirectoryA',\
        'SetDllDirectoryW', 'SetVDMCurrentDirectories', 'SHCreateDirectory', 'SHCreateDirectoryExA', 'SHCreateDirectoryExW']

        # Mutex
        self.mutex = ['Mutx','CreateMutexA', 'CreateMutexW', 'OpenMutexA', 'OpenMutexW', 'ReleaseMutex']

        # Pipe
        self.pipe = [ 'Pipe', 'CallNamedPipeA', 'CallNamedPipeW', 'ConnectNamedPipe', 'CreateNamedPipeA', 'CreateNamedPipeW', 'CreatePipe', 'DisconnectNamedPipe',\
        'GetNamedPipeHandleStateA', 'GetNamedPipeHandleStateW', 'GetNamedPipeInfo', 'PeekNamedPipe', 'SetNamedPipeHandleState', 'TransactNamedPipe',\
        'WaitNamedPipeA', 'WaitNamedPipeW']

        # List of APIs related to HTTP from WinINet
        self.http = [ 'Http', 'HttpAddRequestHeadersA', 'HttpAddRequestHeadersW', 'HttpCheckDavCompliance', 'HttpEndRequestA', 'HttpEndRequestW',\
        'HttpOpenRequestA', 'HttpOpenRequestW', 'HttpQueryInfoA', 'HttpQueryInfoW', 'HttpSendRequestA', 'HttpSendRequestExA', \
        'HttpSendRequestExW', 'HttpSendRequestW' ]

        # enum process
        self.enum = [ 'Enum', 'CreateToolhelp32Snapshot', 'Process32First' ,'Process32FirstW' ,'Process32Next' ,'Process32NextW']

        # List of APIs related to hashing files
        self.hash = ['Hash', 'CryptCreateHash' ,'CryptDestroyHash' ,'CryptDuplicateHash' ,'CryptGetHashParam' ,'CryptHashData' ,'CryptHashSessionKey' ,\
        'CryptSetHashParam' ,'CryptSignHashA' ,'CryptSignHashW', 'FreeEncryptionCertificateHashList']

        # List of APIs related to Cryptograpy files
        self.crypt = ['Crypt', 'CryptAcquireContextA' ,'CryptAcquireContextW' ,'CryptContextAddRef' ,'CryptDecrypt' ,'CryptDeriveKey' ,'CryptDestroyKey' ,\
        'CryptDuplicateKey' ,'CryptEncrypt' ,'CryptEnumProviderTypesA' ,'CryptEnumProviderTypesW' ,'CryptEnumProvidersA' ,'CryptEnumProvidersW'\
        ,'CryptExportKey' ,'CryptGenKey' ,'CryptGenRandom' ,'CryptGetDefaultProviderA' ,'CryptGetDefaultProviderW' ,'CryptGetKeyParam' ,\
        'CryptGetProvParam' ,'CryptGetUserKey' ,'CryptImportKey' ,'CryptReleaseContext' ,'CryptSetKeyParam' ,'CryptSetProvParam' ,\
        'CryptSetProviderA' ,'CryptSetProviderExA' ,'CryptSetProviderExW' ,'CryptSetProviderW' ,'CryptVerifySignatureA' ,'CryptVerifySignatureW' ,\
        'DecryptFileA' ,'DecryptFileW', 'EncryptFileA' ,'EncryptFileW' ,'EncryptedFileKeyInfo' ,'EncryptionDisable', 'WriteEncryptedFileRaw', \
        'OpenEncryptedFileRawA' ,'OpenEncryptedFileRawW', 'DuplicateEncryptionInfoFile', 'SetUserFileEncryptionKey', 'ReadEncryptedFileRaw', \
        'RemoveUsersFromEncryptedFile', 'FileEncryptionStatusA', 'FileEncryptionStatusW', 'FreeEncryptedFileKeyInfo', 'CloseEncryptedFileRaw', \
        'AddUsersToEncryptedFile', 'QueryRecoveryAgentsOnEncryptedFile', 'QueryUsersOnEncryptedFile', 'ChainWlxLogoffEvent' ,'CryptAcquireContextU' ,\
        'CryptBinaryToStringA' ,'CryptBinaryToStringW' ,'CryptCloseAsyncHandle' ,'CryptCreateAsyncHandle' ,'CryptDecodeMessage' ,'CryptDecodeObject' ,\
        'CryptDecodeObjectEx' ,'CryptDecryptAndVerifyMessageSignature' ,'CryptDecryptMessage' ,'CryptEncodeObject' ,'CryptEncodeObjectEx' ,\
        'CryptEncryptMessage' ,'CryptEnumKeyIdentifierProperties' ,'CryptEnumOIDFunction' ,'CryptEnumOIDInfo' ,'CryptEnumProvidersU' ,'CryptExportPKCS8' ,\
        'CryptExportPublicKeyInfo' ,'CryptExportPublicKeyInfoEx' ,'CryptFindLocalizedName' ,'CryptFindOIDInfo' ,'CryptFormatObject' ,\
        'CryptFreeOIDFunctionAddress' ,'CryptGetAsyncParam' ,'CryptGetDefaultOIDDllList' ,'CryptGetDefaultOIDFunctionAddress' ,\
        'CryptGetKeyIdentifierProperty' ,'CryptGetMessageCertificates' ,'CryptGetMessageSignerCount' ,'CryptGetOIDFunctionAddress' ,\
        'CryptGetOIDFunctionValue' ,'CryptHashCertificate' ,'CryptHashMessage' ,'CryptHashPublicKeyInfo' ,'CryptHashToBeSigned' ,\
        'CryptImportPKCS8' ,'CryptImportPublicKeyInfo' ,'CryptImportPublicKeyInfoEx' ,'CryptInitOIDFunctionSet' ,'CryptInstallDefaultContext' ,\
        'CryptInstallOIDFunctionAddress' ,'CryptLoadSip' ,'CryptMemAlloc' ,'CryptMemFree' ,'CryptMemRealloc' ,'CryptMsgCalculateEncodedLength' ,\
        'CryptMsgClose' ,'CryptMsgControl' ,'CryptMsgCountersign' ,'CryptMsgCountersignEncoded' ,'CryptMsgDuplicate' ,'CryptMsgEncodeAndSignCTL' ,\
        'CryptMsgGetAndVerifySigner' ,'CryptMsgGetParam' ,'CryptMsgOpenToDecode' ,'CryptMsgOpenToEncode' ,'CryptMsgSignCTL' ,'CryptMsgUpdate' ,\
        'CryptMsgVerifyCountersignatureEncoded' ,'CryptMsgVerifyCountersignatureEncodedEx' ,'CryptProtectData' ,'CryptQueryObject' ,\
        'CryptRegisterDefaultOIDFunction' ,'CryptRegisterOIDFunction' ,'CryptRegisterOIDInfo' ,'CryptSIPAddProvider' ,\
        'CryptSIPCreateIndirectData' ,'CryptSIPGetSignedDataMsg' ,'CryptSIPLoad' ,'CryptSIPPutSignedDataMsg' ,'CryptSIPRemoveProvider' ,\
        'CryptSIPRemoveSignedDataMsg' ,'CryptSIPRetrieveSubjectGuid' ,'CryptSIPRetrieveSubjectGuidForCatalogFile' ,'CryptSIPVerifyIndirectData' ,\
        'CryptSetAsyncParam' ,'CryptSetKeyIdentifierProperty' ,'CryptSetOIDFunctionValue' ,'CryptSetProviderU' ,'CryptSignAndEncodeCertificate' ,\
        'CryptSignAndEncryptMessage' ,'CryptSignCertificate' ,'CryptSignHashU' ,'CryptSignMessage' ,'CryptSignMessageWithKey' ,\
        'CryptStringToBinaryA' ,'CryptStringToBinaryW' ,'CryptUninstallDefaultContext' ,'CryptUnprotectData' ,'CryptUnregisterDefaultOIDFunction' ,\
        'CryptUnregisterOIDFunction' ,'CryptUnregisterOIDInfo' ,'CryptVerifyCertificateSignature' ,'CryptVerifyCertificateSignatureEx' ,\
        'CryptVerifyDetachedMessageHash' ,'CryptVerifyDetachedMessageSignature' ,'CryptVerifyMessageHash' ,'CryptVerifyMessageSignature' ,\
        'CryptVerifyMessageSignatureWithKey' ,'CryptVerifySignatureU' ,'I_CertProtectFunction' ,'I_CertSrvProtectFunction' ,'I_CertSyncStore' ,\
        'I_CertUpdateStore' ,'I_CryptAddRefLruEntry' ,'I_CryptAddSmartCardCertToStore' ,'I_CryptAllocTls' ,'I_CryptCreateLruCache' ,\
        'I_CryptCreateLruEntry' ,'I_CryptDetachTls' ,'I_CryptDisableLruOfEntries' ,'I_CryptEnableLruOfEntries' ,'I_CryptEnumMatchingLruEntries' ,\
        'I_CryptFindLruEntry' ,'I_CryptFindLruEntryData' ,'I_CryptFindSmartCardCertInStore' ,'I_CryptFlushLruCache' ,'I_CryptFreeLruCache' ,\
        'I_CryptFreeTls' ,'I_CryptGetAsn1Decoder' ,'I_CryptGetAsn1Encoder' ,'I_CryptGetDefaultCryptProv' ,'I_CryptGetDefaultCryptProvForEncrypt' ,\
        'I_CryptGetFileVersion' ,'I_CryptGetLruEntryData' ,'I_CryptGetLruEntryIdentifier' ,'I_CryptGetOssGlobal' ,'I_CryptGetTls' ,'I_CryptInsertLruEntry' ,\
        'I_CryptInstallAsn1Module' ,'I_CryptInstallOssGlobal' ,'I_CryptReadTrustedPublisherDWORDValueFromRegistry' ,'I_CryptRegisterSmartCardStore' ,\
        'I_CryptReleaseLruEntry' ,'I_CryptRemoveLruEntry' ,'I_CryptSetTls' ,'I_CryptTouchLruEntry' ,'I_CryptUninstallAsn1Module' ,\
        'I_CryptUninstallOssGlobal' ,'I_CryptUnregisterSmartCardStore' ,'I_CryptWalkAllLruCacheEntries']

        # Look up there has to be more
        self.service = ['Serv', 'ChangeServiceConfig2A' ,'ChangeServiceConfig2W' ,'ChangeServiceConfigA' ,'ChangeServiceConfigW' ,'CloseServiceHandle' ,\
        'ControlService' ,'CreateServiceA' ,'CreateServiceW' ,'DeleteService' ,'EnumDependentServicesA' ,'EnumDependentServicesW' ,\
        'EnumServiceGroupW' ,'EnumServicesStatusA' ,'EnumServicesStatusExA' ,'EnumServicesStatusExW' ,'EnumServicesStatusW' ,\
        'GetServiceDisplayNameA' ,'GetServiceDisplayNameW' ,'GetServiceKeyNameA' ,'GetServiceKeyNameW' ,'I_ScPnPGetServiceName' ,\
        'I_ScSetServiceBitsA' ,'I_ScSetServiceBitsW' ,'LockServiceDatabase' ,'OpenServiceA' ,'OpenServiceW' ,'PrivilegedServiceAuditAlarmA' ,\
        'PrivilegedServiceAuditAlarmW' ,'QueryServiceConfig2A' ,'QueryServiceConfig2W' ,'QueryServiceConfigA' ,'QueryServiceConfigW' ,\
        'QueryServiceLockStatusA' ,'QueryServiceLockStatusW' ,'QueryServiceObjectSecurity' ,'QueryServiceStatus' ,'QueryServiceStatusEx' ,\
        'RegisterServiceCtrlHandlerA' ,'RegisterServiceCtrlHandlerExA' ,'RegisterServiceCtrlHandlerExW' ,'RegisterServiceCtrlHandlerW' ,\
        'SetServiceBits' ,'SetServiceObjectSecurity' ,'SetServiceStatus' ,'StartServiceA' ,'StartServiceCtrlDispatcherA' ,'StartServiceCtrlDispatcherW' ,\
        'StartServiceW' ,'UnlockServiceDatabase' ,'WdmWmiServiceMain']

        # Generic File Operations
        self.file = ['File', 'CompareFileTime' ,'CopyFileA' ,'CopyFileExA' ,'CopyFileExW' ,'CopyFileW' ,'CopyLZFile' ,'CreateFileA' ,'CreateFileMappingA' ,\
        'CreateFileMappingW' ,'CreateFileW' ,'DeleteFileA' ,'DeleteFileW' ,'DosDateTimeToFileTime' ,'FileTimeToDosDateTime' ,\
        'FileTimeToLocalFileTime' ,'FileTimeToLocalFileTime' ,'FileTimeToSystemTime' ,'FlushFileBuffers' ,'FlushViewOfFile' ,\
        'GetCPFileNameFromRegistry' ,'GetCompressedFileSizeA' ,'GetCompressedFileSizeW' ,'GetFileAttributesA' ,'GetFileAttributesExA' ,\
        'GetFileAttributesExW' ,'GetFileAttributesW' ,'GetFileInformationByHandle' ,'GetFileSize' ,'GetFileSizeEx' ,'GetFileTime' ,\
        'GetFileType' ,'GetSystemTimeAsFileTime' ,'GetTempFileNameA' ,'GetTempFileNameW' ,'LZCloseFile' ,'LZCreateFileW' ,'LZOpenFileA',\
        'LZOpenFileW' ,'LocalFileTimeToFileTime' ,'LocalFileTimeToFileTime' ,'LockFile' ,'LockFileEx' ,'MapViewOfFile' ,'MapViewOfFileEx' ,\
        'MoveFileA' ,'MoveFileExA' ,'MoveFileExW' ,'MoveFileW' ,'MoveFileWithProgressA' ,'MoveFileWithProgressW' ,'OpenDataFile' ,'OpenFile' ,\
        'OpenFileMappingA' ,'OpenFileMappingW' ,'OpenProfileUserMapping' ,'PrivCopyFileExW' ,'PrivMoveFileIdentityW' ,'ReadFile' ,'ReadFileEx' ,\
        'ReplaceFile' ,'ReplaceFileA' ,'ReplaceFileW' ,'SetEndOfFile' ,'SetFileAttributesA' ,'SetFileAttributesW' ,'SetFilePointer' ,\
        'SetFilePointerEx' ,'SetFileShortNameA' ,'SetFileShortNameW' ,'SetFileTime' ,'SetFileValidData' ,'SystemTimeToFileTime' ,\
        'UnlockFile' ,'UnlockFileEx' ,'UnmapViewOfFile' ,'WriteFile' ,'WriteFileEx' ,'WriteFileGather' ,'GetFileSecurityA' ,\
        'GetFileSecurityW' ,'SetFileSecurityA' ,'SetFileSecurityW', 'CreateFileU']

        # APIs related to Collecting information about the host OS
        self.os_info = [ 'Info', 'GetComputerNameA' ,'GetComputerNameExA' ,'GetComputerNameExW' ,'GetComputerNameW' ,'GetDiskFreeSpaceA' ,\
        'GetDiskFreeSpaceExA' ,'GetDiskFreeSpaceExW' ,'GetDiskFreeSpaceW' ,'GetDriveTypeA' ,'GetDriveTypeW', 'GetVersion' ,\
        'GetVersionExA' ,'GetVersionExW', 'GetSystemInfo', 'GetSystemMetrics', 'CheckTokenMembership']

        # List of APIs related to hashing files
        self.cert = ['Cert','CertAddCRLContextToStore' ,'CertAddCRLLinkToStore' ,'CertAddCTLContextToStore' ,'CertAddCTLLinkToStore' ,\
        'CertAddCertificateContextToStore' ,'CertAddCertificateLinkToStore' ,'CertAddEncodedCRLToStore' ,'CertAddEncodedCertificateToStore' ,\
        'CertAddEncodedCertificateToSystemStoreA' ,'CertAddEncodedCertificateToSystemStoreW' ,'CertAddEnhancedKeyUsageIdentifier' ,\
        'CertAddSerializedElementToStore' ,'CertAddStoreToCollection' ,'CertAlgIdToOID' ,'CertCloseStore' ,'CertCompareCertificate' ,\
        'CertCompareCertificateName' ,'CertCompareIntegerBlob' ,'CertComparePublicKeyInfo' ,'CertControlStore' ,'CertCreateCTLContext' ,\
        'CertCreateCTLEntryFromCertificateContextProperties' ,'CertCreateCertificateChainEngine' ,'CertCreateCertificateContext' ,'CertCreateContext',\
        'CertCreateSelfSignCertificate' ,'CertDeleteCTLFromStore' ,'CertDeleteCertificateFromStore' ,'CertDuplicateCTLContext' ,\
        'CertDuplicateCertificateChain' ,'CertDuplicateCertificateContext' ,'CertDuplicateStore' ,'CertEnumCRLContextProperties' ,\
        'CertEnumCRLsInStore' ,'CertEnumCTLContextProperties' ,'CertEnumCTLsInStore' ,'CertEnumCertificateContextProperties' ,\
        'CertEnumCertificatesInStore' ,'CertEnumPhysicalStore' ,'CertEnumSubjectInSortedCTL' ,'CertEnumSystemStore' ,\
        'CertEnumSystemStoreLocation' ,'CertFindAttribute' ,'CertFindCRLInStore' ,'CertFindCertificateInCRL' ,'CertFindCertificateInStore',\
        'CertFindChainInStore' ,'CertFindExtension' ,'CertFindRDNAttr' ,'CertFindSubjectInCTL' ,'CertFindSubjectInSortedCTL' ,\
        'CertFreeCRLContext' ,'CertFreeCertificateChain' ,'CertFreeCertificateChainEngine' ,'CertFreeCertificateContext' ,'CertGetCRLContextProperty' ,\
        'CertGetCRLFromStore' ,'CertGetCTLContextProperty' ,'CertGetCertificateChain' ,'CertGetCertificateContextProperty' ,'CertGetEnhancedKeyUsage' ,\
        'CertGetIssuerCertificateFromStore' ,'CertGetNameStringA' ,'CertGetNameStringW' ,'CertGetPublicKeyLength' ,'CertGetStoreProperty' ,\
        'CertGetSubjectCertificateFromStore' ,'CertGetValidUsages' ,'CertIsRDNAttrsInCertificateName' ,'CertIsValidCRLForCertificate' ,\
        'CertNameToStrA' ,'CertNameToStrW' ,'CertOIDToAlgId' ,'CertOpenStore' ,'CertOpenSystemStoreA' ,'CertOpenSystemStoreW' ,'CertRDNValueToStrA',\
        'CertRDNValueToStrW' ,'CertRegisterPhysicalStore' ,'CertRegisterSystemStore' ,'CertRemoveEnhancedKeyUsageIdentifier' ,\
        'CertRemoveStoreFromCollection' ,'CertResyncCertificateChainEngine' ,'CertSaveStore' ,'CertSerializeCRLStoreElement' ,\
        'CertSerializeCertificateStoreElement' ,'CertSetCRLContextProperty' ,'CertSetCertificateContextPropertiesFromCTLEntry' ,\
        'CertSetCertificateContextProperty' ,'CertSetEnhancedKeyUsage' ,'CertSetStoreProperty' ,'CertStrToNameA' ,'CertStrToNameW' ,\
        'CertUnregisterPhysicalStore' ,'CertUnregisterSystemStore' ,'CertVerifyCRLRevocation' ,'CertVerifyCRLTimeValidity' ,\
        'CertVerifyCTLUsage' ,'CertVerifyCertificateChainPolicy' ,'CertVerifyCertificateChainPolicy' ,'CertVerifyRevocation' ,\
        'CertVerifySubjectCertificateContext','CertVerifyTimeValidity' ,'CertVerifyValidityNesting' ,'CloseCertPerformanceData' ,\
        'CollectCertPerformanceData' ,'CryptAcquireCertificatePrivateKey' ,'CryptFindCertificateKeyProvInfo' ,'CryptGetMessageCertificates' ,\
        'CryptHashCertificate' ,'CryptSignAndEncodeCertificate' ,'CryptSignCertificate' ,'CryptVerifyCertificateSignature' ,\
        'CryptVerifyCertificateSignatureEx' ,'I_CertProtectFunction' ,'I_CertSrvProtectFunction' ,'I_CertSyncStore' ,'I_CertUpdateStore' ,\
        'I_CryptAddSmartCardCertToStore' ,'I_CryptFindSmartCardCertInStore' ,'OpenCertPerformanceData' ,'PFXExportCertStore' ,'PFXExportCertStoreEx' ,'PFXImportCertStore']

        # APIs realted to searching
        self.file_s = ['FSear', 'FindFirstFileW', 'FindNextFileW', 'FindClose']

        # Possible Hook or Injection functions
        self.modify = ['Mod', 'WriteProcessMemory', 'ReadProcessMemory']    

        self.virtual = ['Virt', 'VirtualAlloc' ,'VirtualAllocEx' ,'VirtualBufferExceptionHandler' ,'VirtualFree' ,'VirtualFreeEx' ,'VirtualLock' ,
        'VirtualProtect' ,'VirtualProtectEx' ,'VirtualQuery' ,'VirtualQueryEx' ,'VirtualUnlock']

        self.critical_section = [ 'CrSec', 'DeleteCriticalSection' ,'EnterCriticalSection' ,'InitializeCriticalSection' ,
        'InitializeCriticalSectionAndSpinCount' ,'LeaveCriticalSection' ,'SetCriticalSectionSpinCount' ,'TryEnterCriticalSection']

        #the list are appended to make a matrix. Removes the problem of tracking each list of apis names in the code. Hopefully makes it easier to update.
        self.api_matrix = [ self.reg, self.winsock, self.WinINet, self.cache, self.ftp, self.gopher, self.dir, self.mutex, self.pipe, self.http, self.enum, self.hash, self.crypt, self.service, self.file, self.cert, self.os_info, self.file_s, self.modify, self.virtual, self.critical_section ]

    def generic(self):
        for api_row in self.api_matrix:
            l = api_row[0]
            apis = api_row[1:]
            for api in apis:
                ref_addrs = CodeRefsTo(LocByName(api),0)
                for ref in ref_addrs:
                    func_addr = LocByName(GetFunctionName(ref))
                    func_name = GetFunctionName(ref)
                    if self.rename == True:
                        if l not in func_name:
                            MakeNameEx(func_addr , l + '_' + func_name , SN_NOWARN)
                    if self.comment == True:
                        cmt = ''
                        cmt = GetFunctionCmt(func_addr,1)
                        if l not in cmt:
                            cmt = cmt + ' ' + l 
                            SetFunctionCmt(func_addr, cmt, 1)

def InitialAnalysis_main():
    revokeAnalysis()

    # find unidentified functions
    findUnidentifiedFunctions()

    global InstructionCallbacks
    InstructionCallbacks = []
    InstructionCallbacks.append(CryptoTester())
    InstructionCallbacks.append(CallTester())

    iterateInstructions()

    # refresh ida view to display our results
    idaapi.refresh_idaview_anyway()

    #SetColorSiko
    setcolorsiko()

    #FuncRename
    ah = GENAPI()
    # To disable comments assign to False
    ah.comment = True
    # To disable renaming of functions to False
    ah.rename = True    
    ah.generic()

    print "InitialAnalysis completed"

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


repository: paradiseduo/ida-swift-demangle
file: testSwift.py
description: A tool to demangle Swift function names in IDA.
wanted_name: Swift Dump
comment: Swift Dump plugin for IDA Pro 7.0
created_at: 2021-07-15 08:44:17+00:00
pushed_at: 2021-07-30 07:34:38+00:00
url:
https://github.com/paradiseduo/ida-swift-demangle/blob/fc2dd8a05fb1d9b1afd7ce6f2632da859802b0b4/testSwift.py

content:
import idaapi
import idc
import idautils
import subprocess

class Swift_Dump_t(idaapi.plugin_t):
    comment = "Swift Dump plugin for IDA Pro 7.0"
    help = ""
    wanted_name = "Swift Dump"
    wanted_hotkey = "Ctrl-Alt-D"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        return idaapi.PLUGIN_OK

    def term(self):
        return idaapi.PLUGIN_OK

    def patcher(self):
        for func in idautils.Functions():
            oldname = GetFunctionName(func)
            if oldname[:2] == "_$":
                oldname = oldname[2:]
                p = subprocess.Popen('xcrun swift-demangle --simplified --compact '+oldname, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                newname = str(p.communicate()[0].decode('utf-8', 'ignore').strip())
                if newname != oldname and len(newname) > 2:
                    MakeNameEx(func, newname, SN_NOCHECK|SN_NOWARN)
                    SetFunctionCmt(func, "%s \n%s" % (oldname, newname), 1)

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


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


repository: poohzaza166/pooh-plugin-assistant
file: plugins/userinput/main.py
description: This is a showcase of pytohn plugin system for creating voice assistant
created_at: 2025-03-01 22:01:41+00:00
pushed_at: 2025-03-23 16:19:46+00:00
url:
https://github.com/poohzaza166/pooh-plugin-assistant/blob/21decaa3651a0191e1006927a31c097dd3c8bde3/plugins/userinput/main.py

content:
from .pooh_lib import (PluginBase, get_data, loop_method, provide_data,
                       push_event, register_command, subscribe)


def plugin_entry(message_bus, parser):
    print("loading plugin")
    print("checking parser object")
    print(parser)
    plugin = GetUserInput(message_bus, parser)
    return  plugin


class GetUserInput(PluginBase):
    def __init__(self, message_bus, parser):
        super().__init__(message_bus, parser)
        self.input_count = 0
    
    # @provide_data("text_input")
    # def get_user_input(self, param):
    #     print("doing stdsadasdsadaddsdadasdsdasdasdauff")
    #     return "hello world"
    
    @provide_data("test")
    def return_str(self):
        return "aaaaaaa"

    @get_data("conversation_history")
    def print_latest_conversation(self, obj, **kwargs):
        print(obj)
        print(kwargs)

    @loop_method(delay=2)
    def run_publish(self):
        # self.input_count += 1
        # print(self.input_count)
        self.publish("text_input", input("please insert a querry"))
        print(self.print_latest_conversation())
        # ok = self.message_bus.get_data("test")  
        # print(ok)


repository: richor1042/IDAFuncOutline
file: funcOutline.py
description: Resume FuncOutline by idapython
wanted_name: funcOutline UI Plugin
comment: funcOutline UI Plugin
created_at: 2024-08-13 02:25:34+00:00
pushed_at: 2024-08-20 06:30:28+00:00
url:
https://github.com/richor1042/IDAFuncOutline/blob/3b0670de63458576a9dc236d58f309fdbbe9c58b/funcOutline.py

content:
import idautils
import ida_name
import ida_funcs
import idaapi
import idc
import re

ins_num_limit=7
arm64_ins_space=4
def set_outline_func_flag(func_addr):
    pfn = ida_funcs.get_func(func_addr)
    pfn.flags |= idaapi.FUNC_OUTLINE
    ida_funcs.update_func(pfn)

def check_func(func_addr):
    start_ea = idc.get_func_attr(func_addr, idc.FUNCATTR_START)
    end_ea = idc.get_func_attr(func_addr, idc.FUNCATTR_END)
    # Limit number of instructions
    if end_ea-start_ea>ins_num_limit*arm64_ins_space:
        return False
    #Limit the last instruction
    ins = idc.GetDisasm(end_ea-arm64_ins_space)
    opt_code = ins.split("             ")
    if (opt_code[0]!="B")and(opt_code[0]!="RET"):
        print("last ins no rule",hex(func_addr))
        return False
    return True

def is_hex_string(string):
    try:
        int(string, 16)
        return True
    except ValueError:
        return False

def visit_func(blk):
    pat = "_[0-9a-fA-f]+"
    curr_addr = blk.start_ea
    while curr_addr < blk.end_ea:
        ins=idc.GetDisasm(curr_addr)
        opt_code = ins.split("             ")

        if len(opt_code)==2:
            if ((opt_code[0].find("BL")==0)and(opt_code[0]!="BLR")) and (opt_code[1].replace(" ","").find("_")!=0):
                pat = "_[0-9a-fA-f]+"
                tmpcont=re.findall(pat,opt_code[1])
                func_son_addr=tmpcont[0][1:]
                if is_hex_string(func_son_addr):
                    if check_func(int(func_son_addr,16)):
                        # print("set_outline_func_flag", func_son_addr)
                        set_outline_func_flag(int(func_son_addr,16))

        curr_addr = idc.next_head(curr_addr, blk.end_ea)

def outline_resume(fn_addr):
    print("Extract Resume",hex(fn_addr))
    f_blocks = idaapi.FlowChart(ida_funcs.get_func(fn_addr))
    for block in f_blocks:
        visit_func(block)
    ida_funcs.update_func(fn_addr)
    print("Extract Resume finish",hex(fn_addr))


PSEUDOCODE_CLICK = "PSEUDOCODEClick"
MENU_PATH = "FuncOutlineUI/"

class FuncOutlineUIHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, widget, popup):
        if idaapi.get_widget_type(widget) == idaapi.BWN_PSEUDOCODE:
            idaapi.attach_action_to_popup(widget, popup, PSEUDOCODE_CLICK, MENU_PATH)

class FuncOutlineUIActions(idaapi.action_handler_t):
    def __init__(self, name):
        idaapi.action_handler_t.__init__(self)
        self.name = name

    def activate(self, ctx):
        if self.name == PSEUDOCODE_CLICK:
            ea = idc.here()
            now_func = ida_funcs.get_func(ea)
            outline_resume(now_func.start_ea)
        return 1

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

class FuncOutlineUI(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "funcOutline UI Plugin"
    help = "This is a funcOutline UI plugin"
    wanted_name = "funcOutline UI Plugin"
    wanted_hotkey = ""
    def init(self):
        self.hooks = FuncOutlineUIHooks()
        self.hooks.hook()


        # Register action for PSEUDOCODEClick
        action_desc = idaapi.action_desc_t(
            PSEUDOCODE_CLICK,
            "FuncOutline resume",
            FuncOutlineUIActions(PSEUDOCODE_CLICK),
            '',
            "FuncOutline resume",
            50)
        idaapi.register_action(action_desc)

        return idaapi.PLUGIN_KEEP

    def term(self):
        self.hooks.unhook()
        idaapi.unregister_action(PSEUDOCODE_CLICK)
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return FuncOutlineUI()


repository: Rhythm113/Tools
file: ioctl_plugin.py
description: Simple Tools For Game Hacking Purpose & Others
url: https://github.com/Rhythm113/Tools/blob/93d34e21760d38b11caa6e0d67f7b2659f008ad7/ioctl_plugin.py
wanted_name: Find IOCTL
comment: Find IOCTL functions and analyze cases
created_at: 2021-04-28 09:07:42+00:00
pushed_at: 2023-10-31 08:16:07+00:00
content:
#Date: 2023-06-20 09:41 PM
#Simple IOCTL Analyzer helper by @Rhythm113


import idaapi
import idc
import idautils

def find_ioctl_functions():
    ioctl_functions = []
    for func_ea in Functions():
        func_name = GetFunctionName(func_ea)
        if "ioctl" in func_name:
            ioctl_functions.append(func_ea)

    return ioctl_functions

def analyze_ioctl_function(func_ea):
    print("Analyzing function at address: 0x{:X}".format(func_ea))
    case_map = {}  # Dictionary to store case connections
    root_case = None

    for (startea, endea) in idautils.Chunks(func_ea):
        for head in idautils.Heads(startea, endea):
            if idc.isCode(GetFlags(head)):
                insn = idc.GetDisasm(head)
                if "case" in insn:
                    case_num = idc.GetOperandValue(head, 0)
                    case_map[case_num] = []  # Initialize an empty list for case connections

                    if not root_case:
                        root_case = case_num  # Set the first case encountered as the root case

                elif "j" in insn and insn.startswith("jumptable"):
                    targets = idautils.CodeRefsFrom(head, 1)
                    for target in targets:
                        target_insn = idc.GetDisasm(target)
                        if "case" in target_insn:
                            target_case_num = idc.GetOperandValue(target, 0)
                            case_map[case_num].append(target_case_num)  # Add case connection

    return root_case, case_map

def find_connected_classes():
    connected_classes = []
    for class_ea in idautils.ClassInstances():
        class_name = idc.GetClassName(class_ea)
        lib_name = idc.GetLongPrm(idc.INF_SHORT_LPSNAME)
        if lib_name not in class_name:
            connected_classes.append((class_name, class_ea))

    return connected_classes

class FindIoctlPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Find IOCTL functions and analyze cases"
    help = "This plugin searches for functions containing 'ioctl' in their names, analyzes the case connections, and identifies classes connected to other libraries."
    wanted_name = "Find IOCTL"
    wanted_hotkey = "Alt-F9"
    author = "Rhythm113"

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 730:  # Check IDA SDK version for compatibility
            return idaapi.PLUGIN_OK
        else:
            return idaapi.PLUGIN_SKIP

    def run(self, arg):
        ioctl_functions = find_ioctl_functions()
        if len(ioctl_functions) == 0:
            print("No functions containing 'ioctl' found.")
        else:
            print("Found {} functions containing 'ioctl':".format(len(ioctl_functions)))
            for func_ea in ioctl_functions:
                root_case, case_map = analyze_ioctl_function(func_ea)
                print("Root case number: {}".format(root_case))
                for case_num, connections in case_map.items():
                    print("Case {}: connected to cases {}".format(case_num, connections))

            connected_classes = find_connected_classes()
            if len(connected_classes) == 0:
                print("No classes connected to other libraries found.")
            else:
                print("Found {} classes connected to other libraries:".format(len(connected_classes)))
                for class_name, class_ea in connected_classes:
                    print("Class: {} at address 0x{:X}".format(class_name, class_ea))

    def term(self):
        pass

def PLUGIN_ENTRY():
    return FindIoctlPlugin()
    


repository: gistable/gistable
file: all-gists/c6936d219db8e6635d25/snippet.py
description: GitHub gists and the dockerfiles needed to run them.
wanted_name: simpa64
comment: Simplifier
created_at: 2018-02-01 16:19:24+00:00
pushed_at: 2023-02-11 18:20:10+00:00
url:
https://github.com/gistable/gistable/blob/665d39a2bd82543d5196555f0801ef8fd4a3ee48/all-gists/c6936d219db8e6635d25/snippet.py

content:
#  AArch64 mov simplifier IDA plugin
#
#  Copyright (c) 2015 xerub
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# based on Rolf Rolles x86 deobfuscator http://www.msreverseengineering.com

import idaapi
import idc

ARM64_MOVE_I = idaapi.ARM_mov

def dump_cmd(cmd):
	print "cs = %lx" % cmd.cs
	print "ip = %lx" % cmd.ip
	print "ea = %lx" % cmd.ea
	print "itype = %lx" % cmd.itype
	print "size = %lx" % cmd.size
	print "auxpref = %lx" % cmd.auxpref
	print "segpref = %lx" % cmd.segpref
	print "insnpref = %lx" % cmd.insnpref
	print "flags = %lx" % cmd.flags

def dump_op(op):
	print "n = %lx" % op.n
	print "type = %lx" % op.type
	print "offb = %lx" % op.offb
	print "offo = %lx" % op.offo
	print "flags = %lx" % op.flags
	print "dtyp = %lx" % op.dtyp
	print "reg = %lx" % op.reg
	print "phrase = %lx" % op.phrase
	print "value = %lx" % op.value
	print "addr = %lx" % op.addr
	print "specval = %lx" % op.specval
	print "specflag1 = %lx" % op.specflag1
	print "specflag2 = %lx" % op.specflag2
	print "specflag3 = %lx" % op.specflag3
	print "specflag4 = %lx" % op.specflag4

def HighestSetBit(N, imm):
	i = N - 1
	while i >= 0:
		if imm & (1 << i):
			return i
		i -= 1
	return -1

def ZeroExtendOnes(M, N):				# zero extend M ones to N width
	return (1 << M) - 1

def RORZeroExtendOnes(M, N, R):
	val = ZeroExtendOnes(M, N)
	return ((val >> R) & ((1 << (N - R)) - 1)) | ((val & ((1 << R) - 1)) << (N - R))

def Replicate(val, bits):
	ret = val
	shift = bits
	while shift < 64:				# XXX actually, it is either 32 or 64
		ret |= (val << shift)
		shift += bits
	return ret

def DecodeBitMasks(immN, imms, immr, immediate):
	len = HighestSetBit(7, (immN << 6) | (~imms & 0x3F))
	if len < 1:
		return None
	levels = ZeroExtendOnes(len, 6)
	if immediate and (imms & levels) == levels:
		return None
	S = imms & levels
	R = immr & levels
	esize = 1 << len
	return Replicate(RORZeroExtendOnes(S + 1, esize, R), esize)

def DecodeMov(opcode, total, first):
	# opc
	o = (opcode >> 29) & 3
	# constant
	k = (opcode >> 23) & 0x3F

	if k == 0x24 and o == 1:			# MOV (bitmask imm) <=> ORR (immediate)
		# sf
		s = (opcode >> 31) & 1
		# N
		N = (opcode >> 22) & 1
		if s == 0 and N != 0:
			return None
		# rn
		rn = (opcode >> 5) & 0x1F
		if rn == 31:
			imms = (opcode >> 10) & 0x3F
			immr = (opcode >> 16) & 0x3F
			return DecodeBitMasks(N, imms, immr, True)
	elif k == 0x25:					# MOVN/MOVZ/MOVK
		# sf
		s = (opcode >> 31) & 1
		# hw
		h = (opcode >> 21) & 3
		# imm16
		i = (opcode >> 5) & 0xFFFF
		if s == 0 and h > 1:
			return None
		h *= 16
		i <<= h
		if o == 0:				# MOVN
			return ~i
		elif o == 2:				# MOVZ
			return i
		elif o == 3 and not first:		# MOVK
			return (total & ~(0xFFFF << h)) | i
	elif (k | 1) == 0x23 and not first:		# ADD (immediate)
		# shift
		h = (opcode >> 22) & 3
		if h > 1:
			return None
		# rn
		rd = opcode & 0x1F
		rn = (opcode >> 5) & 0x1F
		if rd != rn:
			return None
		# imm12
		i = (opcode >> 10) & 0xFFF
		h *= 12
		i <<= h
		if o & 2:				# SUB
			return total - i
		else:					# ADD
			return total + i

	return None

def check_mov_sequence(ea):
	oldea = ea
	reg = -1
	total = 0
	is64 = False
	while idaapi.getseg(ea).use64():
		d = idaapi.get_long(ea)
		# reg
		r = d & 0x1F
		if reg >= 0 and reg != r:
			break
		newval = DecodeMov(d, total, reg < 0)
		if newval is None:
			break
		if reg >= 0 and idaapi.get_first_fcref_to(ea) != idaapi.BADADDR:
			break
		if (d >> 31) & 1:
			is64 = True
		total = newval
		reg = r
		ea += 4
	return ea - oldea, reg, is64, total

def is_my_mov(cmd):
	if cmd.itype == ARM64_MOVE_I and cmd.flags == idaapi.INSN_MACRO and cmd.size > 4:
		return True
	return False

def check_ubfm_shift(ea):
	if idaapi.getseg(ea).use64():
		opcode = idaapi.get_long(ea)
		# opc
		o = (opcode >> 29) & 3
		# constant
		k = (opcode >> 23) & 0x3F
		if (o & 1) == 0 and k == 0x26:
			# sf
			s = (opcode >> 31) & 1
			# N
			N = (opcode >> 22) & 1
			if s == N:
				# imm
				imms = (opcode >> 10) & 0x3F
				immr = (opcode >> 16) & 0x3F
				mask = 0x1F | ((s & N) << 5)
				if imms == mask:
					return idaapi.ARM_lsr if o else idaapi.ARM_asr, opcode, s, immr
				elif immr == imms + 1:
					return idaapi.ARM_lsl if o else idaapi.ARM_null, opcode, s, mask - imms
	return idaapi.ARM_null, 0, 0, 0

class simpA64Hook(idaapi.IDP_Hooks):
	def __init__(self):
		idaapi.IDP_Hooks.__init__(self)
		self.n = idaapi.netnode("$ A64 Simplifier",0,1)

	def custom_ana(self):
		len, reg, is64, imm = check_mov_sequence(idaapi.cmd.ea)
		if len > 4:
			#print "0x%x: MOV/MOVK %c%d, #0x%x" % (idaapi.cmd.ea, 'X' if is64 else 'W', reg, imm)
			#dump_cmd(idaapi.cmd)
			#dump_op(idaapi.cmd.Op1)
			#dump_op(idaapi.cmd.Op2)
			idaapi.cmd.itype = ARM64_MOVE_I
			idaapi.cmd.segpref = 14 # ARM Condition = ALways
			idaapi.cmd.Op1.type = idaapi.o_reg
			idaapi.cmd.Op1.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op1.reg = reg + 129 # Use Wn/Xn registers instead of Rn
			idaapi.cmd.Op2.type = idaapi.o_imm
			idaapi.cmd.Op2.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op2.value = imm
			idaapi.cmd.flags = idaapi.INSN_MACRO
			idaapi.cmd.size = len
			return True
		insn, regs, is64, shift = check_ubfm_shift(idaapi.cmd.ea)
		if insn != idaapi.ARM_null:
			idaapi.cmd.itype = insn
			idaapi.cmd.segpref = 14
			idaapi.cmd.Op1.type = idaapi.o_reg
			idaapi.cmd.Op1.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op1.reg = (regs & 0x1F) + 129
			idaapi.cmd.Op2.type = idaapi.o_reg
			idaapi.cmd.Op2.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op2.reg = ((regs >> 5) & 0x1F) + 129
			idaapi.cmd.Op3.type = idaapi.o_imm
			idaapi.cmd.Op3.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op3.value = shift
			idaapi.cmd.size = 4
			return True
		return False

	def custom_mnem(self): # totally optional
		if is_my_mov(idaapi.cmd):
			return "MOVE"
		return None

#	def custom_out(self): # XXX ida would just append .EQ
#		if is_my_mov(idaapi.cmd):
#			buf = idaapi.init_output_buffer(1024)
#			idaapi.OutMnem(16, "")
#			idaapi.out_one_operand(0)
#			idaapi.out_symbol(',')
#			idaapi.OutChar(' ')
#			idaapi.out_one_operand(1)
#			idaapi.term_output_buffer()
#			idaapi.MakeLine(buf)
#			return True
#		return False

#	def custom_outop(self, op): # XXX ida would just use Rn
#		if is_my_mov(idaapi.cmd) and op.type == idaapi.o_reg:
#			idaapi.out_register("%c%d" % ('X' if op.dtyp == idaapi.dt_qword else 'W', op.reg))
#			return True
#		return False

class simpa64_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_PROC
	comment = "Simplifier"
	wanted_hotkey = "Alt-Z"
	help = "Runs transparently"
	wanted_name = "simpa64"
	hook = None
	enabled = 1

	def init(self):
		self.hook = None
		if idaapi.ph_get_id() != idaapi.PLFM_ARM or idaapi.BADADDR <= 0xFFFFFFFF:
			return idaapi.PLUGIN_SKIP

		self.hook = simpA64Hook()
		flag = self.hook.n.altval(0)
		if flag:
			self.enabled = flag - 1
		print "%s is %sabled" % (self.wanted_name, "en" if self.enabled else "dis")
		if self.enabled:
			self.hook.hook()
		return idaapi.PLUGIN_KEEP

	def run(self, arg):
		print "%sabling %s" % ("dis" if self.enabled else "en", self.wanted_name)
		if self.enabled:
			self.hook.unhook()
		else:
			self.hook.hook()
		self.enabled = self.enabled ^ 1
		self.hook.n.altset(0, self.enabled + 1)
		idc.Refresh()

	def term(self):
		if self.hook:
			self.hook.unhook()

def PLUGIN_ENTRY():
	return simpa64_t()


repository: clovme/WTools
file: software/Reverse Cracking/IDA Professional 9.0/python/examples/hexrays/vds17.py
wanted_name: Structure offsets (IDAPython)
comment: Sample plugin17 for Hex-Rays decompiler
created_at: 2025-03-27 00:14:54+00:00
pushed_at: 2025-05-11 10:18:21+00:00
url:
https://github.com/clovme/WTools/blob/f12f01c833a6a6241799a9d01e9117cc28a65dbc/software/Reverse%20Cracking/IDA%20Professional%209.0/python/examples/hexrays/vds17.py

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

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

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

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

import ida_idaapi
import ida_hexrays
import ida_lines
import ida_typeinf
import ida_kernwin

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

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

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

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

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

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

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

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

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

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

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

        return 1

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


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

def PLUGIN_ENTRY():
    return my_plugin_t()

"""
# A few notes about the VDS17 sample

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

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

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


## How the user interface works

Let us suppose that we have the following local types:

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

and the decompiler generates the following pseudocode:

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

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

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

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

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

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

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

The right pane will change to:

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

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

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


### What do we need this plugin for?

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


## API details

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

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

Then the magic begins:

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


### The devil is in the details

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

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

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

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

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

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

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


repository: yoavst/Graffiti
file: backends/ida/graffiti.py
description: Create customized callgraph directly from your favorite IDE
wanted_name: Graffiti
comment: Create customized callgraph directly from IDA
created_at: 2022-06-11 10:50:40+00:00
pushed_at: 2025-03-18 18:33:14+00:00
url:
https://github.com/yoavst/Graffiti/blob/07a075fb5f4406479ad55b564256049422f6ed09/backends/ida/graffiti.py

content:
import json
import socket
import os
import sys
import struct
import threading

import idaapi
import ida_kernwin
import ida_nalt
import idc
import idautils
import ida_funcs
import ida_segment

from authentication import (
    get_token_or_else,
    get_last_connected_server,
    save_last_connected_server_to_file,
)


HAVE_WIN32_LIBS = False
if sys.platform == "win32":
    try:
        import win32gui, win32con

        HAVE_WIN32_LIBS = True
    except:
        print(
            (
                "********************************************\n"
                "*   Install the following for better GUI   *\n"
                "*          pip install pywin32             *\n"
                "********************************************"
            )
        )


def to_bytes(s):
    if sys.version_info.major == 3:
        return s.encode("utf-8")
    else:
        return s


def from_bytes(b):
    if sys.version_info.major == 3:
        return b.decode("utf-8")
    else:
        return b


def readexactly(sock, num_bytes):
    buf = b""
    while len(buf) < num_bytes:
        buf += sock.recv(num_bytes - len(buf))

    return buf


def lengthy_send(sock, data):
    sock.send(struct.pack(">i", len(data)))
    sock.send(data)


sock = None


class AddToGraphHandler(idaapi.action_handler_t):
    def __init__(self, force_func=True):
        idaapi.action_handler_t.__init__(self)
        self.force_func = force_func

    # Say hello when invoked.
    def activate(self, ctx):
        global sock
        if ctx.cur_func or not self.force_func:
            my_sock = sock
            if my_sock is not None:
                payload = self.create_payload(ctx)
                if payload:
                    data = to_bytes(json.dumps(payload))
                    lengthy_send(my_sock, data)
                else:
                    print("Graffiti: Empty payload")
            else:
                print("Graffiti: Not connected to server")
        else:
            print("Graffiti: Not currently in a function.")

        return True

    def create_payload(self, ctx):
        if ctx.cur_func:
            # Get the name of the current function
            current_func_name = idc.get_func_name(ctx.cur_func.start_ea)
            demangled_func_name = idc.demangle_name(current_func_name, idc.INF_SHORT_DN)
            func_name = demangled_func_name or current_func_name

            # Get the name of the current IDA database
            db_path = ida_nalt.get_input_file_path()
            db_filename = os.path.basename(db_path)

            return {
                "type": "addData",
                "node": {
                    "project": "IDA: " + db_filename,
                    "address": str(ctx.cur_func.start_ea),
                    "baseName": func_name,
                    "baseAddress": str(ctx.cur_func.start_ea),
                    "computedProperties": [
                        {"name": "label", "format": "{0}", "replacements": ["baseName"]}
                    ],
                },
            }

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


class AddToGraphWithEdgeInfoHandler(AddToGraphHandler):
    def create_payload(self, ctx):
        base = AddToGraphHandler.create_payload(self, ctx)
        edge_text = ida_kernwin.ask_str("", 2, "Enter label for edge")
        if edge_text and edge_text.strip():
            base["edge"] = {"label": edge_text}
        return base


class AddToGraphLineHandler(AddToGraphHandler):
    def __init__(self):
        AddToGraphHandler.__init__(self, False)

    def create_payload(self, ctx):
        base = AddToGraphHandler.create_payload(self, ctx)
        if base:
            base["node"]["computedProperties"] = [
                {
                    "name": "label",
                    "format": "{0}+{1}",
                    "replacements": ["baseName", "line"],
                }
            ]
            base["node"]["address"] = str(ctx.cur_ea)
            base["node"]["line"] = "{0:x}".format(ctx.cur_ea - ctx.cur_func.start_ea)
            return base
        else:
            # Non-function address
            db_path = ida_nalt.get_input_file_path()
            db_filename = os.path.basename(db_path)

            return {
                "type": "addData",
                "node": raw_node_from_ea(ctx.cur_ea, db_filename),
            }


class AddXrefsHandler(AddToGraphHandler):
    def __init__(self, line_xrefs=False):
        AddToGraphHandler.__init__(self)
        self.line_xrefs = line_xrefs

    def create_payload(self, ctx):
        db_path = ida_nalt.get_input_file_path()
        db_filename = os.path.basename(db_path)

        nodes = dict()
        for xref in idautils.XrefsTo(ctx.cur_func.start_ea):
            f = ida_funcs.get_func(xref.frm)
            if f is None:
                # Raw address
                nodes[str(xref.frm)] = raw_node_from_ea(xref.frm, db_filename)
            else:
                original_func_name = idc.get_func_name(f.start_ea)
                demangled_func_name = idc.demangle_name(
                    original_func_name, idc.INF_SHORT_DN
                )
                func_name = demangled_func_name or original_func_name

                node = {
                    "project": "IDA: " + db_filename,
                    "baseName": func_name,
                    "baseAddress": str(f.start_ea),
                }
                if not self.line_xrefs:
                    node["address"] = str(f.start_ea)
                    node["computedProperties"] = [
                        {"name": "label", "format": "{0}", "replacements": ["baseName"]}
                    ]
                else:
                    node["address"] = str(xref.frm)
                    node["line"] = "{0:x}".format(xref.frm - f.start_ea)
                    node["computedProperties"] = [
                        {
                            "name": "label",
                            "format": "{0}+{1}",
                            "replacements": ["baseName", "line"],
                        }
                    ]
                nodes[node["address"]] = node

        if nodes:
            return {"type": "addDataBulk", "nodes": list(nodes.values()), "direction": "n2e"}


def raw_node_from_ea(ea, db_filename):
    seg = ida_segment.getseg(ea)
    if seg:
        seg_name = ida_segment.get_segm_name(seg) + ":"
    else:
        seg_name = ""

    name = idc.get_name(ea, idc.GN_DEMANGLED)
    if not name or name == idc.BADADDR:
        name = "{0:x}".format(ea)

    return {
        "project": "IDA: " + db_filename,
        "baseAddress": str(ea),
        "address": str(ea),
        "line": "{0:x}".format(ea),
        "baseName": name,
        "seg": seg_name,
        "computedProperties": [
            {"name": "label", "format": "{0}{1}", "replacements": ["seg", "baseName"]}
        ],
    }


class AddXrefLinesHandler(AddXrefsHandler):
    def __init__(self):
        AddXrefsHandler.__init__(self, True)


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

    # Say hello when invoked.
    def activate(self, ctx):
        global sock
        if sock is not None:
            sock.close()
            sock = None
        server = get_last_connected_server() or "localhost:8501"
        addr = ida_kernwin.ask_str(
            server, 2, "What is the address of the grafiti server?"
        )
        if not addr:
            return

        addr, port = addr.split(":")
        sock = socket.socket()
        sock.connect((addr, int(port)))
        save_last_connected_server_to_file(addr + ":" + port)

        db_path = ida_nalt.get_input_file_path()
        db_filename = os.path.basename(db_path)

        thread = threading.Thread(target=sync_read_thread, args=(db_filename,))
        thread.daemon = True
        thread.start()

        return True

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


def bring_ida_to_foreground():
    # according to https://hex-rays.com/blog/plugin-focus-heimdallr/
    # Tested only on macos, should be adapted easily to Linux/Windows
    if sys.platform in ["darwin", "win32"]:
        # https://www.riverbankcomputing.com/static/Docs/PyQt5/
        qtwidget = ida_kernwin.PluginForm.TWidgetToPyQtWidget(
            ida_kernwin.get_current_viewer()
        )
        window = qtwidget.window()

        # UnMinimize
        WindowMinimized = 0x00000001  # https://www.riverbankcomputing.com/static/Docs/PyQt5/api/qtcore/qt.html#WindowState
        cur_state = window.windowState()
        new_state = cur_state & (~WindowMinimized)
        window.setWindowState(new_state)

        # Switch desktop / give keyboard control

        if sys.platform == "darwin":
            window.show()
            window.raise_()
        elif sys.platform == "win32":
            # https://stackoverflow.com/a/25970104
            if HAVE_WIN32_LIBS:
                win32gui.SetWindowPos(
                    window.winId(),
                    win32con.HWND_TOPMOST,
                    0,
                    0,
                    0,
                    0,
                    win32con.SWP_NOMOVE | win32con.SWP_NOSIZE | win32con.SWP_SHOWWINDOW,
                )
                win32gui.SetWindowPos(
                    window.winId(),
                    win32con.HWND_NOTOPMOST,
                    0,
                    0,
                    0,
                    0,
                    win32con.SWP_NOMOVE | win32con.SWP_NOSIZE | win32con.SWP_SHOWWINDOW,
                )
            window.raise_()
            window.show()
            window.activateWindow()


def sync_read_thread(db_filename):
    global sock
    print("Background thread running")
    try:
        while True:
            length = struct.unpack(">i", readexactly(sock, 4))[0]
            data = json.loads(readexactly(sock, length))

            if "type" in data and data["type"] == "auth_req_v1":
                token_storage = []

                def ask_for_token():
                    def ui_ask_for_token():
                        token = ida_kernwin.ask_str(
                            "", 2, "Enter the UUID token from graffiti website"
                        )
                        if token and token.strip():
                            token_storage.append(token)

                    idaapi.execute_sync(ui_ask_for_token, idaapi.MFF_FAST)
                    return token_storage[0] if token_storage else None

                token = get_token_or_else(ask_for_token)
                if token is not None:
                    lengthy_send(
                        sock,
                        to_bytes(json.dumps({"type": "auth_resp_v1", "token": token})),
                    )
                    continue
                else:
                    break

            if "project" in data:
                if data["project"] != "IDA: {}".format(db_filename):
                    continue

            addr = int(data["address"])

            def on_ui():
                bring_ida_to_foreground()
                ida_kernwin.jumpto(addr)
                return False

            ida_kernwin.execute_ui_requests([on_ui])

        if sock is not None:
            sock.close()
    except socket.error:
        print("Socket is closed")
        sock = None


# 4) Register hooks
class GraffitiUIHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        if (
            idaapi.get_widget_type(form) == idaapi.BWN_DISASM
            or idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE
        ):
            idaapi.attach_action_to_popup(
                form, popup, "graffiti:addToGraph", "Graffiti/"
            )
            idaapi.attach_action_to_popup(
                form, popup, "graffiti:addToGraphWithEdgeInfo", "Graffiti/"
            )
            idaapi.attach_action_to_popup(
                form, popup, "graffiti:addLineToGraph", "Graffiti/"
            )
            idaapi.attach_action_to_popup(
                form, popup, "graffiti:addXrefsToGraph", "Graffiti/"
            )
            idaapi.attach_action_to_popup(
                form, popup, "graffiti:addXrefLinesToGraph", "Graffiti/"
            )


class GrafitiIDBHooks(idaapi.IDB_Hooks):
    def renamed(self, ea, new_name, is_local):
        global sock
        if sock is not None and new_name:
            payload = {
                "type": "updateNodes",
                "selection": [["baseAddress", str(ea)]],
                "update": {"baseName": new_name},
            }
            lengthy_send(sock, to_bytes(json.dumps(payload)))


class GraffitiPlugin(idaapi.plugin_t, idaapi.UI_Hooks):
    flags = idaapi.PLUGIN_MOD | idaapi.PLUGIN_HIDE
    comment = "Create customized callgraph directly from IDA"
    help = ""
    wanted_name = "Graffiti"
    wanted_hotkey = ""

    def __init__(self):
        super(idaapi.plugin_t, self).__init__()
        self.hooks = []

    def init(self):
        """plugin_t init() function"""
        self.hooks = [GraffitiUIHooks(), GrafitiIDBHooks()]
        for hook in self.hooks:
            hook.hook()

        # 2) Describe the actions
        add_to_graph_action_desc = idaapi.action_desc_t(
            "graffiti:addToGraph",
            "Graffiti: Add to graph",  # The action text.
            AddToGraphHandler(),  # The action handler.
            "Ctrl+Shift+A",  # Optional: the action shortcut
            "Add to graph",
        )
        add_to_graph_with_edge_info_action_desc = idaapi.action_desc_t(
            "graffiti:addToGraphWithEdgeInfo",
            "Graffiti: Add to graph with edge comment",  # The action text.
            AddToGraphWithEdgeInfoHandler(),  # The action handler.
            "Ctrl+Shift+X",  # Optional: the action shortcut
            "Add to graph with edge comment",
        )
        add_line_to_graph_action_desc = idaapi.action_desc_t(
            "graffiti:addLineToGraph",
            "Graffiti: Add current line to graph",  # The action text.
            AddToGraphLineHandler(),  # The action handler.
            "Ctrl+Alt+A",  # Optional: the action shortcut
            "Add line to graph",
        )
        add_xrefs_to_graph_action_desc = idaapi.action_desc_t(
            "graffiti:addXrefsToGraph",
            "Graffiti: Add function xrefs to graph",  # The action text.
            AddXrefsHandler(),  # The action handler.
            "Ctrl+Shift+Q",  # Optional: the action shortcut
            "Add function xrefs to graph",
        )
        add_xrefs_lines_to_graph_action_desc = idaapi.action_desc_t(
            "graffiti:addXrefLinesToGraph",
            "Graffiti: Add line xrefs to graph",  # The action text.
            AddXrefLinesHandler(),  # The action handler.
            "Ctrl+Alt+Shift+Q",  # Optional: the action shortcut
            "Add line xrefs to graph",
        )

        enable_graffiti_sync_action_desc = idaapi.action_desc_t(
            "graffiti:ConnectToServer",
            "Graffiti: Connect to server",  # The action text.
            EnableSyncHandler(),  # The action handler.
            None,  # Optional: the action shortcut
            "Connect To Server",
        )

        # 3) Register the action
        idaapi.unregister_action("graffiti:addToGraph")
        idaapi.register_action(add_to_graph_action_desc)
        idaapi.unregister_action("graffiti:addToGraphWithEdgeInfo")
        idaapi.register_action(add_to_graph_with_edge_info_action_desc)
        idaapi.unregister_action("graffiti:addLineToGraph")
        idaapi.register_action(add_line_to_graph_action_desc)
        idaapi.unregister_action("graffiti:addXrefsToGraph")
        idaapi.register_action(add_xrefs_to_graph_action_desc)
        idaapi.unregister_action("graffiti:addXrefLinesToGraph")
        idaapi.register_action(add_xrefs_lines_to_graph_action_desc)
        idaapi.unregister_action("graffiti:ConnectToServer")
        idaapi.register_action(enable_graffiti_sync_action_desc)

        idaapi.attach_action_to_menu(
            "Options/Source paths...", "graffiti:ConnectToServer", idaapi.SETMENU_APP
        )
        return idaapi.PLUGIN_KEEP

    def run(self, arg=0):
        """plugin_t run() implementation"""
        return

    def term(self):
        """plugin_t term() implementation"""
        for hook in self.hooks:
            hook.unhook()
        return


def PLUGIN_ENTRY():
    return GraffitiPlugin()


repository: AzzOnFire/yarka
file: yarka.py
description: IDA plugin for YARA signature creation
url: https://github.com/AzzOnFire/yarka/blob/4e159e51c52484d6c97186b867327e740558e98e/yarka.py
wanted_name: Yarka
comment: Yarka Plugin
created_at: 2023-11-25 21:13:56+00:00
pushed_at: 2024-10-19 13:08:33+00:00
content:
from itertools import chain

import idaapi
import ida_kernwin
import ida_funcs

from yarka.adapter import YaraExtractor
from yarka.extractor import RangeExtractor, FunctionExtractor
from yarka.yara import RuleBuilder, RulesetBuilder
from yarka import ui
from yarka import utils


# rule parameters
WRAP_CURLY_BRACE = False        # wrap curly brace on a new line
INDENT_HEADERS = False          # indent for 'meta', 'strings', 'condition'

# editor default settings
DEFAULT_SHOW_COMMENTS = True    # show comments (checkbox)
DEFAULT_STRICT_RULE = False     # don't wildcard relative offsets (checkbox)
DEFAULT_INDENT = 2              # default indent

# values can be string or lambda
DEFAULT_META_FIELDS = {
    'description': '',
    'source': 'Internal research',
    'date': lambda: utils.get_current_time('%Y-%m-%d'),
    'hash': lambda: utils.get_file_md5(),
    # 'hash': lambda: utils.get_file_sha256(),
}

# values can be string or lambda (optional argument - strings count)
DEFAULT_CONDITIONS = [
    lambda: '(PE)' in idaapi.get_file_type_name() and 'uint16(0) == 0x5A4D',
    lambda: 'ELF' in idaapi.get_file_type_name() and 'uint32(0) == 0x464C457F',
    lambda count: f'{count // 2} of them' if count > 2 else 'all of them'
]


def generate_rule(name, entities) -> RuleBuilder:
    rule = RuleBuilder(
        name,
        indent_headers=INDENT_HEADERS,
        wrap_curly_brace=WRAP_CURLY_BRACE,
    )

    for header, value in DEFAULT_META_FIELDS.items():
        rule.add_meta(header, value)
    
    existing = set()
    for i, entity in enumerate(entities, start=1):
        if str(entity) not in existing:
            rule.add_string(f's{i}', entity)
            existing.add(str(entity))

    for condition in DEFAULT_CONDITIONS:
        rule.add_condition(condition)

    return rule


class YaraDialog(ui.ClosableDialog):
    def __init__(self, rule: RuleBuilder):
        self.rule = rule

        self.strict_checkbox = ui.Checkbox(
            label='strict',
            default=DEFAULT_STRICT_RULE,
            on_click=self.handle_strict_checkbox_click,
        )
        self.comments_checkbox = ui.Checkbox(
            label='show comments',
            default=DEFAULT_SHOW_COMMENTS,
            on_click=self.handle_comments_checkbox_click,
        )
        self.indent_edit_label = ui.Label('indentation')
        self.indent_edit = ui.NumberInput(
            default=DEFAULT_INDENT,
            min_value=2,
            max_value=8,
            step=2,
            on_change=self.handle_indent_change,
        )
        self.text_edit = ui.YaraTextEdit(rule.build())

        super().__init__(
            title = 'Yarka - Generated Yara Rule',
            width=800,
            top_left_items=[self.strict_checkbox, self.comments_checkbox],
            top_right_items=[self.indent_edit_label, self.indent_edit],
            body_items=[self.text_edit]
        )

    def _rebuild_rule(self):
        content = self.rule.build()
        self.text_edit.set_content(content)

    def handle_strict_checkbox_click(self):
        self.rule.strict = self.strict_checkbox.is_checked()
        self._rebuild_rule()

    def handle_comments_checkbox_click(self):
        self.rule.show_comments = self.comments_checkbox.is_checked()
        self._rebuild_rule()

    def handle_indent_change(self):
        self.rule.indent = self.indent_edit.value()
        self._rebuild_rule()


class YarkaPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = 'Yarka Plugin'
    help = 'Yarka Plugin'
    wanted_name = 'Yarka'
    wanted_hotkey = 'Ctrl+Alt+Y'
    banner = (
        " __     __        _         "
        " \ \   / /       | |        "
        "  \ \_/ /_ _ _ __| | ____ _ "
        "   \   / _` | '__| |/ / _` |"
        "    | | (_| | |  |   < (_| |"
        "    |_|\__,_|_|  |_|\_\__,_|"
        "                            "
        "                            "
    )

    def init(self):
        self.ui_hooks = YarkaUIHooks()
        self.ui_hooks.hook()

        label = 'Yarka: Generate Rule'
        hotkey = 'Ctrl+Y'
        YarkaRule.register(self, label, hotkey)
        YarkaStrings.register(self, label, hotkey)
        YarkaFunctions.register(self, label, hotkey)

        idaapi.msg(f'[{self.wanted_name}] loaded, hotkey - {self.wanted_hotkey}')
        return idaapi.PLUGIN_KEEP

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

    def run(self, arg):
        ida_kernwin.show_wait_box('Generating ruleset ...')

        ruleset = RulesetBuilder()
        functions = list(utils.get_custom_functions())
        for i, function in enumerate(functions, start=1):
            if ida_kernwin.user_cancelled():
                ida_kernwin.hide_wait_box()
                ida_kernwin.warning('Yarka - ruleset generation canceled')
                return None

            ida_kernwin.replace_wait_box(f'Processing ({i}/{len(functions)})')

            try:
                name = ida_funcs.get_func_name(function.start_ea)
                entities = YaraExtractor(FunctionExtractor(function.start_ea))
                rule = generate_rule(name=name, entities=list(entities))
                if not rule.is_empty():
                    ruleset.append(rule)
            except Exception as e:
                ida_kernwin.warning(f'Yarka - unexpected error {e}')

        ida_kernwin.replace_wait_box(f'Building ruleset ...')

        self.dialog = YaraDialog(ruleset)
        self.dialog.show()
        ida_kernwin.hide_wait_box()


class YarkaRule(ui.ActionHandler):
    __windows__ = [idaapi.BWN_DISASM, idaapi.BWN_PSEUDOCODE]

    def activate(self, ctx):
        start, end = utils.get_selection()
        name = f'{utils.get_file_name()}_0x{start:08X}'
        entities = YaraExtractor(RangeExtractor(start, end))
        rule = generate_rule(name=name, entities=entities)

        self.dialog = YaraDialog(rule)
        self.dialog.show()


class YarkaStrings(ui.ActionHandler):
    __windows__ = [idaapi.BWN_STRINGS]

    def activate(self, ctx):
        strings = []
        for idx in ctx.chooser_selection:
            data = ida_kernwin.get_chooser_data(ctx.widget_title, idx)
            location, *_ = data
            address = int(location.split(':')[-1].lstrip('0'), 16)

            string_gen = YaraExtractor(RangeExtractor(address, address + 1))
            strings.append(string_gen)
        
        name = f'{utils.get_file_name()}_strings'
        entities = chain.from_iterable(strings)
        rule = generate_rule(name=name, entities=entities)

        self.dialog = YaraDialog(rule)
        self.dialog.show()


class YarkaFunctions(ui.ActionHandler):
    __windows__ = [idaapi.BWN_FUNCS]

    def activate(self, ctx):
        ruleset = RulesetBuilder()
        for idx in ctx.chooser_selection:
            _, _, location, *_ = ida_kernwin.get_chooser_data(ctx.widget_title, idx)
            address = int(location.lstrip('0'), 16)

            name = ida_funcs.get_func_name(address)
            if not name:
                continue

            entities = YaraExtractor(FunctionExtractor(address))
            rule = generate_rule(name=name, entities=list(entities))
            if not rule.is_empty():
                ruleset.append(rule)

        self.dialog = YaraDialog(ruleset)
        self.dialog.show()


class YarkaUIHooks(idaapi.UI_Hooks):
    def populating_widget_popup(self, form, popup):
        pass

    def finish_populating_widget_popup(self, form, popup):
        YarkaRule.attach_ctx_menu(form, popup, 'Yarka')
        YarkaStrings.attach_ctx_menu(form, popup, 'Yarka')
        YarkaFunctions.attach_ctx_menu(form, popup, 'Yarka')


plugin = YarkaPlugin()


def PLUGIN_ENTRY():
    global plugin
    return plugin


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

content:
from idaapi import *
import idaapi
import idc
from prettytable import PrettyTable

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

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

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

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


# describe arg num of function

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

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

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

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


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

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


def getArgAddr(start_addr, regNum):
    armcondition=[]

    #大部分，这个深度是够的，但是仍有小部分超过这个范围
    scan_deep = 100
    count = 0
    reg = "R" + str(regNum)

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



def getArg(start_addr, regNum):
    arm=['LDR','CMP','STR','MOV']
    arg_addr = getArgAddr(start_addr, regNum)

    if arg_addr != BADADDR:
        Mnemonics = idc.print_insn_mnem(arg_addr)
        if Mnemonics[0:3] == "ADD":
            pc_rx_str = idc.print_operand(arg_addr, 1) + "+" +  idc.print_operand(arg_addr, 2)
            if pc_rx_str[0:4]=="PC+R":
                mid_addr=getArgAddr(arg_addr,int(pc_rx_str[-1]))
                arg =  idc.print_operand(arg_addr, 1) + "+" + idc.print_operand(mid_addr, 1)
            else:
                arg = pc_rx_str
        elif Mnemonics[0:3] == "SUB":
            arg = idc.print_operand(arg_addr, 1) + "-" +  idc.print_operand(arg_addr, 2)
        elif Mnemonics[0:3] in arm:
            arg = idc.print_operand(arg_addr, 1)
        else:
            arg = idc.generate_disasm_line(arg_addr,1).split(" ")[0]
        set_cmt(arg_addr, "addr: 0x%x " % start_addr  + "-------> arg" + str((int(regNum)+1)) + " : " + arg, 0)
        return arg
    else:
        return "get fail"
#-------------------------------------------------------------------#



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


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


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

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


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

    if Rx_arg in ['R9','R11','R4','R6','R10']:
        if len(Rx_arg)==2:
            true_arg_addr=getArgAddr(mid_arg_addr,int(Rx_arg[-1]))
            true_arg_addr=getArgAddr(true_arg_addr,int(Rx_arg[-1]))
        else:
            true_arg_addr = getArgAddr(mid_arg_addr, int(Rx_arg[1:]))
            true_arg_addr = getArgAddr(true_arg_addr, int(Rx_arg[1:]))
    else:
        true_arg_addr=getArgAddr(mid_arg_addr,format_function_offset_dict[func_name])
    string_and_addr =  getFormatString(true_arg_addr)

    if Rx_arg in ['R9','R11','R4','R6','R10']:
        format_arg=string_and_addr[2]
        ret_list.pop(-1)
        ret_list.append(format_arg)

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


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

    # get arg num and set table
    if func_name in one_arg_function:
        arg_num = 1
    elif func_name in two_arg_function:
        arg_num = 2
    elif func_name in three_arg_function:
        arg_num = 3
    elif func_name in format_function_offset_dict:
        arg_num = format_function_offset_dict[func_name] + 1
    else:
        print("The %s function didn't write in the describe arg num of function array,please add it to,such as add to `two_arg_function` arary" % func_name)
        return

    table_head = ["func_name", "addr"]
    for num in range(0,arg_num):
        table_head.append("arg"+str(num+1))

    if func_name in format_function_offset_dict:
        table_head.append("format&value[string_addr, num of '%', fmt_arg...]")
    table_head.append("local_buf_size")
    table = PrettyTable(table_head)

    #循环遍历调用当前函数的地址
    call_addr = get_first_cref_to(func_addr)
    while call_addr != BADADDR:
        idc.set_color(call_addr, idc.CIC_ITEM, 0xffff00)
        #获取当前地址的指令
        Mnemonics = print_insn_mnem(call_addr)
        if Mnemonics[0:1] == "B":
            if func_name in format_function_offset_dict:
                info = auditFormat(call_addr, func_name, arg_num)
            else:
                info = auditAddr(call_addr, func_name, arg_num)
            table.add_row(info)
        call_addr = get_next_cref_to(func_addr, call_addr)
    print(table)



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

    print("Auditing attention function ......")
    for func_name in attention_function:
        audit(func_name)

    print("Auditing command execution function ......")
    for func_name in command_execution_function:
        audit(func_name)

    print("Finished! Enjoy the result ~")

m_initialized = False

class ArmAudit_Plugin_t(idaapi.plugin_t):
    comment = "ARM Audit plugin for IDA Pro" # 描述信息
    help = "todo" # 帮助信息
    wanted_name = "ArmAudit" # 菜单中显示的名字
    wanted_hotkey = "Ctrl-Alt-A" # 希望注册的快捷键
    flags = idaapi.PLUGIN_KEEP # 插件的状态, 当前状态保持在Plugin菜单中

    def init(self):
        global m_initialized

        if m_initialized is False:
            m_initialized = True
            print("=" * 80)
            start = '''
            _                      _             _ _ _   
           / \   _ __ _ __ ___    / \  _   _  __| (_) |_ 
          / _ \ | '__| '_ ` _ \  / _ \| | | |/ _` | | __|
         / ___ \| |  | | | | | |/ ___ \ |_| | (_| | | |_ 
        /_/   \_\_|  |_| |_| |_/_/   \_\__,_|\__,_|_|\__|

                            re-edit by The_Itach1  2022.8.3
            '''
            print(start)
            print("=" * 80)

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        '''
                    每次运行插件时, 执行的具体操作
                    功能代码在此编写
        '''

        info = idaapi.get_inf_structure()
        if 'ARM' in info.procName:
            ArmAudit()
        else:
            print('ArmAudit is not supported on the current arch')

def PLUGIN_ENTRY():
    '''
            插件入口,用于实例对象
            返回的就是插件的功能等
    '''
    return ArmAudit_Plugin_t()

if __name__ == '__main__':
    ArmAudit()


repository: seritools/ida-settings
file: plugins/IDARustDemangler.py
created_at: 2023-01-11 23:20:40+00:00
pushed_at: 2023-06-08 18:29:28+00:00
url:
https://github.com/seritools/ida-settings/blob/79469449f44b57a108208a58526f65a3685a9fdf/plugins/IDARustDemangler.py

content:
import shutil
import subprocess
import os

try:
    import ida_kernwin
    import idautils
    import idaapi
    import idc
except ModuleNotFoundError:
    print("[!] This is an IDA script")
    exit(-1)


class IDARustDemangler():
    """
    IDARustDemangler is a tool that demangles and normalizes symbols for use with the IDA disassembler.
    It replaces or modifies special characters to make them compatible with IDA's syntax, 
    making binary analysis faster and more efficient.
    """

    def __init__(self, debug=False):
        self.debug = debug
        self.num_resolved = 0
        self.hash_prefix = "17h"
        self.delimiters = "><"
        self.badchars = "*,'`" + self.delimiters
        self.queue = {}
        self.resolved = {}
        self.rs_dml = shutil.which("rs-dml")
        assert len(self.rs_dml) > 0, "rs-dml is not installed on your system"

    def add(self, address: int, symbol: str) -> None:
        # Must add support for non legacy symbols

        # If the hash is present
        if symbol.lstrip("_").startswith("Z") and self.hash_prefix in symbol:
            hash = symbol.split(self.hash_prefix)[-1].rstrip("E")

            # If the hash length is not 16 skip as it is not a valid rust legacy symbol
            if len(hash) != 16:
                return

            self.queue[address, hash] = symbol.encode()

    def resolve(self):
        # Pass symbols which should be resolved to the stdin of `rs-dml` and map
        # the results to the correct values
        process_rs_dml = subprocess.Popen(
            [self.rs_dml], stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.STDOUT)
        resolved_symbols, _ = process_rs_dml.communicate(
            input=b"\n".join(self.queue.values()))
        self.resolved = zip(self.queue.keys(),
                            resolved_symbols.decode().splitlines())

    def apply(self):
        # for each symbol resolved normalize it and apply it to the IDA db

        for (address, hash), symbol in self.resolved:
            normalized = self.ida_normalize(symbol)

            # If normalization succeded and all the character in the normalized
            # symbol are valid set the name in IDA

            if any([badchar in normalized for badchar in self.badchars]) and self.debug:
                print(
                    f"[ERROR] {address:#016x} -> sym:'{symbol}', hash: '{hash}', normalized: '{normalized}'")
                continue

            if self.debug:
                print(
                    f"[*] {address:#016x} -> sym:'{symbol}', hash: '{hash}', normalized: '{normalized}'")

            # set the name and add the hash at the end
            idc.set_name(address, normalized + str(len(hash)) + hash)
            self.num_resolved += 1

    def ida_normalize(self, name: str) -> str:
        """
        This function tries to normalize sybmols to be accepted by IDA
        """

        # Replace bad characters with accepted ones
        # unfortunately there is no way known to me
        # to insert these chars into IDA symbol names

        name = name.replace(" ", "_")
        name = name.replace(",", "")
        name = name.replace("{", "<")
        name = name.replace("}", ">")

        i = 0
        output = "_ZN"

        while i < len(name):
            if name[i] == "<":
                # 'I' corresponds to '<'
                output += "I"
                i += 1
            elif name[i] == ">":
                # 'E' corresponds to '>'
                output += "E"
                i += 1
            elif name[i] == "*":
                # 'P' corresponds to pointer-type word
                output += "P"
                i += 1
            else:

                # if it the `word` starts with "::" skip it
                # as IDA automatically adds it
                if name[i:i+2] == "::":
                    i += 2

                # this should find the closest delimiter to
                # recognize the entire word

                idxs = []

                for special_char in self.delimiters:
                    tmp_idx = name[i+1:].find(special_char)

                    if tmp_idx != -1:
                        idxs.append(tmp_idx)

                if len(idxs) >= 1:
                    idx = min(idxs)
                    word = name[i:i+idx + 1]
                else:
                    word = name[i:]

                if len(word) > 0:
                    if "*" in word:
                        # if '*' is present it means it is a pointer just have
                        # to a number of 'P' at the start of the word
                        # corresponding to the number of '*'
                        output += "P" * word.count("*")
                        i += word.count("*")
                        word = word.replace("*", "")

                    output += str(len(word)) + word

                i += len(word)

        output += "E"
        return output


# Register the actual plugin
class IDARustDemangleHandler(idaapi.plugin_t):
    PLUGIN_NAME = "IDA Rust Demangler"
    PLUGIN_DIRECTORY = "IDARustDemangler"
    PLUGIN_DESCRIPTION = "Rust name demangler"

    flags = idaapi.PLUGIN_UNL
    comment = PLUGIN_DESCRIPTION
    help = PLUGIN_DESCRIPTION
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""

    def init(self):
        self.r = IDARustDemangler(debug=False)
        return idaapi.PLUGIN_OK

    def run(self, arg):
        print(f"[*] {self.PLUGIN_NAME} started!")

        # push every symbol name that is found
        # invalid symbols will be discarded
        for address, name in idautils.Names():
            self.r.add(address, name)

        # resolve with rs-dml
        self.r.resolve()

        # apply the results in ida
        self.r.apply()

        print(f"[*] Demangled {self.r.num_resolved} symbols")
        print(f"[*] {self.PLUGIN_NAME} has finished!")

        return 1

    def term(self):
        pass


class IDARustDemanglerHook(ida_kernwin.UI_Hooks):
    """
    this class is only used to install the icon to the corresponding IDA action
    """

    def __init__(self, cb):
        super().__init__()
        self.cb = cb

    def updated_actions(self):
        if self.cb():
            self.unhook()


def install_icon():
    plugin_name = IDARustDemangleHandler.PLUGIN_NAME
    action_name = "Edit/Plugins/" + plugin_name
    LOGO_PATH = None

    # if the action is not present wait for our hook action
    if action_name not in ida_kernwin.get_registered_actions():
        return False

    # check if in any of the IDA plugins directory if there is
    # our plugin directory and take the logo from there
    for plugin_path in idaapi.get_ida_subdirs("plugins"):
        LOGO_PATH = os.path.join(
            plugin_path, f"{IDARustDemangleHandler.PLUGIN_DIRECTORY}\\rust-logo.png")

        # if the file exists use the first one found
        if os.path.isfile(LOGO_PATH):
            break

    if LOGO_PATH is None:
        print("[?] IDA Rust Demangler logo not found")
        return True

    # load the logo and apply it to the action
    icon = idaapi.load_custom_icon(
        LOGO_PATH, format="png")

    ida_kernwin.update_action_icon(action_name, icon)

    return True


def PLUGIN_ENTRY():
    return IDARustDemangleHandler()


h = IDARustDemanglerHook(install_icon)
h.hook()


repository: Cobra-de1/ida-pro-mcp
file: mcp-plugin.py
description: Fork of https://github.com/mrexodia/ida-pro-mcp
wanted_name: MCP
comment: MCP Plugin
created_at: 2025-04-21 10:00:46+00:00
pushed_at: 2025-04-21 10:24:02+00:00
url:
https://github.com/Cobra-de1/ida-pro-mcp/blob/0a8405d3168ce9f1b5a9e5dac43d2cbb1403d6d3/mcp-plugin.py

content:
import sys

if sys.version_info < (3, 11):
    raise RuntimeError("Python 3.11 or higher is required for the MCP plugin")
import re
import json
import struct
import threading
import http.server
from urllib.parse import urlparse
from typing import Any, Callable, get_type_hints, TypedDict, Optional, Annotated, TypeVar, Generic

class JSONRPCError(Exception):
    def __init__(self, code: int, message: str, data: Any = None):
        self.code = code
        self.message = message
        self.data = data

class RPCRegistry:
    def __init__(self):
        self.methods: dict[str, Callable] = {}

    def register(self, func: Callable) -> Callable:
        self.methods[func.__name__] = func
        return func

    def dispatch(self, method: str, params: Any) -> Any:
        if method not in self.methods:
            raise JSONRPCError(-32601, f"Method '{method}' not found")

        func = self.methods[method]
        hints = get_type_hints(func)

        # Remove return annotation if present
        hints.pop("return", None)

        if isinstance(params, list):
            if len(params) != len(hints):
                raise JSONRPCError(-32602, f"Invalid params: expected {len(hints)} arguments, got {len(params)}")

            # Validate and convert parameters
            converted_params = []
            for value, (param_name, expected_type) in zip(params, hints.items()):
                try:
                    if not isinstance(value, expected_type):
                        value = expected_type(value)
                    converted_params.append(value)
                except (ValueError, TypeError):
                    raise JSONRPCError(-32602, f"Invalid type for parameter '{param_name}': expected {expected_type.__name__}")

            return func(*converted_params)
        elif isinstance(params, dict):
            if set(params.keys()) != set(hints.keys()):
                raise JSONRPCError(-32602, f"Invalid params: expected {list(hints.keys())}")

            # Validate and convert parameters
            converted_params = {}
            for param_name, expected_type in hints.items():
                value = params.get(param_name)
                try:
                    if not isinstance(value, expected_type):
                        value = expected_type(value)
                    converted_params[param_name] = value
                except (ValueError, TypeError):
                    raise JSONRPCError(-32602, f"Invalid type for parameter '{param_name}': expected {expected_type.__name__}")

            return func(**converted_params)
        else:
            raise JSONRPCError(-32600, "Invalid Request: params must be array or object")

rpc_registry = RPCRegistry()

def jsonrpc(func: Callable) -> Callable:
    """Decorator to register a function as a JSON-RPC method"""
    global rpc_registry
    return rpc_registry.register(func)

class JSONRPCRequestHandler(http.server.BaseHTTPRequestHandler):
    def send_jsonrpc_error(self, code: int, message: str, id: Any = None):
        response = {
            "jsonrpc": "2.0",
            "error": {
                "code": code,
                "message": message
            }
        }
        if id is not None:
            response["id"] = id
        response_body = json.dumps(response).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", len(response_body))
        self.end_headers()
        self.wfile.write(response_body)

    def do_POST(self):
        global rpc_registry

        parsed_path = urlparse(self.path)
        if parsed_path.path != "/mcp":
            self.send_jsonrpc_error(-32098, "Invalid endpoint", None)
            return

        content_length = int(self.headers.get("Content-Length", 0))
        if content_length == 0:
            self.send_jsonrpc_error(-32700, "Parse error: missing request body", None)
            return

        request_body = self.rfile.read(content_length)
        try:
            request = json.loads(request_body)
        except json.JSONDecodeError:
            self.send_jsonrpc_error(-32700, "Parse error: invalid JSON", None)
            return

        # Prepare the response
        response = {
            "jsonrpc": "2.0"
        }
        if request.get("id") is not None:
            response["id"] = request.get("id")

        try:
            # Basic JSON-RPC validation
            if not isinstance(request, dict):
                raise JSONRPCError(-32600, "Invalid Request")
            if request.get("jsonrpc") != "2.0":
                raise JSONRPCError(-32600, "Invalid JSON-RPC version")
            if "method" not in request:
                raise JSONRPCError(-32600, "Method not specified")

            # Dispatch the method
            result = rpc_registry.dispatch(request["method"], request.get("params", []))
            response["result"] = result

        except JSONRPCError as e:
            response["error"] = {
                "code": e.code,
                "message": e.message
            }
            if e.data is not None:
                response["error"]["data"] = e.data
        except IDAError as e:
            response["error"] = {
                "code": -32000,
                "message": e.message,
            }
        except Exception as e:
            traceback.print_exc()
            response["error"] = {
                "code": -32603,
                "message": "Internal error (please report a bug)",
                "data": traceback.format_exc(),
            }

        try:
            response_body = json.dumps(response).encode("utf-8")
        except Exception as e:
            traceback.print_exc()
            response_body = json.dumps({
                "error": {
                    "code": -32603,
                    "message": "Internal error (please report a bug)",
                    "data": traceback.format_exc(),
                }
            }).encode("utf-8")

        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", len(response_body))
        self.end_headers()
        self.wfile.write(response_body)

    def log_message(self, format, *args):
        # Suppress logging
        pass

class MCPHTTPServer(http.server.HTTPServer):
    allow_reuse_address = False

class Server:
    HOST = "localhost"
    PORT = 13337

    def __init__(self):
        self.server = None
        self.server_thread = None
        self.running = False

    def start(self):
        if self.running:
            print("[MCP] Server is already running")
            return

        self.server_thread = threading.Thread(target=self._run_server, daemon=True)
        self.running = True
        self.server_thread.start()

    def stop(self):
        if not self.running:
            return

        self.running = False
        if self.server:
            self.server.shutdown()
            self.server.server_close()
        if self.server_thread:
            self.server_thread.join()
            self.server = None
        print("[MCP] Server stopped")

    def _run_server(self):
        try:
            # Create server in the thread to handle binding
            self.server = MCPHTTPServer((Server.HOST, Server.PORT), JSONRPCRequestHandler)
            print(f"[MCP] Server started at http://{Server.HOST}:{Server.PORT}")
            self.server.serve_forever()
        except OSError as e:
            if e.errno == 98 or e.errno == 10048:  # Port already in use (Linux/Windows)
                print("[MCP] Error: Port 13337 is already in use")
            else:
                print(f"[MCP] Server error: {e}")
            self.running = False
        except Exception as e:
            print(f"[MCP] Server error: {e}")
        finally:
            self.running = False

# A module that helps with writing thread safe ida code.
# Based on:
# https://web.archive.org/web/20160305190440/http://www.williballenthin.com/blog/2015/09/04/idapython-synchronization-decorator/
import logging
import queue
import traceback
import functools

import ida_hexrays
import ida_kernwin
import ida_funcs
import ida_gdl
import ida_lines
import ida_idaapi
import idc
import idaapi
import idautils
import ida_nalt
import ida_bytes
import ida_typeinf
import ida_xref
import ida_entry
import idautils

class IDAError(Exception):
    def __init__(self, message: str):
        super().__init__(message)

    @property
    def message(self) -> str:
        return self.args[0]

class IDASyncError(Exception):
    pass

# Important note: Always make sure the return value from your function f is a
# copy of the data you have gotten from IDA, and not the original data.
#
# Example:
# --------
#
# Do this:
#
#   @idaread
#   def ts_Functions():
#       return list(idautils.Functions())
#
# Don't do this:
#
#   @idaread
#   def ts_Functions():
#       return idautils.Functions()
#

logger = logging.getLogger(__name__)

# Enum for safety modes. Higher means safer:
class IDASafety:
    ida_kernwin.MFF_READ
    SAFE_NONE = ida_kernwin.MFF_FAST
    SAFE_READ = ida_kernwin.MFF_READ
    SAFE_WRITE = ida_kernwin.MFF_WRITE

call_stack = queue.LifoQueue()

def sync_wrapper(ff, safety_mode: IDASafety):
    """
    Call a function ff with a specific IDA safety_mode.
    """
    #logger.debug('sync_wrapper: {}, {}'.format(ff.__name__, safety_mode))

    if safety_mode not in [IDASafety.SAFE_READ, IDASafety.SAFE_WRITE]:
        error_str = 'Invalid safety mode {} over function {}'\
                .format(safety_mode, ff.__name__)
        logger.error(error_str)
        raise IDASyncError(error_str)

    # No safety level is set up:
    res_container = queue.Queue()

    def runned():
        #logger.debug('Inside runned')

        # Make sure that we are not already inside a sync_wrapper:
        if not call_stack.empty():
            last_func_name = call_stack.get()
            error_str = ('Call stack is not empty while calling the '
                'function {} from {}').format(ff.__name__, last_func_name)
            #logger.error(error_str)
            raise IDASyncError(error_str)

        call_stack.put((ff.__name__))
        try:
            res_container.put(ff())
        except Exception as x:
            res_container.put(x)
        finally:
            call_stack.get()
            #logger.debug('Finished runned')

    ret_val = idaapi.execute_sync(runned, safety_mode)
    res = res_container.get()
    if isinstance(res, Exception):
        raise res
    return res

def idawrite(f):
    """
    decorator for marking a function as modifying the IDB.
    schedules a request to be made in the main IDA loop to avoid IDB corruption.
    """
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        ff = functools.partial(f, *args, **kwargs)
        ff.__name__ = f.__name__
        return sync_wrapper(ff, idaapi.MFF_WRITE)
    return wrapper

def idaread(f):
    """
    decorator for marking a function as reading from the IDB.
    schedules a request to be made in the main IDA loop to avoid
      inconsistent results.
    MFF_READ constant via: http://www.openrce.org/forums/posts/1827
    """
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        ff = functools.partial(f, *args, **kwargs)
        ff.__name__ = f.__name__
        return sync_wrapper(ff, idaapi.MFF_READ)
    return wrapper

def is_window_active():
    """Returns whether IDA is currently active"""
    try:
        from PyQt5.QtWidgets import QApplication
    except ImportError:
        return False

    app = QApplication.instance()
    if app is None:
        return False

    for widget in app.topLevelWidgets():
        if widget.isActiveWindow():
            return True
    return False

class Metadata(TypedDict):
    path: str
    module: str
    base: str
    size: str
    md5: str
    sha256: str
    crc32: str
    filesize: str

def get_image_size():
    try:
        # https://www.hex-rays.com/products/ida/support/sdkdoc/structidainfo.html
        info = idaapi.get_inf_structure()
        omin_ea = info.omin_ea
        omax_ea = info.omax_ea
    except AttributeError:
        import ida_ida
        omin_ea = ida_ida.inf_get_omin_ea()
        omax_ea = ida_ida.inf_get_omax_ea()
    # Bad heuristic for image size (bad if the relocations are the last section)
    image_size = omax_ea - omin_ea
    # Try to extract it from the PE header
    header = idautils.peutils_t().header()
    if header and header[:4] == b"PE\0\0":
        image_size = struct.unpack("<I", header[0x50:0x54])[0]
    return image_size

@jsonrpc
@idaread
def get_metadata() -> Metadata:
    """Get metadata about the current IDB"""
    # Fat Mach-O binaries can return a None hash:
    # https://github.com/mrexodia/ida-pro-mcp/issues/26
    def hash(f):
        try:
            return f().hex()
        except:
            return None
    return {
        "path": idaapi.get_input_file_path(),
        "module": idaapi.get_root_filename(),
        "base": hex(idaapi.get_imagebase()),
        "size": hex(get_image_size()),
        "md5": hash(ida_nalt.retrieve_input_file_md5),
        "sha256": hash(ida_nalt.retrieve_input_file_sha256),
        "crc32": hex(ida_nalt.retrieve_input_file_crc32()),
        "filesize": hex(ida_nalt.retrieve_input_file_size()),
    }

def get_prototype(fn: ida_funcs.func_t) -> Optional[str]:
    try:
        prototype: ida_typeinf.tinfo_t = fn.get_prototype()
        if prototype is not None:
            return str(prototype)
        else:
            return None
    except AttributeError:
        try:
            return idc.get_type(fn.start_ea)
        except:
            tif = ida_typeinf.tinfo_t()
            if ida_nalt.get_tinfo(tif, fn.start_ea):
                return str(tif)
            return None
    except Exception as e:
        print(f"Error getting function prototype: {e}")
        return None

class Function(TypedDict):
    address: str
    name: str
    size: str

def parse_address(address: str) -> int:
    try:
        return int(address, 0)
    except ValueError:
        for ch in address:
            if ch not in "0123456789abcdefABCDEF":
                raise IDAError(f"Failed to parse address: {address}")
        raise IDAError(f"Failed to parse address (missing 0x prefix): {address}")

def get_function(address: int, *, raise_error=True) -> Function:
    fn = idaapi.get_func(address)
    if fn is None:
        if raise_error:
            raise IDAError(f"No function found at address {hex(address)}")
        return None

    try:
        name = fn.get_name()
    except AttributeError:
        name = ida_funcs.get_func_name(fn.start_ea)
    return {
        "address": hex(fn.start_ea),
        "name": name,
        "size": hex(fn.end_ea - fn.start_ea),
    }

DEMANGLED_TO_EA = {}

def create_demangled_to_ea_map():
    for ea in idautils.Functions():
        # Get the function name and demangle it
        # MNG_NODEFINIT inhibits everything except the main name
        # where default demangling adds the function signature
        # and decorators (if any)
        demangled = idaapi.demangle_name(
            idc.get_name(ea, 0), idaapi.MNG_NODEFINIT)
        if demangled:
            DEMANGLED_TO_EA[demangled] = ea

@jsonrpc
@idaread
def get_function_by_name(
    name: Annotated[str, "Name of the function to get"]
) -> Function:
    """Get a function by its name"""
    function_address = idaapi.get_name_ea(idaapi.BADADDR, name)
    if function_address == idaapi.BADADDR:
        # If map has not been created yet, create it
        if len(DEMANGLED_TO_EA) == 0:
            create_demangled_to_ea_map()
        # Try to find the function in the map, else raise an error
        if name in DEMANGLED_TO_EA:
            function_address = DEMANGLED_TO_EA[name]
        else:
            raise IDAError(f"No function found with name {name}")
    return get_function(function_address)

@jsonrpc
@idaread
def get_function_by_address(
    address: Annotated[str, "Address of the function to get"]
) -> Function:
    """Get a function by its address"""
    return get_function(parse_address(address))

@jsonrpc
@idaread
def get_current_address() -> str:
    """Get the address currently selected by the user"""
    return hex(idaapi.get_screen_ea())

@jsonrpc
@idaread
def get_current_function() -> Optional[Function]:
    """Get the function currently selected by the user"""
    return get_function(idaapi.get_screen_ea())

class ConvertedNumber(TypedDict):
    decimal: str
    hexadecimal: str
    bytes: str
    ascii: Optional[str]
    binary: str

@jsonrpc
def convert_number(
    text: Annotated[str, "Textual representation of the number to convert"],
    size: Annotated[Optional[int], "Size of the variable in bytes"],
) -> ConvertedNumber:
    """Convert a number (decimal, hexadecimal) to different representations"""
    try:
        value = int(text, 0)
    except ValueError:
        raise IDAError(f"Invalid number: {text}")

    # Estimate the size of the number
    if not size:
        size = 0
        n = abs(value)
        while n:
            size += 1
            n >>= 1
        size += 7
        size //= 8

    # Convert the number to bytes
    try:
        bytes = value.to_bytes(size, "little", signed=True)
    except OverflowError:
        raise IDAError(f"Number {text} is too big for {size} bytes")

    # Convert the bytes to ASCII
    ascii = ""
    for byte in bytes.rstrip(b"\x00"):
        if byte >= 32 and byte <= 126:
            ascii += chr(byte)
        else:
            ascii = None
            break

    return {
        "decimal": str(value),
        "hexadecimal": hex(value),
        "bytes": bytes.hex(" "),
        "ascii": ascii,
        "binary": bin(value)
    }

T = TypeVar("T")

class Page(TypedDict, Generic[T]):
    data: list[T]
    next_offset: Optional[int]

def paginate(data: list[T], offset: int, count: int) -> Page[T]:
    if count == 0:
        count = len(data)
    next_offset = offset + count
    if next_offset >= len(data):
        next_offset = None
    return {
        "data": data[offset:offset+count],
        "next_offset": next_offset,
    }

@jsonrpc
@idaread
def list_functions(
    offset: Annotated[int, "Offset to start listing from (start at 0)"],
    count: Annotated[int, "Number of functions to list (100 is a good default, 0 means remainder)"],
) -> Page[Function]:
    """List all functions in the database (paginated)"""
    functions = [get_function(address) for address in idautils.Functions()]
    return paginate(functions, offset, count)

class String(TypedDict):
    address: str
    length: int
    type: str
    string: str

def get_strings() -> list[String]:
    strings = []
    for item in idautils.Strings():
        string_type = "C" if item.strtype == 0 else "Unicode"
        try:
            string = str(item)
            if string:
                strings.append({
                    "address": hex(item.ea),
                    "length": item.length,
                    "type": string_type,
                    "string": string
                })
        except:
            continue
    return strings

@jsonrpc
@idaread
def list_strings(
    offset: Annotated[int, "Offset to start listing from (start at 0)"],
    count: Annotated[int, "Number of strings to list (100 is a good default, 0 means remainder)"],
) -> Page[String]:
    """List all strings in the database (paginated)"""
    strings = get_strings()
    return paginate(strings, offset, count)

@jsonrpc
@idaread
def search_strings(
        pattern_str: Annotated[str, "The regular expression to match((The generated regular expression includes case by default))"],
        offset: Annotated[int, "Offset to start listing from (start at 0)"],
        count: Annotated[int, "Number of strings to list (100 is a good default, 0 means remainder)"],
) -> Page[String]:
    """Search for strings that satisfy a regular expression"""
    strings = get_strings()
    try:
        pattern = re.compile(pattern_str)
    except Exception as e:
        raise ValueError(f"Regular expression syntax error, reason is {e}")
    try:
        matched_strings = [s for s in strings if s["string"] and re.search(pattern, s["string"])]
    except Exception as e:
        raise ValueError(f"The regular match failed, reason is {e}")
    return paginate(matched_strings, offset, count)

@jsonrpc
@idaread
def search_strings(
    pattern: Annotated[str, "Substring to search for in strings"],
    offset: Annotated[int, "Offset to start listing from (start at 0)"],
    count: Annotated[int, "Number of strings to list (100 is a good default, 0 means remainder)"],
) -> Page[String]:
    """Search for strings containing the given pattern (case-insensitive)"""
    strings = get_strings()
    matched_strings = [s for s in strings if pattern.lower() in s["string"].lower()]
    return paginate(matched_strings, offset, count)



def decompile_checked(address: int) -> ida_hexrays.cfunc_t:
    if not ida_hexrays.init_hexrays_plugin():
        raise IDAError("Hex-Rays decompiler is not available")
    error = ida_hexrays.hexrays_failure_t()
    cfunc: ida_hexrays.cfunc_t = ida_hexrays.decompile_func(address, error, ida_hexrays.DECOMP_WARNINGS)
    if not cfunc:
        message = f"Decompilation failed at {hex(address)}"
        if error.str:
            message += f": {error.str}"
        if error.errea != idaapi.BADADDR:
            message += f" (address: {hex(error.errea)})"
        raise IDAError(message)
    return cfunc

@jsonrpc
@idaread
def decompile_function(
    address: Annotated[str, "Address of the function to decompile"]
) -> str:
    """Decompile a function at the given address"""
    address = parse_address(address)
    cfunc = decompile_checked(address)
    if is_window_active():
        ida_hexrays.open_pseudocode(address, ida_hexrays.OPF_REUSE)
    sv = cfunc.get_pseudocode()
    pseudocode = ""
    for i, sl in enumerate(sv):
        sl: ida_kernwin.simpleline_t
        item = ida_hexrays.ctree_item_t()
        addr = None if i > 0 else cfunc.entry_ea
        if cfunc.get_line_item(sl.line, 0, False, None, item, None):
            ds = item.dstr().split(": ")
            if len(ds) == 2:
                try:
                    addr = int(ds[0], 16)
                except ValueError:
                    pass
        line = ida_lines.tag_remove(sl.line)
        if len(pseudocode) > 0:
            pseudocode += "\n"
        if not addr:
            pseudocode += f"/* line: {i} */ {line}"
        else:
            pseudocode += f"/* line: {i}, address: {hex(addr)} */ {line}"

    return pseudocode

@jsonrpc
@idaread
def disassemble_function(
    start_address: Annotated[str, "Address of the function to disassemble"]
) -> str:
    """Get assembly code (address: instruction; comment) for a function"""
    start = parse_address(start_address)
    func = idaapi.get_func(start)
    if not func:
        raise IDAError(f"No function found containing address {start_address}")
    if is_window_active():
        ida_kernwin.jumpto(start)

    # TODO: add labels and limit the maximum number of instructions
    disassembly = ""
    for address in ida_funcs.func_item_iterator_t(func):
        if len(disassembly) > 0:
            disassembly += "\n"
        disassembly += f"{hex(address)}: "
        disassembly += idaapi.generate_disasm_line(address, idaapi.GENDSM_REMOVE_TAGS)
        comment = idaapi.get_cmt(address, False)
        if not comment:
            comment = idaapi.get_cmt(address, True)
        if comment:
            disassembly += f"; {comment}"
    return disassembly

class Xref(TypedDict):
    address: str
    type: str
    function: Optional[Function]

@jsonrpc
@idaread
def get_xrefs_to(
    address: Annotated[str, "Address to get cross references to"]
) -> list[Xref]:
    """Get all cross references to the given address"""
    xrefs = []
    xref: ida_xref.xrefblk_t
    for xref in idautils.XrefsTo(parse_address(address)):
        xrefs.append({
            "address": hex(xref.frm),
            "type": "code" if xref.iscode else "data",
            "function": get_function(xref.frm, raise_error=False),
        })
    return xrefs

@jsonrpc
@idaread
def get_entry_points() -> list[Function]:
    """Get all entry points in the database"""
    result = []
    for i in range(ida_entry.get_entry_qty()):
        ordinal = ida_entry.get_entry_ordinal(i)
        address = ida_entry.get_entry(ordinal)
        func = get_function(address, raise_error=False)
        if func is not None:
            result.append(func)
    return result

@jsonrpc
@idawrite
def set_comment(
    address: Annotated[str, "Address in the function to set the comment for"],
    comment: Annotated[str, "Comment text"]
):
    """Set a comment for a given address in the function disassembly and pseudocode"""
    address = parse_address(address)

    if not idaapi.set_cmt(address, comment, False):
        raise IDAError(f"Failed to set disassembly comment at {hex(address)}")

    # Reference: https://cyber.wtf/2019/03/22/using-ida-python-to-analyze-trickbot/
    # Check if the address corresponds to a line
    cfunc = decompile_checked(address)

    # Special case for function entry comments
    if address == cfunc.entry_ea:
        idc.set_func_cmt(address, comment, True)
        cfunc.refresh_func_ctext()
        return

    eamap = cfunc.get_eamap()
    if address not in eamap:
        print(f"Failed to set decompiler comment at {hex(address)}")
        return
    nearest_ea = eamap[address][0].ea

    # Remove existing orphan comments
    if cfunc.has_orphan_cmts():
        cfunc.del_orphan_cmts()
        cfunc.save_user_cmts()

    # Set the comment by trying all possible item types
    tl = idaapi.treeloc_t()
    tl.ea = nearest_ea
    for itp in range(idaapi.ITP_SEMI, idaapi.ITP_COLON):
        tl.itp = itp
        cfunc.set_user_cmt(tl, comment)
        cfunc.save_user_cmts()
        cfunc.refresh_func_ctext()
        if not cfunc.has_orphan_cmts():
            return
        cfunc.del_orphan_cmts()
        cfunc.save_user_cmts()
    print(f"Failed to set decompiler comment at {hex(address)}")

def refresh_decompiler_widget():
    widget = ida_kernwin.get_current_widget()
    if widget is not None:
        vu = ida_hexrays.get_widget_vdui(widget)
        if vu is not None:
            vu.refresh_ctext()

def refresh_decompiler_ctext(function_address: int):
    error = ida_hexrays.hexrays_failure_t()
    cfunc: ida_hexrays.cfunc_t = ida_hexrays.decompile_func(function_address, error, ida_hexrays.DECOMP_WARNINGS)
    if cfunc:
        cfunc.refresh_func_ctext()

@jsonrpc
@idawrite
def rename_local_variable(
    function_address: Annotated[str, "Address of the function containing the variable"],
    old_name: Annotated[str, "Current name of the variable"],
    new_name: Annotated[str, "New name for the variable (empty for a default name)"]
):
    """Rename a local variable in a function"""
    func = idaapi.get_func(parse_address(function_address))
    if not func:
        raise IDAError(f"No function found at address {function_address}")
    if not ida_hexrays.rename_lvar(func.start_ea, old_name, new_name):
        raise IDAError(f"Failed to rename local variable {old_name} in function {hex(func.start_ea)}")
    refresh_decompiler_ctext(func.start_ea)

@jsonrpc
@idawrite
def rename_global_variable(
    old_name: Annotated[str, "Current name of the global variable"],
    new_name: Annotated[str, "New name for the global variable (empty for a default name)"]
):
    """Rename a global variable"""
    ea = idaapi.get_name_ea(idaapi.BADADDR, old_name)
    if not idaapi.set_name(ea, new_name):
        raise IDAError(f"Failed to rename global variable {old_name} to {new_name}")
    refresh_decompiler_ctext(ea)

@jsonrpc
@idawrite
def set_global_variable_type(
    variable_name: Annotated[str, "Name of the global variable"],
    new_type: Annotated[str, "New type for the variable"]
):
    """Set a global variable's type"""
    ea = idaapi.get_name_ea(idaapi.BADADDR, variable_name)
    tif = ida_typeinf.tinfo_t(new_type, None, ida_typeinf.PT_SIL)
    if not tif:
        raise IDAError(f"Parsed declaration is not a variable type")
    if not ida_typeinf.apply_tinfo(ea, tif, ida_typeinf.PT_SIL):
        raise IDAError(f"Failed to apply type")

@jsonrpc
@idawrite
def rename_function(
    function_address: Annotated[str, "Address of the function to rename"],
    new_name: Annotated[str, "New name for the function (empty for a default name)"]
):
    """Rename a function"""
    func = idaapi.get_func(parse_address(function_address))
    if not func:
        raise IDAError(f"No function found at address {function_address}")
    if not idaapi.set_name(func.start_ea, new_name):
        raise IDAError(f"Failed to rename function {hex(func.start_ea)} to {new_name}")
    refresh_decompiler_ctext(func.start_ea)

@jsonrpc
@idawrite
def set_function_prototype(
    function_address: Annotated[str, "Address of the function"],
    prototype: Annotated[str, "New function prototype"]
) -> str:
    """Set a function's prototype"""
    func = idaapi.get_func(parse_address(function_address))
    if not func:
        raise IDAError(f"No function found at address {function_address}")
    try:
        tif = ida_typeinf.tinfo_t(prototype, None, ida_typeinf.PT_SIL)
        if not tif.is_func():
            raise IDAError(f"Parsed declaration is not a function type")
        if not ida_typeinf.apply_tinfo(func.start_ea, tif, ida_typeinf.PT_SIL):
            raise IDAError(f"Failed to apply type")
        refresh_decompiler_ctext(func.start_ea)
    except Exception as e:
        raise IDAError(f"Failed to parse prototype string: {prototype}")

class my_modifier_t(ida_hexrays.user_lvar_modifier_t):
    def __init__(self, var_name: str, new_type: ida_typeinf.tinfo_t):
        ida_hexrays.user_lvar_modifier_t.__init__(self)
        self.var_name = var_name
        self.new_type = new_type

    def modify_lvars(self, lvars):
        for lvar_saved in lvars.lvvec:
            lvar_saved: ida_hexrays.lvar_saved_info_t
            if lvar_saved.name == self.var_name:
                lvar_saved.type = self.new_type
                return True
        return False

# NOTE: This is extremely hacky, but necessary to get errors out of IDA
def parse_decls_ctypes(decls: str, hti_flags: int) -> tuple[int, str]:
    if sys.platform == "win32":
        import ctypes
        assert isinstance(decls, str), "decls must be a string"
        assert isinstance(hti_flags, int), "hti_flags must be an int"
        c_decls = decls.encode("utf-8")
        c_til = None
        ida_dll = ctypes.CDLL("ida")
        ida_dll.parse_decls.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_int]
        ida_dll.parse_decls.restype = ctypes.c_int

        messages = []
        @ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_char_p, ctypes.c_char_p)
        def magic_printer(fmt: bytes, arg1: bytes):
            if fmt.count(b"%") == 1 and b"%s" in fmt:
                formatted = fmt.replace(b"%s", arg1)
                messages.append(formatted.decode("utf-8"))
                return len(formatted) + 1
            else:
                messages.append(f"unsupported magic_printer fmt: {repr(fmt)}")
                return 0

        errors = ida_dll.parse_decls(c_til, c_decls, magic_printer, hti_flags)
    else:
        # NOTE: The approach above could also work on other platforms, but it's
        # not been tested and there are differences in the vararg ABIs.
        errors = ida_typeinf.parse_decls(None, decls, False, hti_flags)
        messages = []
    return errors, messages

@jsonrpc
@idawrite
def declare_c_type(
    c_declaration: Annotated[str, "C declaration of the type. Examples include: typedef int foo_t; struct bar { int a; bool b; };"],
):
    """Create or update a local type from a C declaration"""
    # PT_SIL: Suppress warning dialogs (although it seems unnecessary here)
    # PT_EMPTY: Allow empty types (also unnecessary?)
    # PT_TYP: Print back status messages with struct tags
    flags = ida_typeinf.PT_SIL | ida_typeinf.PT_EMPTY | ida_typeinf.PT_TYP
    errors, messages = parse_decls_ctypes(c_declaration, flags)

    pretty_messages = "\n".join(messages)
    if errors > 0:
        raise IDAError(f"Failed to parse type:\n{c_declaration}\n\nErrors:\n{pretty_messages}")
    return f"success\n\nInfo:\n{pretty_messages}"

@jsonrpc
@idawrite
def set_local_variable_type(
    function_address: Annotated[str, "Address of the function containing the variable"],
    variable_name: Annotated[str, "Name of the variable"],
    new_type: Annotated[str, "New type for the variable"]
):
    """Set a local variable's type"""
    try:
        # Some versions of IDA don't support this constructor
        new_tif = ida_typeinf.tinfo_t(new_type, None, ida_typeinf.PT_SIL)
    except Exception:
        try:
            new_tif = ida_typeinf.tinfo_t()
            # parse_decl requires semicolon for the type
            ida_typeinf.parse_decl(new_tif, None, new_type+";", ida_typeinf.PT_SIL)
        except Exception:
            raise IDAError(f"Failed to parse type: {new_type}")
    func = idaapi.get_func(parse_address(function_address))
    if not func:
        raise IDAError(f"No function found at address {function_address}")
    if not ida_hexrays.rename_lvar(func.start_ea, variable_name, variable_name):
        raise IDAError(f"Failed to find local variable: {variable_name}")
    modifier = my_modifier_t(variable_name, new_tif)
    if not ida_hexrays.modify_user_lvars(func.start_ea, modifier):
        raise IDAError(f"Failed to modify local variable: {variable_name}")
    refresh_decompiler_ctext(func.start_ea)

class MCP(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "MCP Plugin"
    help = "MCP"
    wanted_name = "MCP"
    wanted_hotkey = "Ctrl-Alt-M"

    def init(self):
        self.server = Server()
        hotkey = MCP.wanted_hotkey.replace("-", "+")
        if sys.platform == "darwin":
            hotkey = hotkey.replace("Alt", "Option")
        print(f"[MCP] Plugin loaded, use Edit -> Plugins -> MCP ({hotkey}) to start the server")
        return idaapi.PLUGIN_KEEP

    def run(self, args):
        self.server.start()

    def term(self):
        self.server.stop()

def PLUGIN_ENTRY():
    return MCP()


repository: AlexAltea/ida
file: plugins/alleycat/alleycat.py
description: Collection of IDA Python plugins/scripts/modules
wanted_name: AlleyCat
created_at: 2017-12-23 02:45:40+00:00
pushed_at: 2017-12-24 15:07:38+00:00
url:
https://github.com/AlexAltea/ida/blob/79ce319e729ea65a02adca9985d9150a4b681b0c/plugins/alleycat/alleycat.py

content:
import idc
import idaapi
import idautils
import time

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

class AlleyCatException(Exception):
    pass

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

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

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

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

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

    def _name(self, ea):
        name = idc.Name(ea)
        if not name:
            name = idc.GetFuncOffset(ea)
            if not name:
                name = '0x%X' % ea
        return name

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

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

        # Loop while there are still unresolve paths and while all path sizes have not exceeded ALLEYCAT_LIMIT
        while partial_paths and len(self.paths) < self.limit and len(partial_paths) < self.limit:
            # Initialize a unique set of callers for this iteration
            callers = set()

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

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

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

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

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

class AlleyCatFunctionPaths(AlleyCat):

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

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

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

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

class AlleyCatCodePaths(AlleyCat):

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

        if not start_func:
            raise AlleyCatException("Address 0x%X is not part of a function!" % start_ea)
        if not end_func:
            raise AlleyCatException("Address 0x%X is not part of a function!" % end_ea)
        if start_func.startEA != end_func.startEA:
            raise AlleyCatException("The start and end addresses are not part of the same function!")

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

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

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

        super(AlleyCatCodePaths, self).__init__(start_block.startEA, end_block.startEA, quiet)

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


### Everything below here is just IDA UI/Plugin stuff ###


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

    INCLUDE_ACTION = 0
    EXCLUDE_ACTION = 1

    def __init__(self):
        self.reset()

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

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

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

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

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

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

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

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

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

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

class AlleyCatGraph(idaapi.GraphViewer):
    '''
    Displays the graph and manages graph actions.
    '''

    def __init__(self, results, title="AlleyCat Graph"):
        idaapi.GraphViewer.__init__(self, title)
        self.results = results

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

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

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

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

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

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

        for path in self.results:
            parent_node = None

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

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

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

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

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

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

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

        return True

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

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

        return (self[node_id], color)

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

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

        return hint

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

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

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

        if self.edges.has_key(node_id):
            for edge_node_id in self.edges[node_id]:

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

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

        if xref_locations:
            xref_locations.sort()

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

            idc.Jump(xref_locations[0][0])
        else:
            idc.Jump(node_ea)

    def OnClose(self):
        # TODO: Add a 'do not ask again' feature?
        if idc.AskYN(1, "Path nodes have been highlighted in the disassembly window. Undo highlighting?") == 1:
            self.unhighlight_all()

    def match_xref_source(self, xref, source):
        # TODO: This must be modified if support for graphing function blocks is added.
        return ((xref.type != idc.fl_F) and (idc.GetFunctionAttr(xref.frm, idc.FUNCATTR_START) == source))

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

        @name - Location name

        Returns the address of the named location, or idc.BADADDR on failure.
        '''
        # This allows support of the function offset style names (e.g., main+0C)
        # TODO: Is there something in the IDA API that does this already??
        if '+' in name:
            (func_name, offset) = name.split('+')
            base_ea = idc.LocByName(func_name)
            if base_ea != idc.BADADDR:
                try:
                    ea = base_ea + int(offset, 16)
                except:
                    ea = idc.BADADDR
        else:
            ea = idc.LocByName(name)
            if ea == idc.BADADDR:
                try:
                    ea = int(name, 0)
                except:
                    ea = idc.BADADDR

        return ea

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

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

        @ea - Address.

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

    def colorize_node(self, ea, color):
        # Colorizes an entire code block
        func = idaapi.get_func(ea)
        if func:
            for block in idaapi.FlowChart(func):
                if block.startEA <= ea and block.endEA > ea:
                    ea = block.startEA
                    while ea < block.endEA:
                        idaapi.set_item_color(ea, color)
                        ea = idc.NextHead(ea)
                    break

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

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

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

# Helpers
def current_function():
    return idaapi.get_func(ScreenEA()).startEA

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

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

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

    if results:
        return AlleyCatGraph(results, 'Path Graph')
    return None

def get_user_selected_functions(many=False):
    functions = []
    ea = idc.ScreenEA()
    try:
        current_function = idc.GetFunctionAttr(ea, idc.FUNCATTR_START)
    except:
        current_function = None

    while True:
        function = idc.ChooseFunction("Select a function and click 'OK' until all functions have been selected. When finished, click 'Cancel' to display the graph.")
        # ChooseFunction automatically jumps to the selected function
        # if the enter key is pressed instead of clicking 'OK'. Annoying.
        if idc.ScreenEA() != ea:
            idc.Jump(ea)

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

        if not many:
            break

    return functions

# Handlers
class FindPathsFromManyHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
        self.graph = None

    def activate(self, arg):
        if self.graph:
            self.graph.Close()
        target = current_function()
        if target:
            sources = get_user_selected_functions(many=True)
            if sources:
                self.graph = find_and_plot_paths(sources, [target])
        if self.graph:
            self.graph.Show()

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

class FindPathsToManyHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
        self.graph = None
        
    def activate(self, ctx):
        if self.graph:
            self.graph.Close()
        source = current_function()
        if source:
            targets = get_user_selected_functions(many=True)
            if targets:
                self.graph = find_and_plot_paths([source], targets)
        if self.graph:
            self.graph.Show()

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

class FindPathsToCodeBlockHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
        self.graph = None

    def activate(self, ctx):
        if self.graph:
            self.graph.Close()
        target = idc.ScreenEA()
        source = current_function()
        if source:
            self.graph = find_and_plot_paths(
                [source], [target], klass=AlleyCatCodePaths)
        if self.graph:
            self.graph.Show()

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

# Actions
action_alleycat_from = 'alleycat:from'
action_alleycat_to = 'alleycat:to'
action_alleycat_block = 'alleycat:block'

class idapathfinder_t(idaapi.plugin_t):

    flags = 0
    comment = ''
    help = ''
    wanted_name = 'AlleyCat'
    wanted_hotkey = ''

    def init(self):
        ui_path = "View/Graphs/"
        self.menu_contexts = []

        # IDA 7.x
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.register_action(idaapi.action_desc_t(
                action_alleycat_from,
                'Find paths to the current function from...',
                FindPathsFromManyHandler(),
                None,
                'Find paths from other function to the current function'))
            idaapi.attach_action_to_menu(ui_path,
                action_alleycat_from, idaapi.SETMENU_APP)
             
            idaapi.register_action(idaapi.action_desc_t(
                action_alleycat_to,
                'Find paths from the current function to...',
                FindPathsToManyHandler(),
                None,
                'Find paths from the current function to other function'))
            idaapi.attach_action_to_menu(ui_path,
                action_alleycat_to, idaapi.SETMENU_APP)
            
            idaapi.register_action(idaapi.action_desc_t(
                action_alleycat_block,
                'Find paths in the current function to the current code block',
                FindPathsToCodeBlockHandler(),
                None,
                'Find paths in the current function to the current code block'))
            idaapi.attach_action_to_menu(ui_path,
                action_alleycat_block, idaapi.SETMENU_APP)
        # IDA 6.x
        else:
            self.menu_contexts.append(idaapi.add_menu_item(ui_path,
                "Find paths to the current function from...",
                "",
                0,
                self.FindPathsFromMany,
                (None,)))
            self.menu_contexts.append(idaapi.add_menu_item(ui_path,
                "Find paths from the current function to...",
                "",
                0,
                self.FindPathsToMany,
                (None,)))
            self.menu_contexts.append(idaapi.add_menu_item(ui_path,
                "Find paths in the current function to the current code block",
                "",
                0,
                self.FindPathsToCodeBlock,
                (None,)))
        
        return idaapi.PLUGIN_KEEP

    def term(self):
        # IDA 7.x
        if idaapi.IDA_SDK_VERSION >= 700:
            if idaapi.unregister_action(action_alleycat_from) is None:
                print "Failed to unregister action."
            if idaapi.unregister_action(action_alleycat_to) is None:
                print "Failed to unregister action."
            if idaapi.unregister_action(action_alleycat_block) is None:
                print "Failed to unregister action."
        # IDA 6.x
        else:
            for context in self.menu_contexts:
                idaapi.del_menu_item(context)
        return None

    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return idapathfinder_t()


repository: zmrbak/IDAPython
file: 配套代码/ida10/python/examples/hexrays/vds6.py
description: 在线课程《IDAPython脚本化软件逆向分析入门》课程源码
wanted_name: Hex-Rays space remover (IDAPython)
comment: Sample plugin6 for Hex-Rays decompiler
created_at: 2022-05-02 02:13:34+00:00
pushed_at: 2022-05-14 10:41:43+00:00
url:
https://github.com/zmrbak/IDAPython/blob/478278339ba0c150b3ace657296233d8a89b26ea/%E9%85%8D%E5%A5%97%E4%BB%A3%E7%A0%81/ida10/python/examples/hexrays/vds6.py

content:
"""
summary: superficially modify the decompilation output

description:
  modifies the decompilation output in a superficial manner,
  by removing some white spaces

  Note: this is rather crude, not quite "pythonic" code.
"""

from __future__ import print_function

import idautils
import idc
import ida_idaapi
import ida_hexrays
import ida_lines

def dbg(msg):
    #print(msg)
    pass

def is_cident_char(c):
    return c.isalnum() or c == '_'

def my_tag_skipcodes(l, storage):
    n = ida_lines.tag_skipcodes(l)
    dbg("Skipping %d chars ('%s')" % (n, l[0:n]))
    storage.append(l[0:n])
    return l[n:]

def remove_spaces(sl):
    dbg("*" * 80)
    l = sl.line
    out = []

    def push(c):
        dbg("Appending '%s'" % c)
        out.append(c)

    # skip initial spaces, do not compress them
    while True:
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        if not c.isspace():
            break
        push(c)
        l = l[1:]

    # remove all spaces except in string and char constants
    delim = None # if not None, then we are skipping until 'delim'
    last = None # last seen character
    while True:
        # go until comments
        if l.startswith("//"):
            push(l)
            break
        dbg("-" * 60)
        nchars = ida_lines.tag_advance(l, 1)
        push(l[0:nchars])
        l = l[nchars:]
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        dbg("c: '%s', last: '%s', l: '%s'" % (c, last, l))
        if delim:
            # we're inside a literal.
            if c == delim:
                delim = None # literal ended
        elif c == '"' or c == "'":
            delim = c # string/char literal started
        elif c.isspace():
            end = l.lstrip()
            nptr = my_tag_skipcodes(end, out)
            dbg("end: '%s', nptr: '%s'" % (end, nptr))
            # do not concatenate idents
            if not is_cident_char(last) or not is_cident_char(nptr[0]):
                l = end
                c = l[0] if l else ''
                dbg("new l: '%s'" % l)
        last = l[0] if l else ''

    sl.line = "".join(out)


class vds6_hooks_t(ida_hexrays.Hexrays_Hooks):
    def func_printed(self, cfunc):
        for sl in cfunc.get_pseudocode():
            remove_spaces(sl);
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays space remover (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin6 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.vds6_hooks = vds6_hooks_t()
            self.vds6_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds6_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: intbjw/gen_frida
file: gen_frida.py
description: IDA插件，用于生成 frida hook 代码
url: https://github.com/intbjw/gen_frida/blob/c2be1e2b97904ba3d1a86d2d54063659edb0f3be/gen_frida.py
wanted_name: GenFridaCode
comment: frida辅助插件
created_at: 2025-01-10 03:27:03+00:00
pushed_at: 2025-01-10 03:30:45+00:00
content:
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ida_hexrays import cfunc_t
    from ida_kernwin import view_mouse_event_t

import idc
import idaapi
import ida_lines
import ida_ida

from ida_idaapi import plugin_t

hook_function_template = '''
function hook_{function_name}() {{
    let base_addr = Module.findBaseAddress("{so_name}");

    Interceptor.attach(base_addr.add(0x{offset:X}), {{
        onEnter: function(args) {{
            console.log(`onEnter {function_name}{args}`);
        }}, onLeave: function(retval) {{
            console.log(`onLeave {function_name}{result}`);
        }}
    }});
}}
'''

inline_hook_template = '''
function hook_0x{offset:X}() {{
    let base_addr = Module.findBaseAddress("{so_name}");

    Interceptor.attach(base_addr.add(0x{offset:X}), {{
        onEnter(args) {{
            console.log(`call 0x{offset:X} {args}`);
        }}
    }});
}}
'''

dlopen_after_template = '''
let android_dlopen_ext = Module.findExportByName(null, "android_dlopen_ext");
if(android_dlopen_ext != null) {{
    Interceptor.attach(android_dlopen_ext, {{
        onEnter: function(args) {{
            let so_name = args[0].readCString();
            if(so_name.indexOf("{so_name}") !== -1) {{
                this.hook = true;
            }}
        }}, onLeave: function(retval) {{
            if(this.hook) {{
                this.hook = false;
                dlopen_todo();
            }}
        }}
    }});
}}

function dlopen_todo() {{
    //todo
}}
'''

init_template = '''
function hook_init() {{
    let linker_name = "linker64";
    let already_hook = false;
    let call_constructor_addr = null;
    if (Process.arch.endsWith("arm")) {{
        linker_name = "linker";
    }}

    let symbols = Module.enumerateSymbolsSync(linker_name);
    for (let i = 0; i < symbols.length; i++) {{
        let symbol = symbols[i];
        if (symbol.name.indexOf("call_constructor") !== -1) {{
            call_constructor_addr = symbol.address;
        }}
    }}

    if (call_constructor_addr != null) {{
        console.log(`get construct address ${{call_constructor_addr}}`);
        Interceptor.attach(call_constructor_addr, {{
            onEnter: function (args) {{
                if(already_hook === false) {{
                    const targetModule = Process.findModuleByName("{so_name}");
                    if (targetModule !== null) {{
                        already_hook = true;
                        init_todo();
                    }}
                }}
            }}
        }});
    }}
}}

function init_todo() {{
    //todo
}}
'''

dump_template = '''
function dump_0x{offset:X}() {{
    let base_addr = Module.findBaseAddress("{so_name}");
    let dump_addr = base_addr.add(0x{offset:X});
    let dump_mem = hexdump(dump_addr, {{length: {length:#x}}});
    console.log(`dump {so_name} + {offset:#x}:\\n${{dump_mem}}`);
}}
'''


def generate_print_args(addr: int):
    args_num = get_args_num(addr)
    tmp = []
    for index in range(args_num):
        tmp.append(f'arg{index}:${{args[{index}]}}')
    return ' '.join(tmp)


def generate_for_func(so_name: str, function_name: str, addr: int):
    # 根据参数个数打印
    args_print = generate_print_args(addr)
    if args_print != '':
        args_print = ' ' + args_print
    # 根据是否有返回值判断是否打印retval
    if has_return(addr):
        ret_print = ' ' + '${retval}'
    else:
        ret_print = ''

    result = hook_function_template.format_map({
        'so_name': so_name,
        'function_name': function_name,
        'offset': get_offset(addr),
        'args': args_print,
        'result': ret_print
    })
    print(result)


def get_offset(addr: int) -> int:
    if ida_ida.idainfo_is_64bit():
        return addr
    else:
        # 可以通过T标志位判断某条指令是 ARM 还是 THUMB
        # ARM 返回 0
        # THUMB 返回 1
        return addr + idc.get_sreg(addr, 'T')


def generate_for_inline(so_name: str, addr: int):
    args_print = '${JSON.stringify(this.context)}'
    if idaapi.is_call_insn(addr):
        # 获取指定索引操作数
        operand = idc.print_operand(addr, 0)
        call_addr = idaapi.get_name_ea(0, operand)
        if call_addr != idaapi.BADADDR:
            # 解析 call 的函数其地址
            # 获取指定索引操作数中的值
            addr = idc.get_operand_value(addr, 0)
            args_print = generate_print_args(addr)
    print(inline_hook_template.format_map(
        {'so_name': so_name, 'offset': get_offset(addr), 'args': args_print}))


def has_return(addr: int) -> bool:
    cfun = idaapi.decompile(addr)  # type: cfunc_t
    has_return = True
    dcl = ida_lines.tag_remove(cfun.print_dcl())  # type: str
    if dcl.startswith('void ') and not dcl.startswith('void *'):
        has_return = False
    return has_return


def get_args_num(addr: int) -> int:
    return len(idaapi.decompile(addr).arguments)


def generate_for_func_by_addr(addr: int):
    so_name = idaapi.get_root_filename()
    function_name = idaapi.get_func_name(addr).replace('.', 'dot_')
    generate_for_func(so_name, function_name, addr)


def generate_for_inline_by_addr(addr: int):
    so_name = idaapi.get_root_filename()
    generate_for_inline(so_name, addr)


def generate_snippet():
    '''
    脚本生成入口
    '''
    # 获取光标所处地址
    addr = idaapi.get_screen_ea()  # type: int
    # 通过获取函数属性取到函数首地址
    start_addr = idc.get_func_attr(addr, idc.FUNCATTR_START)
    if start_addr == addr:
        generate_for_func_by_addr(addr)
    elif start_addr == idc.BADADDR:
        print('不在函数内')
    else:
        generate_for_inline_by_addr(addr)


def generate_init_code():
    so_name = idaapi.get_root_filename()
    print(dlopen_after_template.format_map({'so_name': so_name}))
    print(init_template.format_map({'so_name': so_name}))


def generate_dump_script(start: int, length: int):
    so_name = idaapi.get_root_filename()
    print(dump_template.format_map(
        {'so_name': so_name, 'offset': start, 'length': length}))


class MyViewHooks(idaapi.View_Hooks):

    def view_dblclick(self, view, event: 'view_mouse_event_t'):
        '''
        在汇编界面双击地址
        '''
        widget_type = idaapi.get_widget_type(view)
        if widget_type == idaapi.BWN_DISASM:
            global initialized
            if not initialized:
                initialized = True
                generate_init_code()
            generate_snippet()

    def view_click(self, view, event: 'view_mouse_event_t'):
        '''
        在汇编界面选中一个范围后松开鼠标
        '''
        widget_type = idaapi.get_widget_type(view)
        if widget_type == idaapi.BWN_DISASM:
            start = idc.read_selection_start()
            end = idc.read_selection_end()
            if (start != idaapi.BADADDR) and (end != idaapi.BADADDR):
                length = end - start
                generate_dump_script(start, length)


class GenFridaSnippetHandler(idaapi.action_handler_t):
    '''
    动作处理类，和右键菜单绑定
    - https://hex-rays.com/blog/augmenting-ida-ui-with-your-own-actions/
    '''

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

    def activate(self, ctx):
        '''
        点了右键菜单之后会调用这个函数
        '''
        generate_snippet()
        return 1

    def update(self, ctx):
        '''
        AST_ENABLE_ALWAYS 表示可以点击
        '''
        return idaapi.AST_ENABLE_ALWAYS


class GenFridaInitHandler(idaapi.action_handler_t):

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

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

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


class GenFridaDumpHandler(idaapi.action_handler_t):

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

    def activate(self, ctx):
        start = idc.read_selection_start()
        end = idc.read_selection_end()
        if (start != idaapi.BADADDR) and (end != idaapi.BADADDR):
            length = end - start
            generate_dump_script(start, length)
        return 1

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


class GenFridaPlugin(plugin_t):
    # 关于插件的注释
    # 当鼠标浮于菜单插件上方时，IDA左下角所示
    comment = 'frida辅助插件'
    # 帮助信息，我们选择不填
    help = '右键选择生成hook代码后找地方粘贴即可'
    # 插件的特性，是一直在内存里，还是运行一下就退出，等等
    flags = idaapi.PLUGIN_KEEP
    # 插件的名字
    wanted_name = 'GenFridaCode'
    # 快捷键，我们选择置空不弄
    wanted_hotkey = ''

    # 插件刚被加载到IDA内存中 这里适合做插件的初始化工作
    def init(self):
        print('GenFridaCode init')
        # 初始化的时候将动作绑定到菜单
        # 这个函数不支持 kwargs
        # name, label, handler, shortcut=None, tooltip=None, icon=-1, flags=0
        gen_frida_init_action_desc = idaapi.action_desc_t(
            'my:gen_frida_init', '生成frida hook init代码', GenFridaInitHandler(), '', '生成frida hook init代码')
        idaapi.register_action(gen_frida_init_action_desc)
        gen_frida_snippet_action_desc = idaapi.action_desc_t(
            'my:gen_frida_snippet', '生成frida hook代码', GenFridaSnippetHandler(), '', '生成frida hook代码')
        idaapi.register_action(gen_frida_snippet_action_desc)
        gen_frida_dump_action_desc = idaapi.action_desc_t(
            'my:gen_frida_dump', '生成frida hook dump代码', GenFridaDumpHandler(), '', '生成frida hook dump代码')
        idaapi.register_action(gen_frida_dump_action_desc)
        global my_ui_hooks
        my_ui_hooks = MyUIHooks()
        my_ui_hooks.hook()
        return idaapi.PLUGIN_KEEP

    # 插件运行中 这里是主要逻辑
    def run(self, arg: int):
        print('GenFridaCode run')
        global my_view_hooks
        my_view_hooks = MyViewHooks()
        my_view_hooks.hook()

    # 插件卸载退出的时机 这里适合做资源释放
    def term(self):
        print('GenFridaCode term')


class MyUIHooks(idaapi.UI_Hooks):

    def finish_populating_widget_popup(self, widget, popup_handle, ctx):
        if idaapi.get_widget_type(widget) == idaapi.BWN_DISASM:
            idaapi.attach_action_to_popup(
                widget, popup_handle, 'my:gen_frida_init', 'GenFrida/')
            idaapi.attach_action_to_popup(
                widget, popup_handle, 'my:gen_frida_snippet', 'GenFrida/')
            idaapi.attach_action_to_popup(
                widget, popup_handle, 'my:gen_frida_dump', 'GenFrida/')


initialized = False
# 注册插件


def PLUGIN_ENTRY():
    return GenFridaPlugin()

# 原作者 https://github.com/Pr0214 https://t.zsxq.com/05IEynQVV
# 修改者 https://github.com/SeeFlowerX https://t.zsxq.com/05E2VBuNj
# 地址 https://gist.github.com/SeeFlowerX/41b7f45913eba9d1dff6fd47c2502b13

# IDA插件，用于生成 frida hook 代码，放入plugins目录后，手动在插件菜单激活然后右键双击、选中释放；或者直接在汇编界面右键使用，选择GenFrida


repository: taida957789/ida-mcp-server-plugin
file: ida-mcp-server.py
description: IDA Pro Plugin for serving MCP SSE server for cursor / claude
wanted_name: IDA MCP Server
comment: IDA Model Context Protocol Server
created_at: 2025-03-25 06:30:32+00:00
pushed_at: 2025-03-27 09:33:15+00:00
url:
https://github.com/taida957789/ida-mcp-server-plugin/blob/3df58c12b19fe6baf616607fe9fa72198ad3c893/ida-mcp-server.py

content:
import glob
import json
import os
import ida_bytes
import ida_ua
import ida_funcs
import ida_hexrays
import ida_name
import ida_segment
import idautils
import idc
import ida_idaapi
import ida_kernwin
import idaapi
import threading

from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime
from functools import wraps
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.routing import Mount, Route
from starlette.responses import Response
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from mcp.server import FastMCP
import uvicorn


# Initialize FastMCP server for IDA tools
mcp = FastMCP("IDA MCP Server", port=3000)

# 封裝函數執行在主線程的裝飾器
def execute_on_main_thread(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        result = []
        exception = []
        
        def run_function():
            try:
                result.append(f(*args, **kwargs))
            except Exception as e:
                exception.append(e)
            return 0
        
        ida_kernwin.execute_sync(run_function, ida_kernwin.MFF_FAST)
        
        if exception:
            raise exception[0]
        return result[0]
    return wrapper


@mcp.tool()
@execute_on_main_thread
def get_bytes(ea: int, size: int) -> List[int]:
    """Get bytes at specified address.

    Args:
        ea: Effective address to read from
        size: Number of bytes to read
    """
    try:
        result = [ida_bytes.get_byte(ea + i) for i in range(size)]
        return result
    except Exception as e:
        print(f"Error in get_bytes: {str(e)}")
        return {"error": str(e)}


@mcp.tool()
@execute_on_main_thread
def get_disasm(ea: int) -> str:
    """Get disassembly at specified address.

    Args:
        ea: Effective address to disassemble
    """
    return idc.generate_disasm_line(ea, 0)


@mcp.tool()
@execute_on_main_thread
def get_decompiled_func(ea: int) -> Dict[str, Any]:
    """Get decompiled pseudocode of function containing address.

    Args:
        ea: Effective address within the function
    """
    try:
        func = ida_funcs.get_func(ea)
        if not func:
            return {"error": "No function found at address"}

        decompiler = ida_hexrays.decompile(func.start_ea)
        if not decompiler:
            return {"error": "Failed to decompile function"}

        return {"code": str(decompiler)}
    except Exception as e:
        return {"error": str(e)}


@mcp.tool()
@execute_on_main_thread
def get_function_name(ea: int) -> str:
    """Get function name at specified address.

    Args:
        ea: Effective address of the function
    """
    return ida_name.get_name(ea)


@mcp.tool()
@execute_on_main_thread
def get_segments() -> List[Dict[str, Any]]:
    """Get all segments information.

    @return: List of segments (start, end, name, class, perm, bitness, align, comb, type, sel, flags)
    """
    segments = []
    n = 0
    seg = ida_segment.getnseg(n)
    while seg:
        segments.append(
            {
                "start": seg.start_ea,
                "end": seg.end_ea,
                "name": ida_segment.get_segm_name(seg),
                "class": ida_segment.get_segm_class(seg),
                "perm": seg.perm,
                "bitness": seg.bitness,
                "align": seg.align,
                "comb": seg.comb,
                "type": seg.type,
                "sel": seg.sel,
                "flags": seg.flags,
            }
        )
        n += 1
        seg = ida_segment.getnseg(n)
    return segments


@mcp.tool()
@execute_on_main_thread
def get_functions() -> List[Dict[str, Any]]:
    """Get all functions in the binary."""
    functions = []
    for func_ea in idautils.Functions():
        func_name = ida_name.get_name(func_ea)
        functions.append({"address": func_ea, "name": func_name})
    return functions


@mcp.tool()
@execute_on_main_thread
def get_xrefs_to(ea: int) -> List[Dict[str, Any]]:
    """Get all cross references to specified address.

    Args:
        ea: Effective address to find references to
    """
    xrefs = []
    for xref in idautils.XrefsTo(ea, 0):
        xrefs.append({"from": xref.frm, "type": xref.type})
    return xrefs


@mcp.tool()
@execute_on_main_thread
def get_imports() -> dict[str, list[tuple[int, str, int]]]:
    """Get all imports in the binary.

    Args:
        None

    Returns:
        A dictionary where the keys are module names and the values are lists of tuples.
        Each tuple contains the address of the imported function, the name of the function,
        and the ordinal value of the function.
    """
    tree = {}
    nimps = idaapi.get_import_module_qty()

    for i in range(0, nimps):
        name = idaapi.get_import_module_name(i)
        if not name:
            continue
        # Create a list for imported names
        items = []

        def imports_names_cb(ea, name, ord):
            items.append((ea, "" if not name else name, ord))
            # True -> Continue enumeration
            return True

        # Enum imported entries in this module
        idaapi.enum_import_names(i, imports_names_cb)

        if name not in tree:
            tree[name] = []
        tree[name].extend(items)

    return tree


@mcp.tool()
@execute_on_main_thread
def get_exports() -> List[Tuple[int, int, int, str]]:
    """Get all exports in the binary.

    @return: List of tuples (index, ordinal, ea, name)
    """
    return list(idautils.Entries())


@mcp.tool()
@execute_on_main_thread
def get_entry_point() -> int:
    """Get the entry point of the binary."""
    try:
        import ida_ida
        return ida_ida.inf_get_start_ea()
    except (ImportError, AttributeError):
        try:
            # Alternative method: idc.get_inf_attr to get
            import idc
            return idc.get_inf_attr(idc.INF_START_EA)
        except (ImportError, AttributeError):
            # Last alternative method: use cvar.inf
            return idaapi.cvar.inf.start_ea


@mcp.tool()
@execute_on_main_thread
def make_function(ea: int) -> None:
    """Make a function at specified address."""
    ida_funcs.add_func(ea)


@mcp.tool()
@execute_on_main_thread
def undefine_function(ea: int) -> None:
    """Undefine a function at specified address."""
    ida_funcs.del_func(ea)


@mcp.tool()
@execute_on_main_thread
def get_dword_at(ea: int) -> int:
    """Get the dword at specified address."""
    return idc.get_dword(ea)


@mcp.tool()
@execute_on_main_thread
def get_word_at(ea: int) -> int:
    """Get the word at specified address."""
    return idc.get_word(ea)


@mcp.tool()
@execute_on_main_thread
def get_byte_at(ea: int) -> int:
    """Get the byte at specified address."""
    return idc.get_byte(ea)


@mcp.tool()
@execute_on_main_thread
def get_qword_at(ea: int) -> int:
    """Get the qword at specified address."""
    return idc.get_qword(ea)


@mcp.tool()
@execute_on_main_thread
def get_float_at(ea: int) -> float:
    """Get the float at specified address."""
    return idc.get_float(ea)


@mcp.tool()
@execute_on_main_thread
def get_double_at(ea: int) -> float:
    """Get the double at specified address."""
    return idc.get_double(ea)


@mcp.tool()
@execute_on_main_thread
def get_string_at(ea: int) -> str:
    """Get the string at specified address."""
    return idc.get_strlit_contents(ea)


@mcp.tool()
@execute_on_main_thread
def get_strings():
    strings = []
    for s in idautils.Strings():
        strings.append({"address": s.ea, "string": str(s)})
    return strings

@mcp.tool()
@execute_on_main_thread
def get_current_file_path():
    return idc.get_input_file_path()

@mcp.tool()
@execute_on_main_thread
def list_files_with_relative_path(relative_path: str = ""):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is None or relative_path == "":
        return glob.glob(os.path.join(base_dir, "*"))
    else:
        return glob.glob(os.path.join(base_dir, relative_path, "*"))

@mcp.tool()
@execute_on_main_thread
def read_file(relative_path: str):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "r") as f:
        return f.read()

@mcp.tool()
@execute_on_main_thread
def write_file(relative_path: str, content: str):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "w") as f:
        f.write(content)

@mcp.tool()
@execute_on_main_thread
def read_binary(relative_path: str):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "rb") as f:
        return f.read()

@mcp.tool()
@execute_on_main_thread
def write_binary(relative_path: str , content: bytes):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "wb") as f:
        f.write(content)    

@mcp.tool()
@execute_on_main_thread
def eval_pythoni(script: str):
    return eval(script)

@mcp.tool()
@execute_on_main_thread
def get_instruction_length(address: int) -> int:
    """
    Retrieves the length (in bytes) of the instruction at the specified address.

    Args:
        address: The address of the instruction.

    Returns:
        The length (in bytes) of the instruction.  Returns 0 if the instruction cannot be decoded.
    """
    try:
        # Create an insn_t object to store instruction information.
        insn = ida_ua.insn_t()

        # Decode the instruction.
        length = ida_ua.decode_insn(insn, address)
        if length == 0:
            print(f"Failed to decode instruction at address {hex(address)}")
            return 0

        return length
    except Exception as e:
        print(f"Error getting instruction length: {str(e)}")
        return 0

@mcp.prompt()
def binary_analysis_strategy() -> str:
    """
    Guild for analyzing the binary
    """
    return (
        "IDA Pro MCP Server Tools and Best Practices:\n\n."
        "Tools: \n"
        "- get_bytes: Get bytes at specified address.\n"
        "- get_disasm: Get disassembly at specified address.\n"
        "- get_decompiled_func: Get decompiled pseudocode of function containing address.\n"
        "- get_function_name: Get function name at specified address.\n"
        "- get_segments: Get all segments information.\n"
        "- get_functions: Get all functions in the binary.\n"
        "- get_xrefs_to: Get all cross references to a specified address.\n"
        "- get_imports: Get all imports in the binary.\n"
        "- get_exports: Get all exports in the binary.\n"
        "- get_entry_point: Get the entry point of the binary.\n"
        "- make_function: Make a function at specified address.\n"
        "- undefine_function: Undefine a function at specified address.\n"
        "- get_dword_at: Get the dword at specified address.\n"
        "- get_word_at: Get the word at specified address.\n"
        "- get_byte_at: Get the byte at specified address.\n"
        "- get_qword_at: Get the qword at specified address.\n"
        "- get_float_at: Get the float at specified address.\n"
        "- get_double_at: Get the double at specified address.\n"
        "- get_string_at: Get the string at specified address.\n"
        "- get_strings: Get all strings in the binary.\n"
        "- get_current_file_path: Get the current path of the binary.\n"
        "- list_files_with_relative_path: List all files in the specified relative path in the current directory.\n"
        "- read_file: Read the content of a file.\n"
        "- write_file: Write content to a file.\n"
        "- read_binary: Read the content of a binary file.\n"
        "- write_binary: Write content to a binary file.\n"
        "- eval_python: Evaluate a Python script in IDA Pro.\n"
        "- get_instruction_length: Get the length of the instruction at the specified address.\n"
        "Best Practices: \n"
        "- Initial Analysis Phase\n"
        "   1. Examine the Entry Point\n"
        "       - Use the get_entry_point() tool to locate the program's entry point\n"
        "       - Analyze the code at the entry point to understand the program's startup flow\n"
        "       - Look for unusual instructions or jumps\n"
        "   2. Analyze Import Table\n"
        "       - Use the get_imports() tool to view all imported functions\n"
        "       - Look for suspicious API functions, such as:\n"
        "           - File operations: CreateFile, WriteFile\n"
        "           - Network communication: socket, connect, InternetOpen\n"
        "           - Process manipulation: CreateProcess, VirtualAlloc\n"
        "           - Registry operations: RegOpenKey, RegSetValue\n"
        "           - Cryptography related: CryptEncrypt, CryptDecrypt\n"
        "   3. Review Strings\n"
        "       - Use the get_strings() tool to obtain all strings\n"
        "       - Pay attention to IP addresses, URLs, domain names, file paths\n"
        "       - Look for encrypted or obfuscated string patterns\n"
        "       - Analyze command line parameters and error messages\n"
        "   4. In-Depth Analysis Phase\n"
        "       - Track Key API Calls\n"
        "           - Use get_xrefs_to() to find cross-references to suspicious imported functions\n"
        "           - Use get_decompiled_func() to analyze functions that call these APIs\n"
        "           - Analyze how parameters and return values are handled\n"
        "   5. Identify Main Functional Blocks\n"
        "       - Use get_functions() to get a list of all functions\n"
        "       - Sort functions by size and complexity\n"
        "       - Decompile and analyze large, complex functions\n"
        "       - Look for suspicious function names or unnamed functions\n"
        "   6. Analyze Control Flow\n"
        "       - Observe conditional branches and loop structures\n"
        "       - Analyze function call graphs and execution paths\n"
        "       - Look for anti-debugging and anti-VM detection techniques\n"
        "       - Pay attention to unusual jumps and callback mechanisms\n"
        "   7. Identifying Malicious Behaviors\n"
        "       - Identify Common Malicious Functionality\n"
        "           - Persistence mechanisms: Registry modifications, startup items, service creation\n"
        "           - Data theft: File searching, keylogging, screen capturing\n"
        "           - Communication features: C&C communication, data exfiltration channels\n"
        "           - Evasion techniques: Obfuscation, packing, anti-analysis checks\n"
        "           - Destructive behaviors: File encryption, system damage\n"
        "   8. Analyze Algorithms and Encryption Routines\n"
        "       - Identify encryption and decryption functions\n"
        "       - Look for hardcoded keys and cryptographic constants\n"
        "       - Analyze how data is processed in memory\n"
        "   9. Analyze Network Communication\n"
        "       - Identify network communication functions\n"
        "       - Look for IP addresses, URLs, and domain names\n"
        "       - Analyze how data is sent and received over the network\n"
        "   10. Dump payloads if there is any decryption or encoding\n"
        "       - Generate decryption or decodeing script in python and ida python\n"
        "       - Use eval_python to execute the script in IDA Pro\n"
        "       - Dump payloads in the current directory\n"
        "   11. Document Analysis Results\n"
        "       - Add comments to key functions\n"
        "       - Rename functions to reflect their actual functionality\n"
        "       - Create a logical structure diagram of the code\n"
        "   12. Using Advanced Techniques\n"
        "       - Use IDA Pro's advanced features like IDAPython scripting, IDA SDK, and IDA API\n"
        "       - Implement custom analysis scripts to automate repetitive tasks\n"
        "       - Explore IDA Pro's plugin ecosystem for additional analysis capabilities\n"
    )


def create_starlette_app(mcp_server: Server, *, debug: bool = False) -> Starlette:
    """Create a Starlette application that can serve the provided mcp server with SSE."""

    middleware = [
        Middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_methods=["*"],
            allow_headers=["*"],
        )
    ]
    return Starlette(
        debug=debug,
        middleware=middleware,
        routes=[
            Mount("/", app=mcp.sse_app()),
        ],
    )


class ModelContextProtocolPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_FIX | ida_idaapi.PLUGIN_HIDE
    comment = "IDA Model Context Protocol Server"
    help = "Provides REST API and SSE for IDA Pro analysis"
    wanted_name = "IDA MCP Server"
    wanted_hotkey = ""

    def init(self):
        try:
            print("Initializing IDA Model Context Protocol Server...")
            # app = create_starlette_app(mcp, debug=True)

            def run_server():
                try:
                    # 設置將異常轉換為 JSON 響應
                    mcp.run(transport="sse")
                    # uvicorn.run(app, host="localhost", port=3000, log_level="debug")
                except Exception as e:
                    print(f"Server error: {str(e)}")

            server_thread = threading.Thread(target=run_server)
            server_thread.daemon = True
            server_thread.start()
            print("Server started successfully!")
            return ida_idaapi.PLUGIN_KEEP
        except Exception as e:
            print(f"Failed to start server: {str(e)}")
            return ida_idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass

    def term(self):
        print("Terminating IDA Model Context Protocol Server...")


def PLUGIN_ENTRY():
    return ModelContextProtocolPlugin()


if __name__ == "__main__":
    PLUGIN_ENTRY()


repository: AzzOnFire/emuit
file: emuit.py
description: Easy-to-use IDA plugin for code emulation
url: https://github.com/AzzOnFire/emuit/blob/af36f7b8c6c67cfa165a379f440ffe9005f0e013/emuit.py
comment: Easy-to-use code emulator
created_at: 2022-12-28 20:21:55+00:00
pushed_at: 2024-05-06 20:35:02+00:00
content:
import sys

from emuit import EmuItIda

import idaapi
import ida_kernwin


ida_version = tuple(map(int, idaapi.get_kernel_version().split(".")))
assert (ida_version > (7, 4)), "ERROR: EmuIt requires IDA 7.4+"
assert (sys.version_info >= (3, 5)), "ERROR: EmuIt requires Python 3.6"


PLUGIN_NAME = 'EmuIt'
VERSION = '0.6.1'
PLUGIN_HOTKEY = 'Shift+C'

ACTION_RUN = 'EmuIt:run'
ACTION_RESET = 'EmuIt:reset'
ACTION_TOGGLE_RESET = 'EmuIt:toggle_reset'
ACTION_TOGGLE_SKIP_API_CALLS = 'EmuIt:toggle_skip_api_calls'
ACTION_TOGGLE_BEAUTIFY = 'EmuIt:toggle_beutify'


class EmuItPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = 'Easy-to-use code emulator'
    help = ''
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ''

    def banner(self):
        print('=' * 52, '\n', ' ' * 52)
        print(r'  ______     __    __     __  __     __     ______  ')
        print(r' /\  ___\   /\ "-./  \   /\ \/\ \   /\ \   /\__  _\ ')
        print(r' \ \  __\   \ \ \-./\ \  \ \ \_\ \  \ \ \  \/_/\ \/ ')
        print(r'  \ \_____\  \ \_\ \ \_\  \ \_____\  \ \_\    \ \_\ ')
        print(r'   \/_____/   \/_/  \/_/   \/_____/   \/_/     \/_/ ')
        print(' ' * 52, '\n', '=' * 52)

    def init(self):
        try:
            self.emu = EmuItIda()
        except Exception as e:
            print('EmuIt: an error occurred during initialization:', str(e))
            print('EmuIt: try to rebase program and reopen IDB')
            return idaapi.PLUGIN_SKIP

        self.reset_every_run = True
        self.beautify = False

        action_run = idaapi.action_desc_t(
            ACTION_RUN,
            'EmuIt run',
            action_handler(self.action_run_handler),
            PLUGIN_HOTKEY,
            'Run EmuIt on selection'
        )

        action_reset = idaapi.action_desc_t(
            ACTION_RESET,
            'Reset state',
            action_handler(self.action_reset_handler),
            None,
            'Reset EmuIt state now'
        )

        action_toggle_reset = idaapi.action_desc_t(
            ACTION_TOGGLE_RESET,
            'Reset on every run',
            action_handler(self.action_toggle_reset_handler),
            None,
            'Reset EmuIt state on every run'
        )

        action_toggle_skip_api_calls = idaapi.action_desc_t(
            ACTION_TOGGLE_SKIP_API_CALLS,
            'Skip API calls',
            action_handler(self.action_toggle_skip_api_calls_handler),
            None,
            'Try to skip API calls, you may be lucky!'
        )

        action_toggle_beutify = idaapi.action_desc_t(
            ACTION_TOGGLE_BEAUTIFY,
            'Filter and beautify results',
            action_handler(self.action_toggle_beautify_handler),
            None,
            'Beatify output using metrics and ASCII/Unicode encoding'
        )

        idaapi.register_action(action_run)
        idaapi.register_action(action_reset)

        idaapi.register_action(action_toggle_reset)
        idaapi.update_action_checkable(ACTION_TOGGLE_RESET, True)
        idaapi.update_action_checked(ACTION_TOGGLE_RESET, self.reset_every_run)

        idaapi.register_action(action_toggle_skip_api_calls)
        idaapi.update_action_checkable(ACTION_TOGGLE_SKIP_API_CALLS, True)
        idaapi.update_action_checked(ACTION_TOGGLE_SKIP_API_CALLS,
                                     self.emu.skip_api_calls)

        idaapi.register_action(action_toggle_beutify)
        idaapi.update_action_checkable(ACTION_TOGGLE_BEAUTIFY, True)
        idaapi.update_action_checked(ACTION_TOGGLE_BEAUTIFY, self.beautify)

        self.hooks = EmuItUIHooks()
        self.hooks.hook()

        self.banner()
        print(f'EmuIt {VERSION} started.')
        print(f'EmuIt run shortcut key is {PLUGIN_HOTKEY}.')

        return idaapi.PLUGIN_KEEP

    def action_run_handler(self):
        selection, start_ea, end_ea = ida_kernwin.read_range_selection(None)

        if not selection:
            ea = ida_kernwin.get_screen_ea()
            start_ea = idaapi.get_item_head(ea)
            end_ea = idaapi.get_item_end(ea)

        if self.reset_every_run:
            self.emu.reset()

        print(f'EmuIt: running {start_ea:08X} - {end_ea:08X}')
        results = self.emu.run(start_ea, end_ea)
        if self.beautify:
            results = results.pretty()

        for offset, data in results.items():
            print(hex(offset), data)
        print('EmuIt: finish')

    def action_reset_handler(self):
        self.emu.reset()

    def action_toggle_reset_handler(self):
        self.reset_every_run = bool(not self.reset_every_run)

    def action_toggle_skip_api_calls_handler(self):
        self.emu.skip_api_calls = bool(not self.emu.skip_api_calls)

    def action_toggle_beautify_handler(self):
        self.beautify = bool(not self.beautify)

    def term(self):
        if hasattr(self, 'hooks'):
            self.hooks.unhook()
        print(f'EmuIt {VERSION} terminated.')

    def run(self, arg):
        print('call EmuIt from context menu (right click) in disassembly')


class EmuItUIHooks(idaapi.UI_Hooks):
    def populating_widget_popup(self, widget, popup):
        pass

    def finish_populating_widget_popup(self, widget, popup):
        if (idaapi.get_widget_type(widget) not in {idaapi.BWN_DISASM, idaapi.BWN_PSEUDOCODE}):
            return 0

        tree = PLUGIN_NAME + " settings"
        attach = idaapi.attach_action_to_popup

        attach(widget, popup, ACTION_RUN, tree, idaapi.SETMENU_APP)
        attach(widget, popup, ACTION_RESET, f'{tree}/')
        attach(widget, popup, ACTION_TOGGLE_RESET, f'{tree}/')
        attach(widget, popup, ACTION_TOGGLE_SKIP_API_CALLS, f'{tree}/')
        attach(widget, popup, ACTION_TOGGLE_BEAUTIFY, f'{tree}/')

        return 0


def action_handler(callback):
    class Handler(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

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

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


plugin = EmuItPlugin()


def PLUGIN_ENTRY():
    global plugin
    return plugin


repository: zmrbak/IDAPython
file: 配套代码/ida10/python/examples/core/install_user_defined_prefix.py
description: 在线课程《IDAPython脚本化软件逆向分析入门》课程源码
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2022-05-02 02:13:34+00:00
pushed_at: 2022-05-14 10:41:43+00:00
url:
https://github.com/zmrbak/IDAPython/blob/478278339ba0c150b3ace657296233d8a89b26ea/%E9%85%8D%E5%A5%97%E4%BB%A3%E7%A0%81/ida10/python/examples/core/install_user_defined_prefix.py

content:
"""
summary: inserting information into disassembly prefixes

description:
  By default, disassembly line prefixes contain segment + address
  information (e.g., '.text:08047718'), but it is possible to
  "inject" other bits of information in there, thanks to the
  `ida_lines.user_defined_prefix_t` helper type.
"""

from __future__ import print_function

import ida_lines
import ida_idaapi

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

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


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

    def __init__(self):
        self.prefix = None

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

    def run(self, arg):
        pass

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


def PLUGIN_ENTRY():
    return prefix_plugin_t()


repository: david-lazar/IDAPatternSearch
file: IDAPatternSearch_utils/IDAPatternSearch.py
wanted_name: IDA Pattern Search
created_at: 2021-08-25 11:26:45+00:00
pushed_at: 2021-09-14 20:10:46+00:00
description:
IDAPatternSearch adds a capability of finding functions according to bit-patterns into the well-known IDA Pro disassembler based on Ghidra’s function patterns format.

url:
https://github.com/david-lazar/IDAPatternSearch/blob/8cbef7e07da74b01c58e505858a866eb1b59f7ce/IDAPatternSearch_utils/IDAPatternSearch.py

content:
import ida_funcs
import ida_bytes
import idc
import ida_search
import idaapi
import ida_xref
import ida_segment
import ida_offset
import ida_kernwin
from idaapi import BADADDR
from ida_search import SEARCH_DOWN, SEARCH_UP
import ida_auto

import IDAPatternSearch_utils.ida_common as ida_common

import sys
import os
import pathlib
import struct
import ast

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

PATH_TO_FUNCTION_PATTERNS = str(pathlib.Path(
    __file__).parent.absolute()) + "\\function_patterns\\"

# Each dictionary entry should include: (pattern_file_name, use_pre_patterns, use_post_patterns, arch_name, t_register_value=None)
_SEARCH_PARAMETERS = {
    'THUMB_LE_LOOSENED': ("THUMB_LE_LOOSENED.xml", False, True, "THUMB", 1),
    'THUMB_LE': ("THUMB_LE.xml", False, True, "THUMB", 1),
    'THUMB_BE_LOOSENED': ("THUMB_BE_LOOSENED.xml", False, True, "THUMB", 1),
    'THUMB_BE': ("THUMB_BE.xml", False, True, "THUMB", 1),
    'ARM_32_LE': ("ARM_32_LE.xml", False, True, "ARM_32", 0),
    'ARM_32_BE': ("ARM_32_BE.xml", False, True, "ARM_32", 0),
    'AARCH_64_LE': ("AARCH_64_LE.xml", False, True, "AARCH64"),
    'PPC_BE': ("PPC_BE.xml", False, True, "PowerPC"),
    'V850': ("V850.xml", False, True, "V850")
}


# Alt + G
def _change_t_register(ea, val):
    if val == -1:
        return

    idaapi.split_sreg_range(ea, idaapi.str2reg("T"), val, idaapi.SR_user)


def _parse_ghidra_pattern_item(pattern_item, element_length):
    '''
    This function parses a given Ghidra pattern item (pattern_item) which is one item from a Ghidra pattern.
    The function also uses (element_length) to determine if it is: 
        * Hex item (starting with 0x which was already omitted before calling this function)
          which in this case (element_length) == 4
        * Bitfield item
          which in this case (element_length) == 1
    The '.' characther in both item types represents a wildcard Bit/Byte depends on the item type (Hex/Bitfield).
    This function returns two values: (image, mask) 
        * image - represents the item image to search by the pattern.
        * mask -  represents the item mask that can be used to mask out the matched bytes and check against the image
    '''
    cur_image = 0
    cur_mask = 0
    for element in pattern_item:
        if element == '.':
            # Wildcard element
            # mask should be zeros (in bits), image should be zero (lets say 0)
            cur_mask = cur_mask << element_length
            cur_image = cur_image << element_length
        else:
            # mask should be 1's (in bits), image is same as half_byte
            cur_mask = cur_mask << element_length
            if element_length == 1:
                # bit element
                cur_mask += 1
            else:
                # half byte element
                cur_mask += 0xf

            cur_image = cur_image << element_length
            if element_length == 1:
                # bit element
                cur_image += int(element, 2)
            else:
                # half byte element
                cur_image += int(element, 16)
    return cur_image, cur_mask


def _convert_ghidra_pattern_to_image_and_mask(ghidra_pattern):
    '''
    This function parses a given Ghidra pattern (ghidra_pattern) after extracted from the XML file already.
    Every item is parsed and at the end, all the items are joined into 2 returned byte strings: image and mask.
    The function returns a dictionary with the keys:
        * image - represents the image to search by the pattern.
        * mask - represents the mask that can be used to mask out the matched bytes and check against the image
    The values in the dictionary are byte-strings.
    '''
    ghidra_pattern = ghidra_pattern.split()
    image = [b'']*len(ghidra_pattern)
    mask = [b'']*len(ghidra_pattern)
    for i in range(len(ghidra_pattern)):
        pattern_item = ghidra_pattern[i]
        pattern_byte_len = 0  # Number of bytes presetend by pattern
        cur_image = 0  # Image of current pattern item
        cur_mask = 0  # Mask of current pattern item

        if '0x' in pattern_item[0:2]:
            # Hex parsing
            pattern_item = pattern_item[2:]  # Remove '0x' at start
            if len(pattern_item) == 2 or len(pattern_item) == 4 or len(pattern_item) == 8:
                cur_image, cur_mask = _parse_ghidra_pattern_item(
                    pattern_item, 4)

                # 1 or 2 or 4 byte format
                pattern_byte_len = len(pattern_item)//2
            else:
                print("[-] Unrecognized length while parsing: 0x" +
                      str(pattern_item))
                break
        else:
            # Bit parsing
            if len(pattern_item) == 8:
                cur_image, cur_mask = _parse_ghidra_pattern_item(
                    pattern_item, 1)
                # 1 byte format is the only case for bit parsing
                pattern_byte_len = 1
            else:
                print("[-] Unrecognized length while parsing: " +
                      str(pattern_item))
                break

        # Convert from int to bytes
        image[i] = cur_image.to_bytes(pattern_byte_len, byteorder='big')
        mask[i] = cur_mask.to_bytes(pattern_byte_len, byteorder='big')

    # Join image_list and mask_list
    return {"image": b''.join(image), "mask": b''.join(mask)}


def _parse_ghidra_xml_pattern_list(filepath):
    '''
    Parses an Ghidra pattern XML file resides in (filepath).
    This file contains pattern pairs where each pair consits of pre-patterns and post-patterns.
    Pre-patterns should occur before a function start definition (It is usually a function end in the pattern file).
    Post-patterns should occur in a place where a function start should be defined.

    Returns a dictionary contains two keys:
        * pre_patterns
        * post_patterns
    The values are lists where each item is a Ghidra pattern string.

    Note that only one node for each element type (except data) exists in most cases, 
    but the parsing process treats the elements as lists for a more general case.
    '''
    import xml.dom.minidom

    parsed_pre_patterns = []
    parsed_post_patterns = []

    doc = xml.dom.minidom.parse(filepath)
    patternpairs_list = doc.getElementsByTagName("patternpairs")
    for patternpairs_node in patternpairs_list:
        # Parse pre patterns
        prepatterns_list = patternpairs_node.getElementsByTagName(
            'prepatterns')
        for prepatterns_node in prepatterns_list:
            data_list = prepatterns_node.getElementsByTagName('data')
            for data_node in data_list:
                data_value = data_node.firstChild.nodeValue
                parsed_pre_patterns.append(data_value)

        # Parse post patterns
        postpatterns_list = patternpairs_node.getElementsByTagName(
            'postpatterns')
        for postpatterns_node in postpatterns_list:
            data_list = postpatterns_node.getElementsByTagName('data')
            for data_node in data_list:
                data_value = data_node.firstChild.nodeValue
                parsed_post_patterns.append(data_value)
    return {"pre_patterns": parsed_pre_patterns, "post_patterns": parsed_post_patterns}


def _relax_ghidra_mask(mask):
    '''
    In order to use Ghidra patterns in IDA, we need to make the patterns more premissive
    because IDA only allows byte wildcards, while Ghidra patterns involving bits.

    This function returns the premissive mask given a strict (mask) byte string.

    IDA bin_search mask should be composed of \x00 or \x01 bytes, 
    where \x01 means to perform the comparison and \x00 means not tp perform.
    If the strict mask contains zero bit in one of the bytes, this byte should be \x00 in the relaxed mask.
    Else, the byte contains only one bits, and therefore this byte should be \x01 in the relaxed mask.
    '''
    mask_list = list(mask)
    relaxed_mask = [b'\x00' if b != 0xff else b'\x01' for b in mask_list]
    return b''.join(relaxed_mask)


def parse_and_search(pattern_file_name, use_pre_patterns, use_post_patterns, arch_name, t_register_value=None, **kwargs):
    '''
    Find function prologues based on patterns for given architecture and pattern file.
    '''
    pre_patterns = [None]
    post_patterns = [None]

    pattern_list = _parse_ghidra_xml_pattern_list(
        os.path.join(PATH_TO_FUNCTION_PATTERNS, pattern_file_name))

    if use_post_patterns:
        print("[+] Searching using post-patterns ({})".format(arch_name))
        post_patterns = pattern_list['post_patterns']
    if use_pre_patterns:
        print("[+] Searching using pre-patterns ({})".format(arch_name))
        pre_patterns = pattern_list['pre_patterns']

    _prologue_pattern_search(
        pre_patterns, post_patterns, t_register_value, arch_name, **kwargs)


def _prologue_pattern_search(pre_pattern_list, post_pattern_list, t_register_value, arch_name, **kwargs):
    """
    Find functions prologues in range(**kwargs) based on the given pattern list.

    Patterns are seperated into 2 different lists:  (pre_pattern_list,post_pattern_list). 
    * Pre_pattern is before function definition (e.g. filler or other function end), 
    * Post_pattern is the actual function start.

    The function can use only one of the pattern lists. 
    In this case, the other list (the unused one) should be passed as [None] to the function.
    * As for now we will use only Post Patterns (but usage of pre-patterns is already implemented).

    For each pattern in the pattern list, finding bytes matching the pattern and tries to define a function in IDA.
    Uses the (t_register_value) to set the T register accordingly when tries to define a function. 
    """
    defined_counter = 0  # Number of defined functions
    # Go over all search results for a specific pattern
    for addr_range in ida_common.parse_address_ranges(**kwargs):
        start_ea = addr_range[0]
        end_ea = addr_range[1]
        for pre_pattern in pre_pattern_list:
            for post_pattern in post_pattern_list:
                cur_start_ea = start_ea  # restore cur_start_ea
                while(cur_start_ea < end_ea):
                    cursor, defined = _find_next_pattern_bytes_and_define_function(
                        cur_start_ea, pre_pattern, post_pattern, t_register_value, end_ea)
                    if cursor == BADADDR:
                        break
                    defined_counter += defined
                    cur_start_ea = cursor + 1      # Continue to search in next bytes

    print("[+] Total number of {} functions defined: {}".format(arch_name, defined_counter))


def _find_next_pattern_bytes_and_define_function(start_ea, pre_pattern, post_pattern, t_register_value, end_ea=BADADDR):
    """
    Find bytes starting in range (start_ea,end_ea) based on the given (post_pattern) or/and (pre_pattern).

    There are 3 different cases for the function operation:
    * When only post_pattern exists, the function will be defined at the matched bytes location.
    * When only pre_pattern exists, the function will be defined after the matched bytes location.
    * When both pre_pattern and post_pattern exists, the function will be defined after the pre-pattern matched bytes location.

    When there is a match, first sets the T register with the given value (t_register_value).
    Then, tries to define a function in IDA.
    Returns two values: 
       1. The matched address (in case not found: BADADDR).
       2. An integer that indicates if as a result a function was defined successfully (1 -> Success, 0 -> Failed to define).
    """
    defined = 0  # A function was defined eventually?

    pre_pattern_converted = {}  # in case both pre_pattern and post_pattern exists

    if pre_pattern is None:
        pattern = post_pattern
    elif post_pattern is None:
        pattern = pre_pattern
    else:  # Both pre_pattern and post_pattern
        pattern = pre_pattern + " " + post_pattern
        pre_pattern_converted = _convert_ghidra_pattern_to_image_and_mask(
            pre_pattern)

    pattern_converted = _convert_ghidra_pattern_to_image_and_mask(pattern)

    # Assign pre_pattern_converted in case only pre_pattern exists
    if post_pattern is None:
        pre_pattern_converted = pattern_converted

    pattern_image = pattern_converted['image']
    pattern_mask = pattern_converted['mask']
    # mask from Ghidra should be relaxed to suit IDA search
    relaxed_mask = _relax_ghidra_mask(pattern_mask)
    cursor = ida_bytes.bin_search(start_ea, end_ea, pattern_image,
                                  relaxed_mask, idaapi.BIN_SEARCH_FORWARD, idaapi.BIN_SEARCH_NOCASE)
    if cursor == BADADDR:
        return BADADDR, defined

    # Check if strict mask holds (as defined in Ghidra and includes bits)
    matched_bytes = ida_bytes.get_bytes(cursor, len(pattern_mask))
    masked_bytes = b''.join([int.to_bytes(a & b, 1, 'big')
                            for a, b in zip(matched_bytes, pattern_mask)])
    if masked_bytes != pattern_image:  # because wildcard bit are 0's in the image, we don't need to compare here to pattern_image&pattern_mask
        # strict mask doesn't hold. Continue to next match.
        return cursor, defined

    # Increment cursor if pre_pattern exists
    if pre_pattern is not None:
        cursor = cursor + len(pre_pattern_converted['image'])

    # Only operate on unknown bytes or code bytes that are not part of functions (code bytes marked in red on IDA)
    if ida_bytes.is_unknown(ida_bytes.get_flags(cursor)) or (ida_bytes.is_code(ida_bytes.get_flags(cursor)) and ida_funcs.get_func_num(cursor) == -1):
        # Pattern found.
        # Switch to thumb, try to define a function
        # If failed -> Discard changes
        print("[+] Possible function prologue at {}  -> {}".format(hex(cursor),
              hex(ida_bytes.get_dword(cursor))))

        # Save T register old value (in case we will need to restor it)
        if t_register_value != None:
            t_register_old_value = idc.get_sreg(cursor, "T")

            # Change T register
            _change_t_register(cursor, t_register_value)

        # Try to define a function
        if not ida_funcs.add_func(cursor):
            print("[-]        Could not define a function at {}".format(hex(cursor)))

            # If the definition of the func failed and you have remained with code,
            # undo what you did by "U"'ing it
            if ida_bytes.is_code(ida_bytes.get_flags(cursor)):
                # delete all the chunk you've just tried to define
                idc.del_items(cursor, idc.DELIT_EXPAND, 1)

            if t_register_value != None:
                # Restore T register
                _change_t_register(cursor, t_register_old_value)
        else:
            defined = 1
    return cursor, defined


def explore_using_patterns(**kwargs):
    """
    Identifies the current architecture and explore for functions using predefined patterns.
    The predefined patterns interpreted from Ghidra.

    Supports: 
    - THUMB, ARM 32 bit
    - AARCH64 (ARM 64 bit) (only LE is supported).
    - PowerPC (only BE is supported).
    - V850 (V850E1)
    """

    # Make sure function_patterns directory exists
    if not os.path.exists(PATH_TO_FUNCTION_PATTERNS):
        print("[-] function_patterns directory does not exist")
        print("    Please make sure to include the function_patterns directory in path: " +
              str(PATH_TO_FUNCTION_PATTERNS))
        return

    info = idaapi.get_inf_structure()

    # Handle ARM/Thumb
    if info.procname == "ARM" or info.procname == "ARMB":
        # 64 bit architectures
        # Turns out that AARCH64 bit will return true in is_32bit(). So we need to first check for is_64bit().
        if info.is_64bit():
            parse_and_search(*(_SEARCH_PARAMETERS["AARCH_64_LE"]), **kwargs)

        # 32 bit architectures
        else:
            if info.is_be():
                # Big endian
                parse_and_search(*(_SEARCH_PARAMETERS["ARM_32_BE"]), **kwargs)
                parse_and_search(*(_SEARCH_PARAMETERS["THUMB_BE"]), **kwargs)
                parse_and_search(
                    *(_SEARCH_PARAMETERS["THUMB_BE_LOOSENED"]), **kwargs)
            else:
                # Little endian
                parse_and_search(*(_SEARCH_PARAMETERS["ARM_32_LE"]), **kwargs)
                parse_and_search(*(_SEARCH_PARAMETERS["THUMB_LE"]), **kwargs)
                parse_and_search(
                    *(_SEARCH_PARAMETERS["THUMB_LE_LOOSENED"]), **kwargs)

            # _undefine_filler_functions() # Check for filler functions in ARM or THUMB, in case pre-patterns are used.

    # Handle PowerPC
    elif info.procname == "PPC":
        parse_and_search(*(_SEARCH_PARAMETERS["PPC_BE"]), **kwargs)

    # Handle V850
    elif "V850" in info.procname:
        parse_and_search(*(_SEARCH_PARAMETERS["V850"]), **kwargs)

    else:
        print("[-] Unsupported architecture: " + str(info.procname))

    print("[+] Done exploration")


def _undefine_filler_functions():
    '''
    Searching for all functions that were wrongly defined due to filler pattern ("00 00") and contains all zero-bytes.
    These functions will be undefined.
    Currently doesn't give much added value because we don't use the filler pattern "00 00"
    '''
    print("[+] Searching for filler functions")  # Currently filler functions is all zeros functions.

    # Go over all functions
    for func_addr in ida_common.get_function_iterator():
        # Get function bytes and check for all zeros.
        cur_func = ida_funcs.get_func(func_addr)
        func_bytes = ida_bytes.get_bytes(
            cur_func.start_ea, cur_func.end_ea-cur_func.start_ea)
        is_filler_func = True
        for b in func_bytes:
            if b != 0:
                is_filler_func = False
                break

        # If function bytes are only zeros -> undefine_function
        if is_filler_func:
            _undefine_function(cur_func.start_ea, "Zero filler")


def _undefine_function(start_ea, reason):
    '''
    Undefine the function in address (address_ea).
    Also prints the (reason) for the deletion.
    '''
    # Delete all the function code chunk you've defined
    if (idc.del_items(start_ea, idc.DELIT_EXPAND, 1)):
        print("[+] Deleted function at {}. Reason: {}".format(hex(start_ea), reason))


#--------------------------------------------------------------------------
# IDA Plugin Stuff
#--------------------------------------------------------------------------

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

class IDAPatternSearchPlugin(idaapi.plugin_t):

    flags = 0
    comment = "Identifies the current architecture and look for functions prologues in undefined code using Ghidra patterns"
    help = ""
    wanted_name = "IDA Pattern Search"
    wanted_hotkey = ""

    def init(self):
        """
        What happens when the plugin is loaded.
        """

        if sys.version_info[0] <= 2:
            print('[IDA Pattern Search] IDA Pattern Search supports Python 3+, you are using older version.')
            return idaapi.PLUGIN_SKIP

        print("""[IDA Pattern Search] Plugin loaded, use Edit -> Plugins -> IDA Pattern Search to search for functions """)

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        Run plugin.
        """

        yn_value = ida_kernwin.ask_yn(1,"Search for function prologues in all possible undefined code?")
        if yn_value == 1:
            explore_using_patterns()
        
        elif yn_value == 0:
            answer = ida_kernwin.ask_str("{\"address_ranges\":(0x0, 0xFFFFFFFF),\"segments\":[\".text\"]}",\
            0,\
            """Please write target addresses.Target addresses can be specified in two non-exclusive ways using a python dictionary:
            1) Address range(s) - either a tuple or a list of tuples specifying 
            a start address and an exclusive end address, passed in the address_ranges argument. 
            Passing an empty list would result in including NO addresses.
            EXAMPLE: {\"address_ranges\":[(0, 0x1000), (0xFFFF0000, 0xFFFFFFFF)]}
            2) Segment name(s) - either a segment name or a list of segment names, 
            passed in the segments argument. Passing an empty list would include all arguments in the IDB.
            EXAMPLE: {\"segments\":[\".text\", \".bss\"]}
            """)
            if answer is not None:
                explore_using_patterns(**(ast.literal_eval(answer)))

    def term(self):
        pass

comment:
Identifies the current architecture and look for functions prologues in undefined code using Ghidra patterns


repository: L4ys/LazyIDA
file: LazyIDA.py
description: Make your IDA Lazy!
url: https://github.com/L4ys/LazyIDA/blob/426b3aab008c14870f7ff6a142b0cd232f9a669a/LazyIDA.py
wanted_name: LazyIDA
comment: LazyIDA
created_at: 2016-06-12 14:47:06+00:00
pushed_at: 2025-04-03 14:25:51+00:00
content:
from __future__ import division
from __future__ import print_function
from struct import unpack
import idaapi
import idautils
import idc

from PyQt5.Qt import QApplication

ACTION_CONVERT = ["lazyida:convert%d" % i for i in range(10)]
ACTION_SCANVUL = "lazyida:scanvul"
ACTION_COPYEA = "lazyida:copyea"
ACTION_COPYFO = "lazyida:copyfo"
ACTION_GOTOCLIPEA = "lazyida:gotoclipea"
ACTION_GOTOCLIPFO = "lazyida:gotoclipfo"
ACTION_XORDATA = "lazyida:xordata"
ACTION_FILLNOP = "lazyida:fillnop"

ACTION_HX_REMOVERETTYPE = "lazyida:hx_removerettype"
ACTION_HX_COPYEA = "lazyida:hx_copyea"
ACTION_HX_COPYFO = "lazyida:hx_copyfo"
ACTION_HX_COPYNAME = "lazyida:hx_copyname"
ACTION_HX_GOTOCLIPEA = "lazyida:hx_gotoclipea"
ACTION_HX_GOTOCLIPFO = "lazyida:hx_gotoclipfo"

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

ARCH = 0
BITS = 0

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

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

def parse_location(loc, is_fo=False):
    is_named = False
    ascii_text = ""
    try:
        loc = int(loc, 16)
        if is_fo:
            loc = idaapi.get_fileregion_ea(loc)
    except ValueError:
        try:
            ascii_text = loc.encode(encoding="ascii",errors="replace").decode(encoding="ascii").strip()
            loc = idc.get_name_ea_simple(ascii_text)
            is_named = True
        except:
            return idaapi.BADADDR
    return loc, is_named, ascii_text

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 (EA) has been copied to clipboard" % ea)
        elif self.action == ACTION_COPYFO:
            ea = idc.get_screen_ea()
            if ea != idaapi.BADADDR:
                fo = idaapi.get_fileregion_offset(ea)
                if fo != idaapi.BADADDR:
                    copy_to_clip("0x%X" % fo)
                    print("Address 0x%X (FO) has been copied to clipboard" % fo)
        elif self.action == ACTION_GOTOCLIPEA:
            loc, is_named, name = parse_location(clip_text(), False)
            if loc != idaapi.BADADDR:
                if is_named:
                    print("Goto named location '%s' 0x%X" % (name, loc))
                else:
                    print("Goto location 0x%X (EA)" % loc)
                idc.jumpto(loc)
        elif self.action == ACTION_GOTOCLIPFO:
            loc, is_named, name = parse_location(clip_text(), True)
            if loc != idaapi.BADADDR:
                if is_named:
                    print("Goto named location '%s' 0x%X" % (name, loc))
                else:
                    print("Goto location 0x%X (FO)" % idaapi.get_fileregion_offset(loc))
                idc.jumpto(loc)
        return 1

    def update(self, ctx):
        if idaapi.IDA_SDK_VERSION >= 770:
            target_attr = "widget_type"
        else:
            target_attr = "form_type"

        if idaapi.IDA_SDK_VERSION >= 900:
            try:
                dump_type = idaapi.BWN_HEXVIEW
            except:
                dump_type = idaapi.BWN_DUMP
        else:
            dump_type = idaapi.BWN_DUMP

        if ctx.__getattribute__(target_attr) in (idaapi.BWN_DISASM, dump_type):
            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 (dump as)
            t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(), idaapi.get_current_viewer()
            if idaapi.read_selection(view, t0, t1):
                start, end = t0.place(view).toea(), t1.place(view).toea()
                size = end - start + 1
            elif idc.get_item_size(idc.get_screen_ea()) > 1:
                start = idc.get_screen_ea()
                size = idc.get_item_size(start)
                end = start + size
            else:
                return False

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

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

        return 1

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

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

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

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

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

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

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

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

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

class hexrays_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for hexrays actions
    """
    def __init__(self, action):
        idaapi.action_handler_t.__init__(self)
        self.action = action
        self.ret_type = {}

    def activate(self, ctx):
        if self.action == ACTION_HX_REMOVERETTYPE:
            vdui = idaapi.get_widget_vdui(ctx.widget)
            self.remove_rettype(vdui)
            vdui.refresh_ctext()
        elif self.action == ACTION_HX_COPYEA:
            ea = idaapi.get_screen_ea()
            if ea != idaapi.BADADDR:
                copy_to_clip("0x%X" % ea)
                print("Address 0x%X (EA) has been copied to clipboard" % ea)
        elif self.action == ACTION_HX_COPYFO:
            ea = idaapi.get_screen_ea()
            if ea != idaapi.BADADDR:
                fo = idaapi.get_fileregion_offset(ea)
                if fo != idaapi.BADADDR:
                    copy_to_clip("0x%X" % fo)
                    print("Address 0x%X (FO) has been copied to clipboard" % fo)
        elif self.action == ACTION_HX_COPYNAME:
            highlight = idaapi.get_highlight(idaapi.get_current_viewer())
            name = highlight[0] if highlight else None
            if name:
                copy_to_clip(name)
                print("'%s' has been copied to clipboard" % name)
        elif self.action == ACTION_HX_GOTOCLIPEA:
            loc, is_named, name = parse_location(clip_text(), False)
            if loc != idaapi.BADADDR:
                if is_named:
                    print("Goto named location '%s' 0x%X" % (name, loc))
                else:
                    print("Goto location 0x%X (EA)" % loc)
                idc.jumpto(loc)
        elif self.action == ACTION_HX_GOTOCLIPFO:
            loc, is_named, name = parse_location(clip_text(), True)
            if loc != idaapi.BADADDR:
                if is_named:
                    print("Goto named location '%s' 0x%X" % (name, loc))
                else:
                    print("Goto location 0x%X (FO)" % idaapi.get_fileregion_offset(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 idaapi.IDA_SDK_VERSION >= 900:
            try:
                dump_type = idaapi.BWN_HEXVIEW
            except:
                dump_type = idaapi.BWN_DUMP
        else:
            dump_type = idaapi.BWN_DUMP

        if form_type == idaapi.BWN_DISASM or form_type == dump_type:
            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, "Dump/")

        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()

        if idaapi.IDA_SDK_VERSION >= 900:
            if idaapi.inf_is_64bit():
                BITS = 64
            elif idaapi.inf_is_32bit_exactly():
                BITS = 32
            elif idaapi.inf_is_16bit():
                BITS = 16
            else:
                raise ValueError
        else:
            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.5) plugin has been loaded.")

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

        # Register hotkey actions
        hotkey_actions = (
            idaapi.action_desc_t(ACTION_COPYEA, "Copy EA", hotkey_action_handler_t(ACTION_COPYEA), "w", "Copy current EA", 0),
            idaapi.action_desc_t(ACTION_COPYFO, "Copy FO", hotkey_action_handler_t(ACTION_COPYFO), "Shift-W", "Copy current FO", 0),
            idaapi.action_desc_t(ACTION_GOTOCLIPEA, "Goto clipboard EA", hotkey_action_handler_t(ACTION_GOTOCLIPEA), "Shift-G"),
            idaapi.action_desc_t(ACTION_GOTOCLIPFO, "Goto clipboard FO", hotkey_action_handler_t(ACTION_GOTOCLIPFO), "Ctrl-Shift-G"),
        )
        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.5"
            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", "Copy current EA", 0),
                idaapi.action_desc_t(ACTION_HX_COPYFO, "Copy FO", hexrays_action_handler_t(ACTION_HX_COPYFO), "Shift-W", "Copy current FO", 0),
                idaapi.action_desc_t(ACTION_HX_GOTOCLIPEA, "Goto clipboard EA", hexrays_action_handler_t(ACTION_HX_GOTOCLIPEA), "Shift-G"),
                idaapi.action_desc_t(ACTION_HX_GOTOCLIPFO, "Goto clipboard FO", hexrays_action_handler_t(ACTION_HX_GOTOCLIPFO), "Ctrl-Shift-G"),
                idaapi.action_desc_t(ACTION_HX_COPYNAME, "Copy name", hexrays_action_handler_t(ACTION_HX_COPYNAME), "c"),
            )
            for action in hx_actions:
                idaapi.register_action(action)
                self.registered_hx_actions.append(action.name)

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

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

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

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

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

def PLUGIN_ENTRY():
    return LazyIDA_t()


repository: zmrbak/IDAPython
file: 配套代码/ida10/python/examples/hexrays/vds21.py
description: 在线课程《IDAPython脚本化软件逆向分析入门》课程源码
wanted_name: Hex-Rays custom prototype provider (IDAPython)
comment: Sample plugin21 for Hex-Rays decompiler
created_at: 2022-05-02 02:13:34+00:00
pushed_at: 2022-05-14 10:41:43+00:00
url:
https://github.com/zmrbak/IDAPython/blob/478278339ba0c150b3ace657296233d8a89b26ea/%E9%85%8D%E5%A5%97%E4%BB%A3%E7%A0%81/ida10/python/examples/hexrays/vds21.py

content:
"""
summary: dynamically provide a custom call type

description:
  This plugin can greatly improve decompilation of indirect calls:

      call    [eax+4]

  For them, the decompiler has to guess the prototype of the called function.
  This has to be done at a very early phase of decompilation because
  the function prototype influences the data flow analysis. On the other
  hand, we do not have global data flow analysis results yet because
  we haven't analyzed all calls in the function. It is a chicked-and-egg
  problem.

  The decompiler uses various techniques to guess the called function
  prototype. While it works very well, it may fail in some cases.

  To fix, the user can specify the call prototype manually, using
  "Edit, Operand types, Set operand type" at the call instruction.

  This plugin illustrates another approach to the problem:
  if you happen to be able to calculate the call prototypes dynamically,
  this is how to inform the decompiler about them.
"""

import ida_idaapi
import ida_nalt
import ida_kernwin
import ida_typeinf
import ida_hexrays

testing = False # only for testing purposes

class callinfo_provider_t(ida_hexrays.Hexrays_Hooks):

    # this callback will be called for all call instructions
    # our plugin may provide the function prototype or even a complete new callinfo
    # object. The callinfo object may be useful if the prototype is not enough
    # to express all details of the call.
    def build_callinfo(self, blk, type):
        # it is a good idea to skip direct calls.
        # note that some indirect calls may be resolved and become direct calls,
        # and will be filtered out here:
        ida_kernwin.msg("%x: got called for: %s\n" % (blk.tail.ea, blk.tail.dstr()))
        tail = blk.tail
        if tail.opcode == ida_hexrays.m_call:
            return

        # also, if the type was specified by the user, do not interfere
        call_ea = tail.ea
        tif = ida_typeinf.tinfo_t()
        if ida_nalt.get_op_tinfo(tif, call_ea, 0):
            return

        global testing
        if not testing:
            # ok, the decompiler really has to guess the type.
            # just for the sake of an example, return a predefined prototype.
            # in real life you will provide the prototype you discovered yourself,
            # using your magic of yours :)
            my_proto = "int f();"
            ida_kernwin.msg("%x: providing prototype %s\n" % (call_ea, my_proto))
            ida_typeinf.parse_decl(type, None, my_proto, 0)
        else:
            # as an alternative to filling 'type', you can
            # choose to return a mcallinfo_t instance.
            mi = ida_hexrays.mcallinfo_t()
            mi.cc = ida_typeinf.CM_CC_STDCALL | ida_typeinf.CM_N32_F48 # let's use stdcall to differentiate from the tinfo_t-filling version
            mi.return_type = ida_typeinf.tinfo_t(ida_typeinf.BT_INT)
            mi.return_argloc.set_reg1(0) # eax
            return mi

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays custom prototype provider (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin21 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.hooks = callinfo_provider_t()
            self.hooks.hook()
            ida_kernwin.warning(
                "Installed callinfo provider sample (vds21.py)\n" +\
                "Please note that it is just an example\n" +\
                "and will spoil your decompilations!")
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: AlexWMF/trash_jump
file: trash_jump.py
description: This plugin allows to jump at specific address inside a string with trash inside
url: https://github.com/AlexWMF/trash_jump/blob/bcafb96186e7e601b1da92e39016d1831e98211f/trash_jump.py
wanted_name: Trash Jump
created_at: 2015-12-28 14:01:14+00:00
pushed_at: 2017-08-25 10:25:20+00:00
content:
# -*- coding: utf-8 -*
__author__ = 'AlexWMF'


"""
    The MIT License (MIT)

    Copyright (c) 2015 AlexWMF
    
    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:
    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.
    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.
"""

import idaapi
import re
import traceback
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtCore import Qt


class UiCallable(object):
    def __init__(self, fn_or_fnlist):
        self.__fn = fn_or_fnlist

    def __call__(self):
        if isinstance(self.__fn, list):
            for item in self.__fn:
                item()
        else:
            self.__fn()
        return False


class TrashJumpDialog(QtWidgets.QDialog):
    def __init__(self, custom_base, parent=None, *args, **kwargs):
        super(TrashJumpDialog, self).__init__(parent, *args, **kwargs)
        self.cbb_addr = None
        self.le_custom_base = None
        self.lbl_diff = None
        self.custom_base = custom_base

        self.setUpUI()

    def setUpUI(self):
        self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint)
        self.setWindowTitle('Trash Jump')

        vlay = QtWidgets.QVBoxLayout(self)
        flay = QtWidgets.QFormLayout()

        lbl = QtWidgets.QLabel('TrashJump Address')
        flay.setWidget(0, QtWidgets.QFormLayout.LabelRole, lbl)

        self.cbb_addr = QtWidgets.QComboBox(self)
        self.cbb_addr.setEditable(True)
        flay.setWidget(0, QtWidgets.QFormLayout.FieldRole, self.cbb_addr)

        lbl = QtWidgets.QLabel('Custom base:', self)
        flay.setWidget(1, QtWidgets.QFormLayout.LabelRole, lbl)

        hlay = QtWidgets.QHBoxLayout()

        self.le_custom_base = QtWidgets.QLineEdit(self)
        self.le_custom_base.setMinimumSize(QtCore.QSize(161, 0))
        self.le_custom_base.setMaximumSize(QtCore.QSize(161, 16777215))

        self.le_custom_base.setText('0' if self.custom_base is None else '%x' % self.custom_base)
        self.le_custom_base.textChanged.connect(self.custom_base_changed)

        hlay.addWidget(self.le_custom_base)

        lbl = QtWidgets.QLabel('diff:', self)
        hlay.addWidget(lbl)

        self.lbl_diff = QtWidgets.QLabel(self)
        self.lbl_diff.setTextInteractionFlags(Qt.TextSelectableByMouse)
        hlay.addWidget(self.lbl_diff)

        hlay.setStretch(2, 1)
        flay.setLayout(1, QtWidgets.QFormLayout.FieldRole, hlay)
        vlay.addLayout(flay)

        bb = QtWidgets.QDialogButtonBox(self)
        bb.setOrientation(Qt.Horizontal)
        bb.setCenterButtons(True)
        bb.setStandardButtons(QtWidgets.QDialogButtonBox.Cancel|QtWidgets.QDialogButtonBox.Ok)
        bb.accepted.connect(self.accept)
        bb.rejected.connect(self.reject)

        vlay.addWidget(bb)

        self.resize(421, 102)
        self.custom_base_changed(self.le_custom_base.text())
        main_window = filter(lambda x: x.metaObject().className() == "IDAMainWindow", QtCore.QCoreApplication.instance().allWidgets())
        if not main_window:
            return
        main_window = main_window[0]
        parent_center = main_window.geometry().center()
        self.move(parent_center - self.rect().center())

    def accept(self):
        try:
            self.custom_base = int(self.le_custom_base.text(), 16)
            super(TrashJumpDialog, self).accept()
        except:
            pass

    def custom_base_changed(self, s):
        try:
            self.custom_base = int(self.le_custom_base.text(), 16)
            self.lbl_diff.setText('%x' % (idaapi.get_imagebase() - self.custom_base))
        except:
            self.lbl_diff.setText('<invalid>')


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

    help = ""
    wanted_name = "Trash Jump"
    wanted_hotkey = "G"
    _RE_ADDR = re.compile(r'\b(0x)?([0-9a-f]+)\b', re.I | re.U | re.M)

    NET_NODE = '$ trash_jump'

    def init(self):
        return idaapi.PLUGIN_OK

    @classmethod
    def parse(cls, s):
        rv = list()
        for _, addr in cls._RE_ADDR.findall(s.strip()):
            if not addr:
                continue
            try:
                if '0x' not in addr.lower():
                    addr = '0x' + addr
                rv.append(long(addr, 16))
            except:
                print 'err: %s' % traceback.format_exc()
        return rv

    def run(self, arg):
        cfn = UiCallable(self._run)
        idaapi.execute_ui_requests((cfn,))

    def get_config(self):
        n = idaapi.netnode()
        n.create(self.NET_NODE)

        return {
            'custom_base': n.altval(0) if n.altval(0) else idaapi.get_imagebase()
        }

    def save_config(self, cfg):
        n = idaapi.netnode()
        n.create(self.NET_NODE)

        n.altset(0, cfg['custom_base'] if cfg['custom_base'] else 0)

    def _run(self):
        try:
            askd = TrashJumpDialog(self.get_config()['custom_base'])
            if askd.exec_() != askd.Accepted:
                return

            s = askd.cbb_addr.currentText().encode('ascii').strip()
            diff = 0
            if askd.custom_base:
                diff = idaapi.get_imagebase() - askd.custom_base
                self.save_config({'custom_base': askd.custom_base})

            eas = self.parse(s)
            for ea in eas:
                ea += diff
                if idaapi.isEnabled(ea) and idaapi.jumpto(ea, idaapi.UIJMP_ACTIVATE | idaapi.UIJMP_IDAVIEW):
                    return

            ea = idaapi.get_name_ea(idaapi.BADADDR, s)
            if idaapi.BADADDR != ea:
                if idaapi.isEnabled(ea) and idaapi.jumpto(ea, idaapi.UIJMP_ACTIVATE | idaapi.UIJMP_IDAVIEW):
                    return

            # last try without custom base
            for ea in eas:
                if idaapi.isEnabled(ea) and idaapi.jumpto(ea, idaapi.UIJMP_ACTIVATE | idaapi.UIJMP_IDAVIEW):
                    return

            idaapi.msg('TrashJump: address not found. Parsed: %r\n' % [hex(ea) for ea in eas])
        except:
            idaapi.msg('TrashJump: address not found\nerror: %s\n' % traceback.format_exc())

    def term(self):
        pass


def PLUGIN_ENTRY():
    return trash_jump_plugin_t()

# if __name__ == '__main__':
#     import unittest
#     class SimpleTest(unittest.TestCase):
#         def test_1(self):
#             s = '50A66F10 ; Exported entry 2967. ?detac'
#             self.assertEqual(trash_jump_plugin_t.parse(s)[0], 0x50A66F10)
#     unittest.main()


repository: veritas501/iidada
file: iidada_plugin.py
description: Merge multi binaries into one IDA database.
url: https://github.com/veritas501/iidada/blob/da26a8375ddf092d04b77a8cb6a33b24afc2b5c5/iidada_plugin.py
comment: Merge multi binaries into one IDA database
created_at: 2022-01-11 12:40:35+00:00
pushed_at: 2022-01-18 12:11:45+00:00
content:
# -*- coding: utf8 -*-
# Author: veritas501

"""
IIDADA

Merge multi binaries into one IDA database
"""

import os

import ida_kernwin
import idaapi
from iidada.main import dump_all, load_all
from iidada.inf import get_filename
from iidada.util import DB_SUFFIX
from iidada.version import PROG_NAME, VERSION
import iidada.log

log = iidada.log.get_logger("plugin", iidada.log.INFO)


def iidada_produce():
    default_name = "{}{}".format(get_filename(), DB_SUFFIX)
    fname = ida_kernwin.ask_file(1, default_name, "Save database file as")
    if fname:
        if not fname.endswith(DB_SUFFIX):
            fname += DB_SUFFIX
        if dump_all(fname):
            log.info("Produce database success. {}".format(fname))
        else:
            log.error("Produce database failed.")


def iidada_load():
    default_name = "*{}".format(DB_SUFFIX)
    fname = ida_kernwin.ask_file(0, default_name, "Load database file")
    if fname:
        if load_all(fname):
            log.info("Merge database success. {}".format(fname))
        else:
            log.error("Merge database failed.")


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

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

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


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

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

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


class IIDADAPlugin(idaapi.plugin_t):
    flags = 0
    comment = "Merge multi binaries into one IDA database"
    help = ""
    wanted_name = PROG_NAME
    wanted_hotkey = ""
    produce_action_name = 'produce_iidada:action'
    load_action_name = 'load_iidada:action'
    menu_name = "IIDADA database file..."
    produce_tooltip = "Produce iidada database file."
    load_tooltip = "Load iidada signature file."
    menu_tab = 'File/'
    menu_context = []

    # noinspection PyPropertyAccess
    def init(self):
        produce_desc = idaapi.action_desc_t(
            self.produce_action_name,
            self.menu_name,
            ProduceIIDADAAction(),
            self.wanted_hotkey,
            self.produce_tooltip,
            199
        )

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

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

        idaapi.attach_action_to_menu(
            os.path.join(self.menu_tab, 'Produce file/'),
            self.produce_action_name,
            idaapi.SETMENU_APP
        )
        idaapi.attach_action_to_menu(
            os.path.join(self.menu_tab, 'Load file/'),
            self.load_action_name,
            idaapi.SETMENU_APP
        )

        if idaapi.init_hexrays_plugin():
            addon = idaapi.addon_info_t()
            addon.id = "com.veritas501.iidada"
            addon.name = PROG_NAME
            addon.producer = "veritas501"
            addon.url = "https://github.com/veritas501/iidada"
            addon.version = VERSION
            idaapi.register_addon(addon)

        return idaapi.PLUGIN_KEEP

    def term(self):
        idaapi.detach_action_from_menu(self.menu_tab, self.produce_action_name)
        idaapi.detach_action_from_menu(self.menu_tab, self.load_action_name)
        return None

    # noinspection PyUnusedLocal
    def run(self, arg):
        return None


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


repository: wabiverse/OpenTimelineIO
file: src/py-opentimelineio/opentimelineio/plugins/manifest.py
description: Open Source API and interchange format for editorial timeline information.
created_at: 2024-02-13 23:20:31+00:00
pushed_at: 2024-07-27 21:52:27+00:00
url:
https://github.com/wabiverse/OpenTimelineIO/blob/3e69330ffe37fd7222573d0d58c9ed4e115b4f6a/src/py-opentimelineio/opentimelineio/plugins/manifest.py

content:
# SPDX-License-Identifier: Apache-2.0
# Copyright Contributors to the OpenTimelineIO project

"""OTIO Python Plugin Manifest system: locates plugins to OTIO."""

from importlib import resources
import inspect
import logging
import os
from pathlib import Path

try:
    from importlib import metadata
except ImportError:
    # For python 3.7
    import importlib_metadata as metadata

from .. import (
    core,
    exceptions,
)


# for tracking what kinds of plugins the manifest system supports
OTIO_PLUGIN_TYPES = [
    'adapters',
    'media_linkers',
    'schemadefs',
    'hook_scripts',
    'hooks',
    'version_manifests',
]


def plugin_entry_points():
    """Returns the list of entry points for all available OpenTimelineIO
    plugins.
    """
    try:
        entry_points = metadata.entry_points(group='opentimelineio.plugins')
    except TypeError:
        # For python <= 3.9
        entry_points = metadata.entry_points().get(
            'opentimelineio.plugins', []
        )

    return entry_points


def manifest_from_file(filepath):
    """Read the .json file at filepath into a :py:class:`Manifest` object."""

    result = core.deserialize_json_from_file(filepath)
    absfilepath = os.path.abspath(filepath)
    result.source_files.append(absfilepath)
    result._update_plugin_source(absfilepath)
    return result


def manifest_from_string(input_string):
    """Deserialize the json string into a manifest object."""

    result = core.deserialize_json_from_string(input_string)

    # try and get the caller's name
    name = "unknown"
    stack = inspect.stack()
    if len(stack) > 1 and len(stack[1]) > 3:
        #                     filename     function name
        name = f"{stack[1][1]}:{stack[1][3]}"

    # set the value in the manifest
    src_string = f"call to manifest_from_string() in {name}"
    result.source_files.append(src_string)
    result._update_plugin_source(src_string)

    return result


@core.register_type
class Manifest(core.SerializableObject):
    """Defines an OTIO plugin Manifest.

    This is considered an internal OTIO implementation detail.

    A manifest tracks a collection of plugins and enables finding them by name
    or other features (in the case of adapters, what file suffixes they
    advertise support for).

    For more information, consult the documenation.
    """
    _serializable_label = "PluginManifest.1"

    def __init__(self):
        core.SerializableObject.__init__(self)
        self.adapters = []
        self.schemadefs = []
        self.media_linkers = []
        self.source_files = []

        # hook system stuff
        self.hooks = {}
        self.hook_scripts = []

        self.version_manifests = {}

    adapters = core.serializable_field(
        "adapters",
        type([]),
        "Adapters this manifest describes."
    )
    schemadefs = core.serializable_field(
        "schemadefs",
        type([]),
        "Schemadefs this manifest describes."
    )
    media_linkers = core.serializable_field(
        "media_linkers",
        type([]),
        "Media Linkers this manifest describes."
    )
    hooks = core.serializable_field(
        "hooks",
        type({}),
        "Hooks that hooks scripts can be attached to."
    )
    hook_scripts = core.serializable_field(
        "hook_scripts",
        type([]),
        "Scripts that can be attached to hooks."
    )
    version_manifests = core.serializable_field(
        "version_manifests",
        type({}),
        "Sets of versions to downgrade schemas to."
    )

    def extend(self, another_manifest):
        """
        Aggregate another manifest's plugins into this one.

        During startup, OTIO will deserialize the individual manifest JSON files
        and use this function to concatenate them together.
        """
        if not another_manifest:
            return

        self.adapters.extend(another_manifest.adapters)
        self.schemadefs.extend(another_manifest.schemadefs)
        self.media_linkers.extend(another_manifest.media_linkers)
        self.hook_scripts.extend(another_manifest.hook_scripts)

        for family, label_map in another_manifest.version_manifests.items():
            # because self.version_manifests is an AnyDictionary instead of a
            # vanilla python dictionary, it does not support the .set_default()
            # method.
            if family not in self.version_manifests:
                self.version_manifests[family] = {}
            self.version_manifests[family].update(label_map)

        for trigger_name, hooks in another_manifest.hooks.items():
            # because self.hooks is an AnyDictionary instead of a vanilla
            # python dictionary, it does not support the .set_default() method.
            if trigger_name not in self.hooks:
                self.hooks[trigger_name] = []
            self.hooks[trigger_name].extend(hooks)

        self.source_files.extend(another_manifest.source_files)

    def _update_plugin_source(self, path):
        """Set the source file path for the manifest."""

        for thing in (
            self.adapters
            + self.schemadefs
            + self.media_linkers
            + self.hook_scripts
        ):
            thing._json_path = path

    def from_filepath(self, suffix):
        """Return the adapter object associated with a given file suffix."""

        for adapter in self.adapters:
            if suffix.lower() in adapter.suffixes:
                return adapter
        raise exceptions.NoKnownAdapterForExtensionError(suffix)

    def adapter_module_from_suffix(self, suffix):
        """Return the adapter module associated with a given file suffix."""

        adp = self.from_filepath(suffix)
        return adp.module()

    # @TODO: (breaking change) this should search all plugins by default instead
    #        of just adapters
    def from_name(self, name, kind_list="adapters"):
        """Return the plugin object associated with a given plugin name."""

        for thing in getattr(self, kind_list):
            if name == thing.name:
                return thing

        raise exceptions.NotSupportedError(
            "Could not find plugin: '{}' in kind_list: '{}'."
            " options: {}".format(
                name,
                kind_list,
                getattr(self, kind_list)
            )
        )

    def adapter_module_from_name(self, name):
        """Return the adapter module associated with a given adapter name."""

        adp = self.from_name(name)
        return adp.module()

    def schemadef_module_from_name(self, name):
        """Return the schemadef module associated with a given schemadef name."""

        adp = self.from_name(name, kind_list="schemadefs")
        return adp.module()


_MANIFEST = None


def load_manifest():
    """Walk the plugin manifest discovery systems and accumulate manifests.

    The order of loading (and precedence) is:

       1. Manifests specified via the :term:`OTIO_PLUGIN_MANIFEST_PATH` variable
       2. Entrypoint based plugin manifests
       3. Builtin plugin manifest
       4. Contrib plugin manifest
    """

    result = Manifest()

    # Read plugin manifests defined on the $OTIO_PLUGIN_MANIFEST_PATH
    # environment variable.  This variable is an os.pathsep separated list of
    # file paths to manifest json files. Can be set to "" to indicate no
    # local custom manifest path.
    _local_manifest_path = os.environ.get("OTIO_PLUGIN_MANIFEST_PATH", None)
    if _local_manifest_path:
        for src_json_path in _local_manifest_path.split(os.pathsep):
            json_path = os.path.abspath(src_json_path)
            if (
                not os.path.exists(json_path)
                # the manifest has already been loaded
                or json_path in result.source_files
            ):
                # XXX: In case error reporting is requested
                # print(
                #     "Warning: OpenTimelineIO cannot access path '{}' from "
                #     "$OTIO_PLUGIN_MANIFEST_PATH".format(json_path)
                # )
                continue

            result.extend(manifest_from_file(json_path))

    if not os.environ.get("OTIO_DISABLE_ENTRYPOINTS_PLUGINS"):
        entry_points = plugin_entry_points()

        for plugin in entry_points:
            plugin_name = plugin.name
            try:
                plugin_entry_point = plugin.load()
                try:
                    plugin_manifest = plugin_entry_point.plugin_manifest()

                    # this ignores what the plugin_manifest() function might
                    # put into source_files in favor of using the path to the
                    # python package as the unique identifier

                    manifest_path = os.path.abspath(
                        plugin_entry_point.__file__
                    )

                    if manifest_path in result.source_files:
                        continue

                    plugin_manifest.source_files = [manifest_path]
                    plugin_manifest._update_plugin_source(manifest_path)

                except AttributeError:
                    name = plugin_entry_point.__name__

                    try:
                        filepath = resources.files(name) / "plugin_manifest.json"
                    except AttributeError:
                        # For python <= 3.7
                        with resources.path(name, "plugin_manifest.json") as p:
                            filepath = Path(p)

                    if filepath.as_posix() in result.source_files:
                        continue

                    plugin_manifest = core.deserialize_json_from_string(
                        filepath.read_text()
                    )
                    plugin_manifest._update_plugin_source(filepath.as_posix())
                    plugin_manifest.source_files.append(filepath.as_posix())

            except Exception as e:
                logging.exception(
                    f"could not load plugin: {plugin_name}.  Exception is: {e}"
                )
                continue

            result.extend(plugin_manifest)
    else:
        logging.debug(
            "OTIO_DISABLE_ENTRYPOINTS_PLUGINS is set. "
            "Entry points plugings have been disabled."
        )

    # the builtin plugin manifest
    try:
        builtin_manifest_path = (
            resources.files("opentimelineio.adapters")
            / "builtin_adapters.plugin_manifest.json"
        ).as_posix()
    except AttributeError:
        # For python <= 3.7
        with resources.path(
            "opentimelineio.adapters",
            "builtin_adapters.plugin_manifest.json"
        ) as p:
            builtin_manifest_path = p.as_posix()

    if os.path.abspath(builtin_manifest_path) not in result.source_files:
        plugin_manifest = manifest_from_file(builtin_manifest_path)
        result.extend(plugin_manifest)

    # the contrib plugin manifest (located in the opentimelineio_contrib package)
    try:
        try:
            contrib_manifest_path = (
                resources.files("opentimelineio_contrib.adapters")
                / "contrib_adapters.plugin_manifest.json"
            ).as_posix()
        except AttributeError:
            # For python <= 3.7
            with resources.path(
                "opentimelineio_contrib.adapters",
                "contrib_adapters.plugin_manifest.json"
            ) as p:
                contrib_manifest_path = p.as_posix()

    except ModuleNotFoundError:
        logging.debug("no opentimelineio_contrib.adapters package found")

    else:
        if os.path.abspath(contrib_manifest_path) not in result.source_files:
            contrib_manifest = manifest_from_file(contrib_manifest_path)
            result.extend(contrib_manifest)

    # force the schemadefs to load and add to schemadef module namespace
    for s in result.schemadefs:
        s.module()

    return result


def ActiveManifest(force_reload=False):
    """Return the fully resolved plugin manifest."""

    global _MANIFEST
    if not _MANIFEST or force_reload:
        _MANIFEST = load_manifest()

    return _MANIFEST


repository: tmr232/ida-plugins
file: quick_copy/quick_copy.py
description: IDA Pro Plugins
wanted_name: Quick Copy
comment: Quickly copy addresses and instructions
created_at: 2015-05-27 12:35:29+00:00
pushed_at: 2015-05-31 13:05:09+00:00
url:
https://github.com/tmr232/ida-plugins/blob/c333ebecc9b3d641cdfdc9ab0d66f73e0b551fc3/quick_copy/quick_copy.py

content:
import idaapi
import clipboard
import sark


def copy_current_address():
    start, end = sark.get_selection()
    clipboard.copy("0x{:08X}".format(start))


def copy_current_selection():
    start, end = sark.get_selection()
    buffer = sark.data.read_memory(start, end)
    clipboard.copy(buffer.encode("hex-bytes"))


class QuickCopy(idaapi.plugin_t):
    flags = 0
    comment = "Quickly copy addresses and instructions"
    help = "Quickly copy addresses and instructions"
    wanted_name = "Quick Copy"
    wanted_hotkey = ""

    def init(self):
        self.hotkeys = []
        self.hotkeys.append(idaapi.add_hotkey("Ctrl+Alt+C", copy_current_address))
        self.hotkeys.append(idaapi.add_hotkey("Ctrl+Shift+C", copy_current_selection))
        return idaapi.PLUGIN_KEEP

    def term(self):
        for hotkey in self.hotkeys:
            idaapi.del_hotkey(hotkey)

    def run(self, arg):
        pass


def PLUGIN_ENTRY():
    return QuickCopy()


repository: andreafioraldi/IdaMSDNHelp
file: IdaMSDNHelp.py
description: IdaPython plugin to open MSDN Help Search page
wanted_name: IDA MSDN Help
created_at: 2017-09-19 10:44:22+00:00
pushed_at: 2017-09-19 10:44:50+00:00
url:
https://github.com/andreafioraldi/IdaMSDNHelp/blob/a44def08ba055d1450a1c3cd8c2a0b1b226568ba/IdaMSDNHelp.py

content:
__author__ = "Andrea Fioraldi"
__copyright__ = "Copyright 2017, Andrea Fioraldi"
__license__ = "MIT"
__email__ = "andreafioraldi@gmail.com"

import idaapi
import webbrowser

MENU_PATH = 'Edit/Other'
class IdaMSDNHelpPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = ""

    help = "IdaMSDNHelp: MSDN Help for IDA Pro"
    wanted_name = "IDA MSDN Help"
    wanted_hotkey = "Alt-R"

    def init(self):
        r = idaapi.add_menu_item(MENU_PATH, 'Open MSDN Search', '', 1, self.openHelp, tuple())
        if r is None:
            idaapi.msg("IdaMSDNHelp: add menu failed!\n")
        idaapi.msg("IdaMSDNHelp: initialized, opening webbrowser with '%s' command\n" % webbrowser.get().name)
        return idaapi.PLUGIN_KEEP

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

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

    def openHelp(self):
        webbrowser.open("https://social.msdn.microsoft.com/search/en-US")

    def query(self):
        func = idaapi.get_highlighted_identifier()
        url = "https://social.msdn.microsoft.com/search/en-US?query=" + func
        webbrowser.open(url)

def PLUGIN_ENTRY():
    return IdaMSDNHelpPlugin()


repository: laqieer/laqieer-fehack
file: sync_to_febuilder_disassembler/sync_to_febuilder_disassembler.py
description: Miscellaneous Fire Emblem (Game Boy Advance) hacking stuff
wanted_name: Sync to FEBuilderGBA
comment: Share knowledge with FEBuilderGBA disassembler
created_at: 2016-12-10 06:01:14+00:00
pushed_at: 2024-09-14 11:45:51+00:00
url:
https://github.com/laqieer/laqieer-fehack/blob/578d9d96bfcc34c204eace58d436298e6c226b3e/sync_to_febuilder_disassembler/sync_to_febuilder_disassembler.py

content:
#coding=utf-8

'''
同步IDA数据库信息到FEBuilderGBA的asmmap.
synchronize info in IDA database to FEBuilderGBA disassembler.
by laqieer
2018/02/24
'''

from idaapi import *
from idc import *
from ctypes import *
import os.path
import time

class sync_to_febuilder_disassembler_plugin_t(plugin_t):
	flags = 0
	wanted_name = "Sync to FEBuilderGBA"
	wanted_hotkey = "Alt-8"
	comment = "Share knowledge with FEBuilderGBA disassembler"
	help = "Created by laqieer.\nVisit https://github.com/laqieer for the latest version."
	
	def config(self):
		dialog = '''STARTITEM 0
Config

ROM Info
<##Generation##FE6:R>
<FE7:R>
<FE8:R>>
<##Version##JP:R>
<US:R>>
User Setting
<##Language##Japanese:R>
<English:R>
<Chinese:R>
<Germany:R>>
'''
		gen = c_int(0)
		ver = c_int(0)
		lan = c_int(0)
		while ask_form(dialog, byref(gen), byref(ver), byref(lan)) == 0:
			if ask_yn(0, "Are you really hacking GBA Fire Emblem?") == 0:
				self.n.altset(1, 2);
				break;
		self.n.altset(2, gen.value)
		self.n.altset(3, ver.value)
		self.n.altset(4, lan.value)
	
	def set_asmmap(self):
		#asmmap_path = ""
		asmmap_path = self.n.supval(1)
		#if asmmap_path == "":
		if asmmap_path is None:
		#if self.n.supstr(1, asmmap_path) == 0:
			febuilder_path = 0
			while febuilder_path == 0:
				febuilder_path = ask_file(0, "FEBuilderGBA.exe", "Where is FEBuilderGBA?")
			febuilder_dir, fn = os.path.split(febuilder_path)
			asmmap_path = febuilder_dir + r'\config\data\asmmap_addition_FE'
			gen = self.n.altval(2)
			#ver = self.n.altval(3)
			lan = self.n.altval(4)
			asmmap_path += str(gen + 6)
			if lan == 0:
				asmmap_path += '.txt'
			elif lan == 1:
				asmmap_path += '.en.txt'
			elif lan == 2:
				asmmap_path += '.zh.txt'
			elif lan == 3:
				asmmap_path += '.de.txt'
			else:
				asmmap_path += '.xx.txt'
			self.n.supset(1, asmmap_path)
		print(asmmap_path)
		self.asmmap = open(asmmap_path, 'w')
		while not self.asmmap:
			asmmap_path = 0
			while asmmap_path == 0:
				asmmap_path = ask_file(1, "asmmap_addition_FE*.txt", "Sync to asmmap")
			self.asmmap = open(asmmap_path, 'w')
		self.asmmap.write("// This file is generated by sync_to_nogba_debugger.idc and synchronous with \"" + get_idb_path() + "\".\n")
		self.asmmap.write("// Last sync at " + time.asctime() +"\n\n")
	
	def init(self):
		fn, ext = os.path.splitext(get_root_filename())
		if ext not in ('.gba', '.elf') : return PLUGIN_SKIP
		#if GetProcessorName() != 'ARM' : return PLUGIN_SKIP
		self.n = netnode("$ gba_fe_info", 0, True)
		if self.n.altval(1) == 1: return PLUGIN_OK
		if self.n.altval(1) == 2: return PLUGIN_SKIP
		# self.n.altset(1, AskYN(0, "Are you hacking GBA Fire Emblem?"))
		# if self.n.altval(1) == 0: return PLUGIN_SKIP
		if ask_yn(0, "Are you hacking GBA Fire Emblem?") == 0:
			self.n.altset(1, 2)
			return PLUGIN_SKIP
		self.n.altset(1, 1)
		self.config()
		return PLUGIN_OK
		
	def term(self):
		pass
		
	def run(self,arg):
		self.set_asmmap()
		isUS = self.n.altval(3)
		symbols = Names()
		for ea,name in symbols:
			#self.asmmap.write("%08X\t%s\t%s\t" % (ea, name, GetType(ea)))
			self.asmmap.write("%08X\t%s\t" % (ea, name))
			if GetType(ea) is not None: self.asmmap.write(GetType(ea) + '\t')
			if GetFunctionCmt(ea,1) is not "": self.asmmap.write(GetFunctionCmt(ea,1).replace('\n','. ') + '\t')
			if GetConstCmt(ea,1) is not "": self.asmmap.write(GetConstCmt(ea,1).replace('\n','. ') + '\t')
			if GetConstCmt(ea,0) is not "": self.asmmap.write(GetConstCmt(ea,0).replace('\n','. ') + '\t')
			if isUS:
				self.asmmap.write('{U}\n')
			else:
				self.asmmap.write('{J}\n')
		self.asmmap.close()
		
	# def PLUGIN_ENTRY():
		# return sync_to_febuilder_disassembler_plugin_t()
		
def PLUGIN_ENTRY():
    return sync_to_febuilder_disassembler_plugin_t()


repository: es3n1n/ida-wakatime-py
file: wakatime.py
description: WakaTime integration for IDA Pro
url: https://github.com/es3n1n/ida-wakatime-py/blob/68a13d1fbf13f901c735b266715ac5e467517270/wakatime.py
wanted_name: WakaTime
comment: WakaTime integration for IDA Pro by @es3n1n
created_at: 2022-05-18 15:46:14+00:00
pushed_at: 2024-09-11 12:47:39+00:00
content:
# -*- coding: utf-8 -*-
""" ==========================================================
File:        WakaTime.py
Description: Automatic time tracking for IDA Pro.
Maintainer:  Arsenii Esenin <me@es3n.in>
License:     MIT, see LICENSE for more details.
Website:     https://github.com/es3n1n/ida-wakatime-py
==========================================================="""

import json
import os
import platform
import re
import shutil
import ssl
import subprocess
import sys
import threading
import time
import traceback
import webbrowser
from pathlib import Path
from subprocess import PIPE
from subprocess import STDOUT
from zipfile import ZipFile

import idaapi
import idc
from ida_kernwin import ask_buttons
from ida_kernwin import ask_str

try:
    import Queue as queue  # py2  # noqa
except ImportError:
    import queue  # py3

try:  # @note: @es3n1n: py2
    from ConfigParser import SafeConfigParser as ConfigParser  # noqa
    from ConfigParser import Error as ConfigParserError  # noqa
except ImportError:  # @note: @es3n1n: py3
    from configparser import ConfigParser  # noqa
    from configparser import Error as ConfigParserError  # noqa

try:  # @note: @es3n1n: py2
    from urllib2 import Request  # noqa
    from urllib2 import urlopen  # noqa
    from urllib2 import HTTPError  # noqa
except ImportError:  # @note: @es3n1n: py3
    from urllib.request import Request  # noqa
    from urllib.request import urlopen  # noqa
    from urllib.error import HTTPError  # noqa

# @note: @es3n1n: init py2-py3 stuff
is_py2 = (sys.version_info[0] == 2)
is_py3 = (sys.version_info[0] == 3)

# @note: @es3n1n: plugin-related stuff
VERSION = "1.1"

# @note: @es3n1n: ida-related stuff
ida_ver = idaapi.get_kernel_version()
ida_major, ida_minor = list(map(int, ida_ver.split(".")))
using_ida7api = (ida_major > 6)

# @note: @es3n1n: Log Levels
DEBUG = 'DEBUG'
INFO = 'INFO'
WARNING = 'WARNING'
ERROR = 'ERROR'

# @note: @es3n1n: wakatime stuff
IS_WIN = platform.system() == 'Windows'
HOME_FOLDER = os.path.realpath(
    os.environ.get('WAKATIME_HOME') or os.path.expanduser('~')
)
RESOURCES_FOLDER = os.path.join(HOME_FOLDER, '.wakatime')
CONFIG_FILE = os.path.join(HOME_FOLDER, '.wakatime.cfg')
INTERNAL_CONFIG_FILE = os.path.join(HOME_FOLDER, '.wakatime-internal.cfg')
GITHUB_RELEASES_STABLE_URL = 'https://api.github.com/repos/wakatime/wakatime' \
                             '-cli/releases/latest'
GITHUB_DOWNLOAD_PREFIX = 'https://github.com/wakatime/wakatime-cli/releases' \
                         '/download'
LAST_HEARTBEAT = {
    'time': 0,
    'file': None,
    'is_write': False,
}
LAST_HEARTBEAT_SENT_AT = 0
LAST_FETCH_TODAY_CODING_TIME = 0
FETCH_TODAY_DEBOUNCE_COUNTER = 0
FETCH_TODAY_DEBOUNCE_SECONDS = 60
LATEST_CLI_VERSION = None
WAKATIME_CLI_VERSION = None
WAKATIME_CLI_LOCATION = None
HEARTBEATS = queue.Queue()
HEARTBEAT_FREQUENCY = 2  # minutes between logging heartbeat when editing file
SEND_BUFFER_SECONDS = 30  # seconds between sending buffered heartbeats to API


# @note: @es3n1n: Utils
class Popen(subprocess.Popen):
    """Patched Popen to prevent opening cmd window on Windows platform."""

    def __init__(self, *args, **kwargs):
        if IS_WIN:
            startupinfo = kwargs.get('startupinfo')
            try:
                startupinfo = startupinfo or subprocess.STARTUPINFO()
                startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            except AttributeError:
                pass
            kwargs['startupinfo'] = startupinfo
        super(Popen, self).__init__(*args, **kwargs)


class Config:
    path = Path(__file__).parent.resolve().absolute() / 'wakatime.config.json'
    _cached = None

    @classmethod
    def read(cls):
        if cls._cached is not None:
            return cls._cached

        if not cls.path.exists():
            cls._cached = dict()
            return cls._cached

        try:
            cls._cached = json.loads(cls.path.read_text())
        except json.JSONDecodeError:
            cls._cached = dict()
        return cls._cached

    @classmethod
    def write(cls, value):
        cls.path.write_text(json.dumps(value))

    @classmethod
    def get_var(cls, key, default=None):
        return cls.read().get(key, default)

    @classmethod
    def set_var(cls, key, value):
        cfg = cls.read()
        cfg[key] = value
        cls.write(cfg)


def log(lvl, message, *args, **kwargs):
    if lvl == DEBUG and Config.get_var('debug', 'false') == 'false':
        return

    msg = message

    if len(args) > 0:
        msg = message.format(*args)
    elif len(kwargs) > 0:
        msg = message.format(**kwargs)

    print('WakaTime > {lvl} > {msg}'.format(lvl=lvl, msg=msg))


def parse_config_file(config_file):
    """Returns a configparser.SafeConfigParser instance with configs
    read from the config file. Default location of the config file is
    at ~/.wakatime.cfg.
    """

    kwargs = {} if is_py2 else {'strict': False}
    configs = ConfigParser(**kwargs)
    try:
        with open(config_file, 'r', encoding='utf-8') as fh:
            try:
                if is_py2:
                    configs.readfp(fh)  # noqa
                else:
                    configs.read_file(fh)
                return configs
            except ConfigParserError:
                log(ERROR, traceback.format_exc())
                return None
    except IOError:
        return


def last_modified_and_versions(configs):
    last_modified, last_version = None, None
    if configs.has_option('internal', 'cli_version'):
        last_version = configs.get('internal', 'cli_version')
    if last_version and configs.has_option('internal',
                                           'cli_version_last_modified'):
        last_modified = configs.get('internal', 'cli_version_last_modified')
    if last_modified and last_version and extract_version(last_version):
        return last_modified, last_version
    return None, None


def extract_version(text):
    pattern = re.compile(r"([0-9]+\.[0-9]+\.[0-9]+)")
    match = pattern.search(text)
    if match:
        return 'v{ver}'.format(ver=match.group(1))
    return None


def request(url, last_modified=None):
    req = Request(url)

    if last_modified:
        req.add_header('If-Modified-Since', last_modified)

    try:
        resp = urlopen(req)
        headers = dict(resp.getheaders()) if is_py2 else resp.headers
        return headers, resp.read(), resp.getcode()
    except HTTPError as err:
        if err.code == 304:
            return None, None, 304
        if is_py2:
            with SSLCertVerificationDisabled():
                try:
                    resp = urlopen(req)
                    headers = dict(
                        resp.getheaders()) if is_py2 else resp.headers
                    return headers, resp.read(), resp.getcode()
                except HTTPError as err2:
                    if err2.code == 304:
                        return None, None, 304
                    log(ERROR, err.read().decode())
                    log(ERROR, err2.read().decode())
                    raise
                except IOError:
                    raise
        log(ERROR, err.read().decode())
        raise
    except IOError:
        if is_py2:
            with SSLCertVerificationDisabled():
                try:
                    resp = urlopen(url)
                    headers = dict(
                        resp.getheaders()) if is_py2 else resp.headers
                    return headers, resp.read(), resp.getcode()
                except HTTPError as err:
                    if err.code == 304:
                        return None, None, 304
                    log(ERROR, err.read().decode())
                    raise
                except IOError:
                    raise
        raise


def download(url, file_path):
    req = Request(url)

    with open(file_path, 'wb') as fh:
        try:
            resp = urlopen(req)
            fh.write(resp.read())
        except HTTPError as err:
            if err.code == 304:
                return None, None, 304
            if is_py2:
                with SSLCertVerificationDisabled():
                    try:
                        resp = urlopen(req)
                        fh.write(resp.read())
                        return
                    except HTTPError as err2:
                        log(ERROR, err.read().decode())
                        log(ERROR, err2.read().decode())
                        raise
                    except IOError:
                        raise
            log(ERROR, err.read().decode())
            raise
        except IOError:
            if is_py2:
                with SSLCertVerificationDisabled():
                    try:
                        resp = urlopen(url)
                        fh.write(resp.read())
                        return
                    except HTTPError as err:
                        log(ERROR, err.read().decode())
                        raise
                    except IOError:
                        raise
            raise


def is_symlink(path):
    try:
        return os.is_symlink(path)  # noqa
    except:  # noqa
        return False


def get_cli_location():
    global WAKATIME_CLI_LOCATION

    if not WAKATIME_CLI_LOCATION:
        binary = 'wakatime-cli-{osname}-{arch}{ext}'.format(
            osname=platform.system().lower(),
            arch=architecture(),
            ext='.exe' if IS_WIN else '',
        )
        WAKATIME_CLI_LOCATION = os.path.join(RESOURCES_FOLDER, binary)

    return WAKATIME_CLI_LOCATION


def create_symlink():
    link = os.path.join(RESOURCES_FOLDER, 'wakatime-cli')
    if IS_WIN:
        link = link + '.exe'
    elif os.path.exists(link) and is_symlink(link):
        return  # don't re-create symlink on Unix-like platforms

    try:
        os.symlink(get_cli_location(), link)
    except:  # noqa
        try:
            shutil.copy2(get_cli_location(), link)
            if not IS_WIN:
                os.chmod(link, 509)  # 755
        except:  # noqa
            log(ERROR, traceback.format_exc())


class SSLCertVerificationDisabled(object):

    def __enter__(self):
        self.original_context = ssl._create_default_https_context  # noqa
        ssl._create_default_https_context = ssl._create_unverified_context  # noqa

    def __exit__(self, *args, **kwargs):
        ssl._create_default_https_context = self.original_context


def get_latest_cli_version():
    global LATEST_CLI_VERSION

    if LATEST_CLI_VERSION:
        return LATEST_CLI_VERSION

    configs, last_modified, last_version = None, None, None
    try:
        configs = parse_config_file(INTERNAL_CONFIG_FILE)
        if configs:
            last_modified, last_version = last_modified_and_versions(configs)
    except:  # noqa
        log(DEBUG, traceback.format_exc())

    try:
        headers, contents, code = request(GITHUB_RELEASES_STABLE_URL,
                                          last_modified=last_modified)

        if code == 304:
            LATEST_CLI_VERSION = last_version
            return last_version

        data = json.loads(contents.decode('utf-8'))

        ver = data['tag_name']

        if configs:
            last_modified = headers.get('Last-Modified')
            if not configs.has_section('internal'):
                configs.add_section('internal')
            configs.set('internal', 'cli_version', ver)
            configs.set('internal', 'cli_version_last_modified', last_modified)
            with open(INTERNAL_CONFIG_FILE, 'w', encoding='utf-8') as fh:
                configs.write(fh)

        LATEST_CLI_VERSION = ver
        return ver
    except:  # noqa
        log(ERROR, traceback.format_exc())
        return None


def cli_download_url():
    return '{prefix}/{version}/wakatime-cli-{osname}-{arch}.zip'.format(
        prefix=GITHUB_DOWNLOAD_PREFIX,
        version=get_latest_cli_version(),
        osname=platform.system().lower(),
        arch=architecture(),
    )


class DownloadCLI(threading.Thread):

    def run(self):
        log(INFO, 'Downloading wakatime-cli...')

        if os.path.isdir(os.path.join(RESOURCES_FOLDER, 'wakatime-cli')):
            shutil.rmtree(os.path.join(RESOURCES_FOLDER, 'wakatime-cli'))

        if not os.path.exists(RESOURCES_FOLDER):
            os.makedirs(RESOURCES_FOLDER)

        try:
            url = cli_download_url()
            log(INFO, 'Downloading wakatime-cli from {url}'.format(url=url))
            zip_file = os.path.join(RESOURCES_FOLDER, 'wakatime-cli.zip')
            download(url, zip_file)

            if is_cli_installed():
                try:
                    os.remove(get_cli_location())
                except:  # noqa
                    log(DEBUG, traceback.format_exc())

            log(INFO, 'Extracting wakatime-cli...')
            with ZipFile(zip_file) as zf:
                zf.extractall(RESOURCES_FOLDER)

            if not IS_WIN:
                os.chmod(get_cli_location(), 509)  # 755

            try:
                os.remove(os.path.join(RESOURCES_FOLDER, 'wakatime-cli.zip'))
            except:  # noqa
                log(ERROR, traceback.format_exc())
        except:  # noqa
            log(ERROR, traceback.format_exc())

        create_symlink()

        log(INFO, 'Finished extracting wakatime-cli.')


def architecture():
    arch = platform.machine() or platform.processor()
    if arch == 'armv7l':
        return 'arm'
    if arch == 'aarch64':
        return 'arm64'
    if 'arm' in arch:
        return 'arm64' if sys.maxsize > 2 ** 32 else 'arm'
    return 'amd64' if sys.maxsize > 2 ** 32 else '386'


def is_cli_installed():
    return os.path.exists(get_cli_location())


def is_cli_latest():
    global WAKATIME_CLI_VERSION

    if not is_cli_installed():
        return False

    args = [get_cli_location(), '--version']
    try:
        stdout, stderr = Popen(args, stdout=PIPE, stderr=PIPE).communicate()
    except:  # noqa
        return False
    stdout = (stdout or b'') + (stderr or b'')
    local_ver = extract_version(stdout.decode('utf-8'))
    if not local_ver:
        log(WARNING, 'Local wakatime-cli version not found.')
        return False

    WAKATIME_CLI_VERSION = local_ver

    remote_ver = get_latest_cli_version()

    if not remote_ver:
        return True

    if remote_ver == local_ver:
        return True

    log(INFO, 'Found an updated wakatime-cli %s' % remote_ver)
    return False


def obfuscate_apikey(command_list):
    cmd = list(command_list)
    apikey_index = None
    for num in range(len(cmd)):
        if cmd[num] == '--key':
            apikey_index = num + 1
            break
    if apikey_index is not None and apikey_index < len(cmd):
        v = cmd[apikey_index][-4:]
        cmd[apikey_index] = 'XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXX' + v
    return cmd


def enough_time_passed(now, is_write):
    if now - LAST_HEARTBEAT['time'] > HEARTBEAT_FREQUENCY * 60:
        return True
    if is_write and now - LAST_HEARTBEAT['time'] > 2:
        return True
    return False


class ApiKey(object):
    _key = None

    def read(self):
        if self._key:
            return self._key

        key = Config.get_var('api_key')
        if key:
            self._key = key
            return self._key

        try:
            configs = parse_config_file(CONFIG_FILE)
            if configs:
                if configs.has_option('settings', 'api_key'):
                    key = configs.get('settings', 'api_key')
                    if key:
                        self._key = key
                        return self._key
        except:  # noqa
            pass

        return self._key

    def write(self, key):
        self._key = key
        Config.set_var('api_key', key)


APIKEY = ApiKey()


def build_heartbeat(entity=None, timestamp=None, is_write=None,
                    project=None, idb_path=None):
    """Returns a dict for passing to wakatime-cli as arguments."""

    heartbeat = {
        'entity': entity,
        'timestamp': timestamp,
        'is_write': is_write,
        'idb_path': idb_path
    }

    if project:
        heartbeat['alternate_project'] = project

    return heartbeat


class SendHeartbeatsThread(threading.Thread):
    """Non-blocking thread for sending heartbeats to api.
    """

    def __init__(self, heartbeat):
        threading.Thread.__init__(self)

        self.extra_heartbeats = None
        self.api_key = APIKEY.read() or ''
        self.ignore = None
        self.include = None
        self.hidefilenames = None
        self.proxy = None

        self.heartbeat = heartbeat
        self.has_extra_heartbeats = False

    def add_extra_heartbeats(self, extra_heartbeats):
        self.has_extra_heartbeats = True
        self.extra_heartbeats = extra_heartbeats

    def run(self):
        """Running in background thread."""

        self.send_heartbeats()

    def send_heartbeats(self):
        heartbeat = build_heartbeat(**self.heartbeat)
        ua = 'ida-pro/%s ida-pro-wakatime/%s' % (ida_ver, VERSION)
        cmd = [
            get_cli_location(),
            '--entity', heartbeat['entity'],
            '--time', str('%f' % heartbeat['timestamp']),
            '--plugin', ua,
            '--local-file', heartbeat['idb_path'],
            '--alternate-language', 'IDA',
            '--category', 'code reviewing'  # @note: @es3n1n: lmao
        ]
        if self.api_key:
            cmd.extend(
                ['--key', str(bytes.decode(self.api_key.encode('utf8')))])
        if heartbeat['is_write']:
            cmd.append('--write')
        if heartbeat.get('alternate_project'):
            cmd.extend(['--alternate-project', heartbeat['alternate_project']])
        if heartbeat.get('lineno') is not None:
            cmd.extend(['--lineno', '{0}'.format(heartbeat['lineno'])])
        if heartbeat.get('cursorpos') is not None:
            cmd.extend(['--cursorpos', '{0}'.format(heartbeat['cursorpos'])])
        if heartbeat.get('lines_in_file') is not None:
            cmd.extend(
                ['--lines-in-file', '{0}'.format(heartbeat['lines_in_file'])])
        if self.hidefilenames:
            cmd.append('--hidefilenames')
        if self.proxy:
            cmd.extend(['--proxy', self.proxy])
        stdin = None
        inp = None
        if self.has_extra_heartbeats:
            cmd.append('--extra-heartbeats')
            stdin = PIPE
            extra_heartbeats = json.dumps(
                [build_heartbeat(**x) for x in self.extra_heartbeats])
            inp = "{0}\n".format(extra_heartbeats).encode('utf-8')

        try:
            process = Popen(cmd, stdin=stdin, stdout=PIPE, stderr=STDOUT)
            output, _err = process.communicate(input=inp)
            retcode = process.poll()
            if retcode and retcode not in [102, 112]:
                log(ERROR, 'Error #{0}'.format(retcode))
            if output:
                log(ERROR, 'wakatime-core output: {0}'.format(output))
        except:  # noqa
            log(ERROR, sys.exc_info()[1])


def handle_activity(callback_name, entity=None, is_write=False):
    project = os.path.basename(idc.get_idb_path())
    if not entity:
        entity = project

    timestamp = time.time()
    if entity != LAST_HEARTBEAT['file'] or enough_time_passed(timestamp,
                                                              is_write):
        append_heartbeat(entity, timestamp, is_write, project)
        log(DEBUG, 'heartbeat[{}, {}, {}, {}, {}]'.format(callback_name,
                                                          entity,
                                                          int(timestamp),
                                                          is_write, project))

    return 0


def append_heartbeat(entity, timestamp, is_write, project):
    global LAST_HEARTBEAT

    # add this heartbeat to queue
    heartbeat = {
        'entity': entity,
        'timestamp': timestamp,
        'is_write': is_write,
        'project': project,
        'idb_path': idc.get_idb_path()
    }
    HEARTBEATS.put_nowait(heartbeat)

    # make this heartbeat the LAST_HEARTBEAT
    LAST_HEARTBEAT = {
        'file': entity,
        'time': timestamp,
        'is_write': is_write,
    }

    # process the queue of heartbeats in the future
    process_queue(timestamp)


def process_queue(timestamp):
    global LAST_HEARTBEAT_SENT_AT

    if not is_cli_installed():
        return

    # Prevent sending heartbeats more often than SEND_BUFFER_SECONDS
    now = int(time.time())
    if timestamp != LAST_HEARTBEAT['time'] \
            and LAST_HEARTBEAT_SENT_AT > now - SEND_BUFFER_SECONDS:
        return
    LAST_HEARTBEAT_SENT_AT = now

    try:
        heartbeat = HEARTBEATS.get_nowait()
    except queue.Empty:
        return

    has_extra_heartbeats = False
    extra_heartbeats = []
    try:
        while True:
            extra_heartbeats.append(HEARTBEATS.get_nowait())
            has_extra_heartbeats = True
    except queue.Empty:
        pass

    thread = SendHeartbeatsThread(heartbeat)
    if has_extra_heartbeats:
        thread.add_extra_heartbeats(extra_heartbeats)
    thread.start()


class IDACallbacks(idaapi.Hexrays_Hooks):
    # @note: @es3n1n: New pseudocode view has been opened.
    def open_pseudocode(self, vu):
        return handle_activity('open_pseudocode',
                               idc.get_func_name(vu.cfunc.body.ea), False)

    # @note: @es3n1n: Existing pseudocode text has been refreshed
    def refresh_pseudocode(self, vu):
        return handle_activity('refresh_pseudocode',
                               idc.get_func_name(vu.cfunc.body.ea), False)

    # @note: @es3n1n: Keyboard has been hit.
    def keyboard(self, vu, *args):
        return handle_activity('keyboard',
                               idc.get_func_name(vu.cfunc.body.ea), True)

    # @note: @es3n1n: Mouse right click
    def right_click(self, vu):
        return handle_activity('right_click',
                               idc.get_func_name(vu.cfunc.body.ea), False)

    # @note: @es3n1n: Mouse double click.
    def double_click(self, vu, *args):
        return handle_activity('double_click',
                               idc.get_func_name(vu.cfunc.body.ea), False)

    # @note: @es3n1n: Current cursor position has been changed.
    def curpos(self, vu):
        return handle_activity('curpos',
                               idc.get_func_name(vu.cfunc.body.ea), False)

    # @note: @es3n1n: Local variable got renamed.
    def lvar_name_changed(self, vu, *args):
        return handle_activity('lvar_name_changed',
                               idc.get_func_name(vu.cfunc.body.ea), True)

    # @note: @es3n1n: Local variable type got changed.
    def lvar_type_changed(self, vu, *args):
        return handle_activity('lvar_type_changed',
                               idc.get_func_name(vu.cfunc.body.ea), True)

    # @note: @es3n1n: Local variable comment got changed.
    def lvar_cmt_changed(self, vu, *args):
        return handle_activity('lvar_cmt_changed',
                               idc.get_func_name(vu.cfunc.body.ea), True)

    # @note: @es3n1n: Local variable mapping got changed.
    def lvar_mapping_changed(self, vu, *args):
        return handle_activity('lvar_mapping_changed',
                               idc.get_func_name(vu.cfunc.body.ea), True)

    # @note: @es3n1n: Comment got changed.
    def cmt_changed(self, cfunc, *args):
        return handle_activity('cmt_changed',
                               idc.get_func_name(cfunc.body.ea), True)


def prompt_api_key():
    if APIKEY.read():
        return True

    default_value = 'https://wakatime.com/settings/account'
    if APIKEY.read():
        default_value = APIKEY.read()

    new_key = ask_str(default_value, 1488, 'Enter your wakatime API key: ')
    if not new_key:
        return

    log(INFO, 'Set api key!')
    APIKEY.write(new_key)


class WakaTimePlugin(idaapi.plugin_t):
    comment = "WakaTime integration for IDA Pro by @es3n1n"
    help = ""
    wanted_name = "WakaTime"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP
    callbacks = IDACallbacks()

    def init(self):
        log(INFO, 'Loaded')
        log(INFO, 'v{} by @es3n1n, 2022'.format(VERSION))

        if not is_cli_latest():
            th = DownloadCLI()
            th.start()

        prompt_api_key()
        self.callbacks.hook()

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    @staticmethod
    def run(*args):  # noqa
        dbg = Config.get_var('debug', "false")

        fmt = '''AUTOHIDE NONE
WakaTime integration for IDA Pro
Plugin version: v{}
Plugin author: @es3n1n
Wakatime version: {}
Debug mode: {}'''.format(VERSION, WAKATIME_CLI_VERSION, dbg)
        ret = ask_buttons('~D~ebug', '~G~it repo', '~C~ancel', -1, fmt)

        if ret == 1:
            dbg = "false" if dbg == "true" else "true"
            Config.set_var('debug', dbg)
            log(INFO, 'Set debug to: {}'.format(dbg))

        if ret == 0:
            log(INFO, 'Opening git repo')
            webbrowser.open_new('https://github.com/es3n1n/ida-wakatime-py')


def PLUGIN_ENTRY():  # noqa
    return WakaTimePlugin()


repository: HongThatCong/LazyIDA
file: LazyIDA.py
url: https://github.com/HongThatCong/LazyIDA/blob/217b8669453bfbf67d7c78c7942230b481a0f283/LazyIDA.py
created_at: 2020-08-09 11:52:45+00:00
pushed_at: 2024-03-14 07:12:15+00:00
content:
# pylint: disable=C0301,C0103,C0111

#
# License: Beerware license ;)
#
# Link update:
#   https://github.com/HongThatCong/LazyIDA
# Origin version:
#   https://github.com/L4ys/LazyIDA
#
# Version histories:
# 09/08/2020 - HTC
#   - fix bug read_selection
#   - fix bug parse_location, so Shift-G hotkey can jumpto any ea, name possibles
#   - change "C" hotkey to Shift-C, because C is duplicate function with IDA C and Ctrl-C hotkey
#     Shift-C will copy full name of a current highlight name
#   - add "Shfit-V" hotkey, paste name from clipboard to current highlight name or current addr (fullname)
# 10/08/2020
#   - Fix bug sanitize name from clipboard text - idea of Ngon Nguyen, code by HTC
# 30/08/2020
#   - Rearrange actions and hotkeys - HTC
#   - Refactor, clean code, add nhieu tinh nang. Nhieu qua, lam bieng liet ke, chiu kho xem code hay diff :D
# 31/01/2021
#   - Ver 1.0.5
# 20/12/2021:
#   - Add Revert IDA Decision
#   - Ver 1.0.6
# 26/12/2022:
#   - Port to Python 3
#   - Ver 1.1
#   - Add Shift-X, display refs to a register
#   - Flake8, refactor, fix some bugs
#
# This script is licensed under the "THE BEER-WARE LICENSE" (Revision 42) license :D
#

from __future__ import division
from __future__ import print_function

import os
import string
import re
import base64
from urllib.parse import quote

from struct import unpack

import idaapi
import idautils
import idc
import ida_kernwin
import ida_xref
import ida_name
import ida_loader

from PyQt5.Qt import QApplication

PLUGIN_NAME = "LazyIDA"
PLUGIN_VERSION = "1.1"
PLUGIN_POPUP = "LazyIDA/"

# Popup menus action names
ACTION_MENU_CONVERT = [f"lazyida:convert_{i}" for i in range(12)]
ACTION_MENU_SCAN_VUL = "lazyida:scan_vul"
ACTION_MENU_COPY_DATA = "lazyida:copy_data"     # added by merc
ACTION_MENU_COPY_STR = "lazyida:copystring"     # HTC
ACTION_MENU_DUMP_DATA = "lazyida:dump_data"     # HTC
ACTION_MENU_DUMP_SEG = "lazyida:dump_seg"
ACTION_MENU_XOR_DATA = "lazyida:xor_data"       # HTC & CatBui mod
ACTION_MENU_FILL_NOP = "lazyida:fill_nop"
ACTION_MENU_NOP_HIDER = "lazyida:nop_hider"     # added by HTC
ACTION_MENU_AUTO_OFF = "lazyida:turn_off_ida_decision"
ACTION_MENU_B64STD = "lazyida:base64std_decode"
ACTION_MENU_B64URL = "lazyida:base64url_decode"

#
# HTC: change actions name and hotkeys - 30/08/2020
# Available Shift- hotkeys: B C G H J K O Q V T Y Z
# All actions name and hotkeys, enable allways
# Action is tuple: name, label, shortcut, tooltip/hint, iconid
#
ACTION_HOTKEY_COPY_EA = ("lazyida:copy_ea", "Copy EA", "Shift-Y", "Copy current EA to clipboard", 0x1F)
ACTION_HOTKEY_COPY_RVA = ("lazyida:copy_rva", "Copy RVA", "Shift-Z", "Copy current RVA to clipboard", 0x1F)
ACTION_HOTKEY_COPY_FOFS = ("lazyida:copy_fofs", "Copy File Offset", "Shift-O", "Copy current file offset to clipboard", 0x1F)
ACTION_HOTKEY_COPY_NAME = ("lazyida:copy_name", "Copy highligh name", "Shift-C", "Copy current highlight full name to clipboard", 0x1F)
ACTION_HOTKEY_PASTE_NAME = ("lazyida:paste_name", "Paste to highligh name", "Shift-V", "Change the fullname of current highlight to clipboard text", 0x13)
ACTION_HOTKEY_GOTO_CLIP = ("lazyida:goto_clip", "Goto text in clipboard", "Shift-G", "Goto current name or ea in clipboard", 0x7D)
ACTION_HOTKEY_GOTO_FOFS = ("lazyida:goto_fofs", "Goto file offset", "Shift-J", "Goto file offset", 0x7D)  # IDA already had this action "JumpFileOffset"
ACTION_HOTKEY_GOTO_RVA = ("lazyida:goto_rva", "Goto RVA", "Alt-G", "Goto RVA", 0x7D)
ACTION_HOTKEY_SEARCH_GOOGLE = ("lazyida:search_google", "Search Google", "Ctrl-Shift-G", "Search Google", 0x21)
ACTION_HOTKEY_SEARCH_MSDOC = ("lazyida:search_msdoc", "Search MS Docs", "Ctrl-Shift-S", "Search MS Docs", 0x21)
ACTION_HOTKEY_SEARCH_BING = ("lazyida:search_bing", "Search Bing", "Ctrl-Shift-J", "Search Bing", 0x21)
ACTION_HOTKEY_SEARCH_GITHUB = ("lazyida:search_github", "Search Github", "Ctrl-Shift-H", "Search Github", 0x21)
ACTION_HOTKEY_OPEN_URL = ("lazyida:open_url", "Open URL", "",  "Open URL", 0)

ALL_HOTKEY_ACTIONS = (
    ACTION_HOTKEY_COPY_EA,
    ACTION_HOTKEY_COPY_RVA,
    ACTION_HOTKEY_COPY_FOFS,
    ACTION_HOTKEY_COPY_NAME,
    ACTION_HOTKEY_PASTE_NAME,
    ACTION_HOTKEY_GOTO_CLIP,
    ACTION_HOTKEY_GOTO_FOFS,
    ACTION_HOTKEY_GOTO_RVA,
    ACTION_HOTKEY_SEARCH_GOOGLE,
    ACTION_HOTKEY_SEARCH_MSDOC,
    ACTION_HOTKEY_SEARCH_BING,
    ACTION_HOTKEY_SEARCH_GITHUB,
    ACTION_HOTKEY_OPEN_URL,
)

# Decompiler view hotkeys
# HTC: V duplicate with HexCodeXplorer plugin ctree_item_vew, so change to Alt-R
ACTION_HX_REMOVE_RET_TYPE = ("lazyida::hx_removerettype", "Remove return type", "Alt-R", "Set return type of current function to void")


def calc_lazy_bits():
    info = idaapi.get_inf_structure()
    if info.is_64bit():
        return 64
    elif info.is_32bit():
        return 32
    else:
        return 16


LAZY_ARCH = idaapi.ph_get_id()
LAZY_BITS = calc_lazy_bits()

IS_BE = idaapi.get_inf_structure().is_be()
CAN_NOP = LAZY_ARCH == idaapi.PLFM_386
ENABLE_REMOVE_RETTYPE = "hx:RemoveArg" not in idaapi.get_registered_actions()  # HexRays 7.5+


def u16(x):
    return unpack("<H", x)[0]


def u32(x):
    return unpack("<I", x)[0]


def u64(x):
    return unpack("<Q", x)[0]


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


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


class RangeForm(idaapi.Form):
    """
    Form to prompt for selecting a range: start ea, size or end ea
    """
    def __init__(self, start_ea, end_ea):
        self.intStartEA = None
        self.intSize = None
        self.intEndEA = None

        idaapi.Form.__init__(self, r"""Confirm or change the selected range

{FormChangeCb}
<##Start EA             :{intStartEA}>
<##Size                 :{intSize}>
<##End EA (not included):{intEndEA}>
""",
{
    'intStartEA': idaapi.Form.NumericInput(swidth=20, tp=idaapi.Form.FT_HEX, value=start_ea),
    'intSize': idaapi.Form.NumericInput(swidth=20, tp=idaapi.Form.FT_HEX, value=end_ea - start_ea),
    'intEndEA': idaapi.Form.NumericInput(swidth=20, tp=idaapi.Form.FT_HEX, value=end_ea),
    'FormChangeCb': idaapi.Form.FormChangeCb(self.OnFormChange),
})

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

        start = self.GetControlValue(self.intStartEA)
        size = self.GetControlValue(self.intSize)
        if start and size:
            self.SetControlValue(self.intEndEA, start + size)

        return 1


def plg_print(smsg):
    print(f"[{PLUGIN_NAME}] {smsg}")


def is_valid_addr(ea):
    return idc.get_inf_attr(idc.INF_MIN_EA) <= ea <= idc.get_inf_attr(idc.INF_MAX_EA)


def parse_location(text):
    """Parse text to hex ea or try to get a valid name"""
    plg_print(f'Clipboard text is "{text}"')

    if not text:
        return idaapi.BADADDR

    strs = re.findall(r"[\da-f]+", text, re.IGNORECASE)   # parse hex number
    if strs:
        for s in strs:
            try:
                ea = int(s, 16)
                if is_valid_addr(ea):
                    return ea
            except ValueError:
                pass

    # Fail to find a valid ea, assume text is a name
    ea = idc.get_name_ea_simple(text)
    if is_valid_addr(ea):
        return ea

    # Try to sanitize the name
    san = ida_name.validate_name(text, ida_name.VNT_IDENT)
    if san:
        ea = idc.get_name_ea_simple(san)
        if is_valid_addr(ea):
            return ea

    # Parse text into words, assume every words is a valid name
    strs = re.findall(r"\w+", text)
    if strs:
        for s in strs:
            ea = idc.get_name_ea_simple(s)
            if is_valid_addr(ea):
                return ea

    return idaapi.BADADDR


def lazy_read_selection():
    # HTC - Ignore the byte at end address
    sel, start, end = idaapi.read_range_selection(None)
    if not sel:
        if idc.get_item_size(idc.get_screen_ea()):
            start = idc.get_screen_ea()
            end = start + idc.get_item_size(start)
            sel = True

    if not sel:
        start = idaapi.BADADDR
        end = idaapi.BADADDR
    else:
        rFrm = RangeForm(start, end)
        rFrm.Compile()
        ok = rFrm.Execute()
        if ok == 1:
            # OK
            start = rFrm.intStartEA.value
            end = start + rFrm.intSize.value
        else:
            # Cancel
            sel = False

        rFrm.Free()

    # Ensure we have at least one byte
    if end <= start:
        sel = False

    return sel, start, end


def goto_clip_text():
    loc = parse_location(clip_text())
    if loc != idaapi.BADADDR:
        plg_print(f"Goto location 0x{loc:X}")
        idc.jumpto(loc)
        return 1

    plg_print("Failed to get a valid ea")
    return 0


def str2hex(txt):
    if not txt:
        return 0

    # is it a valid number ?
    val = 0
    if txt.endswith("h"):
        txt = txt[:-1]
    try:
        val = int(txt, 16)
    except ValueError:
        pass

    return val


def get_number_from_highlight():
    view = idaapi.get_current_viewer()
    thing = ida_kernwin.get_highlight(view)
    if thing and thing[1]:
        # we have a highlight
        val = str2hex(thing[0])
        if val != 0:
            # have a valid hex number
            return val

    return str2hex(clip_text())


def get_screen_module():
    if not idaapi.is_debugger_on():
        return (idaapi.get_imagebase(), idaapi.get_input_file_path())
    else:
        ea = idc.get_screen_ea()
        mod = idaapi.modinfo_t()
        while mod:
            if mod.base <= ea < mod.base + mod.size:
                return (mod.base, mod.name)
            if not idaapi.get_next_module(mod):
                break
        return idaapi.BADADDR


def goto_rva():
    rva = get_number_from_highlight()
    rva = ida_kernwin.ask_addr(rva, "Enter the RVA to jump to")
    if rva is None:
        return 0

    base, _ = get_screen_module()
    ea = base + rva
    plg_print(f"Base = 0x{base:X}, RVA = 0x{rva:X} => EA = 0x{ea:X}")
    idc.jumpto(ea)
    return 1


def goto_file_ofs():
    fofs = get_number_from_highlight()
    fofs = ida_kernwin.ask_addr(fofs, "Enter the file offset to jump to")
    if fofs is None:
        return 0

    ea = idaapi.get_fileregion_ea(fofs)
    if ea != idc.BADADDR:
        plg_print(f"File offset = 0x{fofs:X} -> EA = 0x{ea:X}")
        idc.jumpto(ea)
        return 1
    else:
        plg_print(f"Could not goto file offset 0x{fofs:X}")
        return 0


def copy_rva():
    ea = idc.get_screen_ea()
    base, _ = get_screen_module()
    rva = ea - base
    plg_print(f"EA = 0x{ea:X}, Base = 0x{base:X} => RVA = 0x{rva:X} copied to clipboard")
    copy_to_clip(f"0x{rva:X}")


def copy_file_offset():
    ea = idc.get_screen_ea()
    fofs = ida_loader.get_fileregion_offset(ea)
    if fofs == -1:
        plg_print("File offset unknown")
        return 0

    plg_print(f"EA = 0x{ea:X} -> file offset = 0x{fofs:X} copied to clipboard")
    copy_to_clip(f"0x{fofs:X}")
    return 1


# Org code of William Ballethin (FireEye) - hints-call plugin
# Thanks Willi :)
def get_ea_from_highlight():
    view = idaapi.get_current_viewer()
    thing = ida_kernwin.get_highlight(view)
    if thing and thing[1]:
        # we have a highligh, is it a valid name ?
        ea = idc.get_name_ea_simple(thing[0])
        if ea != idaapi.BADADDR:
            return ea

        # get name at screen ea
        ea = idc.get_screen_ea()
        name = idc.get_name(ea, idaapi.GN_DEMANGLED)
        if name and thing[0] in name:
            return ea

        # are we at end of function ?
        fn = idaapi.get_func(ea)
        if fn:
            if ea == fn.end_ea or idc.next_head(ea) == fn.end_ea:
                name = idc.get_name(fn.start_ea, idaapi.GN_DEMANGLED)
                if name and thing[0] in name:
                    return fn.start_ea

        # Try to get full highlight name
        place = idaapi.get_custom_viewer_place(view, False)
        if place and len(place) == 3:   # (plate_t, x, y)
            ea = place[0].toea()
            far_code_refs = [xref.to for xref in idautils.XrefsFrom(ea, ida_xref.XREF_FAR)]
            if far_code_refs:
                return far_code_refs[0]     # First xref

    # Reach now, we do not have any valid name, return current screen ea
    return idc.get_screen_ea()


def copy_highlight_name():
    ea = get_ea_from_highlight()
    if ea != idaapi.BADADDR:
        name = idc.get_name(ea)
        if not name:
            name = f"0x{ea:X}"  # copy ea
        copy_to_clip(name)
        plg_print(f"'{name}' copied to clipboard")
        return True
    else:
        plg_print("Invalid ea to copy")
        return False


def paste_highlight_name():
    ea = get_ea_from_highlight()
    if ea != idaapi.BADADDR:
        name = clip_text()
        if name:
            if not ida_name.force_name(ea, name):
                name = ida_name.validate_name(name, ida_name.VNT_IDENT)
                if not ida_name.force_name(ea, name):
                    plg_print(f"FAILED to set name '{name}' to 0x{ea:X}")
                    return False

            plg_print(f"Set name '{name}' to 0x{ea:X}")
            return True
        else:
            plg_print("Clipboard is empty")
    else:
        plg_print("Invalid ea to paste")

    return False


def get_selected_text():
    """ Get the highlight text. If none, force IDA copy text and we will get from clipboard """
    text = ""
    old_text = clip_text()

    view = idaapi.get_current_viewer()
    if view:
        thing = ida_kernwin.get_highlight(view)
        if thing and thing[1]:
            text = thing[0]

    # We not have a highlight text
    if not text:
        for action in idaapi.get_registered_actions():
            if "Copy" in action:
                shortcut = idaapi.get_action_shortcut(action)
                state = idaapi.get_action_state(action)
                if ("Ctrl-C" in shortcut) and (state and state[0] and (state[1] <= idaapi.AST_ENABLE)):
                    idaapi.process_ui_action(action)
                    text = clip_text()
                    if text != old_text:
                        break

    if not text:
        plg_print("Could not get any highlight/auto copied text\n"
                  f"Search with old clipboard text: '{old_text}'")
        text = old_text

    return text


def search_web(idx):
    urls = ["https://www.google.com/search?q=%s",
            "https://docs.microsoft.com/en-us/search/?terms=%s",
            "https://www.bing.com/search?q=%s",
            "https://github.com/search?q=%s&type=Code"]
    assert idx in range(len(urls))

    text = get_selected_text()
    if text:
        copy_to_clip(text)
        idaapi.open_url(urls[idx] % quote(text))


def dump_data_to_file(fName, data):
    defPath = os.path.dirname(idaapi.get_input_file_path())
    defPath = os.path.join(defPath, fName)
    dumpPath = idaapi.ask_file(1, defPath, "*.dump")
    if dumpPath:
        try:
            with open(dumpPath, "wb") as f:
                f.write(data)
            plg_print(f"Dump {len(data)} bytes to file {dumpPath} successed")
        except IOError as e:
            plg_print(str(e))


def process_data_result(start, data):
    # 16 bytes on a line
    # one byte take 4 char: 2 hex char, a space and a char if isalnum
    # one line take 3 char addtion: two space and \n, and ea hex address

    BYTES_PER_LINE = 16
    MAX_BYTES_HEX_DUMP = BYTES_PER_LINE * 64    # 64 lines

    printLen = len(data)
    if printLen > MAX_BYTES_HEX_DUMP:
        printLen = MAX_BYTES_HEX_DUMP
        plg_print(f"Only hexdump first {MAX_BYTES_HEX_DUMP} bytes")

    nLines = printLen // BYTES_PER_LINE     # Number of lines
    nOdd = printLen % BYTES_PER_LINE        # Number of bytes at last line

    isStr = True
    sHex = str()
    for i in range(printLen):
        # Accept NULL char in string
        if isStr and (chr(data[i]) not in string.printable) and (data[i] != 0):
            isStr = False

        if i % BYTES_PER_LINE == 0:
            sHex += f"{idaapi.ea2str(start + i)}: "

        sHex += f"{data[i]:02X} "

        if (i % BYTES_PER_LINE == BYTES_PER_LINE - 1) or (i == printLen - 1):
            # add the end of data or end of a line
            if nLines:
                lineIdx = i // BYTES_PER_LINE   # current line number
                low = lineIdx * BYTES_PER_LINE
                high = i + 1
            else:
                low = 0
                high = printLen

            sHex += " "

            # Padding last line
            if i == printLen - 1 and nLines and nOdd:
                sHex += " " * (BYTES_PER_LINE - nOdd) * 3

            for j in range(low, high):
                ch = chr(data[j])
                sHex += ch if ch.isalnum() else "."

            sHex += "\n"

    # Print out the hexdump string
    print(sHex)

    if isStr:
        txt = str(data).rstrip(chr(0))  # remove NULL chars at end txt
        print(f"String result: '{txt}'")
        idaapi.set_cmt(start, f"'{txt}'", 1)

    ret = idaapi.ask_yn(idaapi.ASKBTN_NO, "AUTOHIDE SESSION\nDo you want to patch selected range with result data ?")
    if ret != idaapi.ASKBTN_CANCEL:
        if ret == idaapi.ASKBTN_YES:
            idaapi.patch_bytes(start, bytes(data))
        ret = idaapi.ask_yn(idaapi.ASKBTN_NO, "AUTOHIDE SESSION\nDo you want to dump result data to file ?")
        if ret == idaapi.ASKBTN_YES:
            dump_data_to_file(f"{idaapi.get_root_filename()}_Dump_At_0x{start:X}_Size_{len(data)}.dump", data)


def base64_decode(std):
    addr = idc.BADADDR
    ea = idc.get_screen_ea()
    flags = idaapi.get_flags(ea)

    if idc.is_strlit(flags):
        addr = ea   # cursor is on the string
    elif idc.is_code(flags):
        addr = idc.get_first_dref_from(ea)  # get data reference from the instruction

    if addr == idc.BADADDR:
        plg_print("No string or reference to the string found\n")
        return

    b64str_enc = idc.get_strlit_contents(addr, -1, idc.get_str_type(addr))
    if not b64str_enc:
        plg_print(f"Could not get string at address 0x{addr:X}")
        return

    try:
        b64str_dec = base64.standard_b64decode(b64str_enc) if std else base64.urlsafe_b64decode(b64str_enc)
    except ValueError as e:
        plg_print(f"Could not decode.\n{str(e)}")
        return

    if b64str_dec:
        plg_print(f"Base64 decode of string '{b64str_enc}':")
        process_data_result(ea, bytearray(b64str_dec))


def str_to_bytes(sInput):
    """ str -> bytearray """
    try:
        s = sInput.strip()  # remove trailing white spaces
        if s.startswith('"') or s.startswith("'"):
            s = s[1:]
        if s.endswith('"') or s.endswith("'"):
            s = s[:-1]
        return bytearray(s.encode("utf-8"))
    except ValueError:
        return None


def hex_to_bytes(sInput):
    """ hex str to bytearray """
    s = sInput.lower().replace('0x', '').replace('\\x', '')
    s = ''.join("0" + c if len(c) % 2 else c for c in s.split())     # remove all white spaces
    try:
        s = bytes.fromhex(s)
    except ValueError:
        plg_print(f"Invalid hex string input '{sInput}'")
        s = None
    return s


def is_str(s):
    return s.startswith(('"', "'")) and s.endswith(('"', "'"))


def xor_data(data, key):
    """
    data: bytes
    key: bytes

    return: bytes
    """
    output = bytearray(len(data))
    for i, b in enumerate(data):
        output[i] = b ^ key[i % len(key)]
    return output


def nop_hider():
    hides = []
    in_nop_sled = False
    curr_pos = 0
    sled_len = 0

    for fn_ea in idautils.Functions():
        pfn = idaapi.get_func(fn_ea)
        if not pfn:
            continue
        for ea in range(pfn.start_ea, pfn.end_ea):
            b = idaapi.get_byte(ea)
            if b in (0x90, 0xCC):
                sled_len += 1
                if not in_nop_sled:
                    in_nop_sled = True
                    curr_pos = ea
            else:
                if in_nop_sled:
                    in_nop_sled = False
                    hides.append([curr_pos, sled_len])
                    curr_pos = 0
                    sled_len = 0

    # at end of function
    if in_nop_sled and sled_len > 1:
        hides.append([curr_pos, sled_len])

    if len(hides) == 0:
        plg_print("Found nothing NOPs block")

    for h in hides:
        if h[1] > 1:
            plg_print(f"Hide range: 0x{h[0]:X} - 0x{h[0] + h[1] -1:X}")
            idaapi.del_hidden_range(h[0] + h[1] - 1)
            idc.add_hidden_range(h[0], h[0] + h[1], '[NOPs]', '', '', idc.DEFCOLOR)
            idc.update_hidden_range(h[0], False)

    plg_print(f"Hidding {len(hides)} NOPs block")


def turn_off_ida_decision():
    sel, start, end = lazy_read_selection()
    if not sel:
        return 0

    plg_print(f"Turn off IDA auto analysis for range 0x{start:X} - 0x{end - 1:X}")
    idaapi.revert_ida_decisions(start, end)
    return 1


def lazy_get_bytes(ea, size):
    if idaapi.is_debugger_on():
        return idaapi.dbg_read_memory(ea, size)
    else:
        return idaapi.get_bytes(ea, size)


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_HOTKEY_COPY_EA[0]:
            ea = idc.get_screen_ea()
            copy_to_clip(f"0x{ea:X}")
            plg_print(f"Address '0x{ea:X}' copied to clipboard")
        elif self.action == ACTION_HOTKEY_COPY_RVA[0]:
            copy_rva()
        elif self.action == ACTION_HOTKEY_COPY_FOFS[0]:
            copy_file_offset()
        elif self.action == ACTION_HOTKEY_COPY_NAME[0]:
            copy_highlight_name()
        elif self.action == ACTION_HOTKEY_PASTE_NAME[0]:
            paste_highlight_name()
        elif self.action == ACTION_HOTKEY_GOTO_CLIP[0]:
            goto_clip_text()
        elif self.action == ACTION_HOTKEY_GOTO_FOFS[0]:
            goto_file_ofs()
        elif self.action == ACTION_HOTKEY_GOTO_RVA[0]:
            goto_rva()
        elif self.action == ACTION_HOTKEY_SEARCH_GOOGLE[0]:
            search_web(0)
        elif self.action == ACTION_HOTKEY_SEARCH_MSDOC[0]:
            search_web(1)
        elif self.action == ACTION_HOTKEY_SEARCH_BING[0]:
            search_web(2)
        elif self.action == ACTION_HOTKEY_SEARCH_GITHUB[0]:
            search_web(3)
        elif self.action == ACTION_HOTKEY_OPEN_URL[0]:
            text = get_selected_text()
            if text:
                idaapi.open_url(text)
        else:
            return 0

        return 1

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


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
        self.last_hex = "AA BB CC DD"

    def activate(self, ctx):
        if self.action in ACTION_MENU_CONVERT:
            sel, start, end = lazy_read_selection()
            if not sel:
                plg_print("Nothing to convert.")
                return 0

            size = end - start
            data = lazy_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:
                output = None
                plg_print(f"Dump from 0x{start:X} to 0x{end - 1:X} ({size} bytes):")
                if self.action == ACTION_MENU_CONVERT[0]:
                    # escaped string
                    output = "".join(f"\\x{b:02X}" for b in data)

                elif self.action == ACTION_MENU_CONVERT[1]:
                    # hex string, space
                    output = " ".join(f"{b:02X}" for b in data)

                elif self.action == ACTION_MENU_CONVERT[2]:
                    # C array
                    output = f"unsigned char {name}[{size}] = {{"
                    for i in range(size):
                        if i % 16 == 0:
                            output += "\n    "
                        output += f"0x{data[i]:02X}, "
                    output = output[:-2] + "\n};"

                elif self.action == ACTION_MENU_CONVERT[3]:
                    # C array word
                    data += b"\x00"
                    array_size = (size + 1) // 2
                    output = f"unsigned short {name}[{array_size}] = {{"
                    for i in range(0, size, 2):
                        if i % 16 == 0:
                            output += "\n    "
                        output += f"0x{u16(data[i:i+2]):04X}, "
                    output = output[:-2] + "\n};"

                elif self.action == ACTION_MENU_CONVERT[4]:
                    # C array dword
                    data += b"\x00" * 3
                    array_size = (size + 3) // 4
                    output = f"unsigned int {name}[{array_size}] = {{"
                    for i in range(0, size, 4):
                        if i % 32 == 0:
                            output += "\n    "
                        output += f"0x{u32(data[i:i+4]):08X}, "
                    output = output[:-2] + "\n};"

                elif self.action == ACTION_MENU_CONVERT[5]:
                    # C array qword
                    data += b"\x00" * 7
                    array_size = (size + 7) // 8
                    output = f"unsigned __int64 {name}[{array_size}] = {{"
                    for i in range(0, size, 8):
                        if i % 32 == 0:
                            output += "\n    "
                        output += f"0x{u64(data[i:i+8]):016X}, "
                    output = output[:-2] + "\n};"

                elif self.action == ACTION_MENU_CONVERT[6]:
                    # python list
                    output = f"{name} = [{', '.join(f'0x{b:02X}' for b in data)}]"

                elif self.action == ACTION_MENU_CONVERT[7]:
                    # python list word
                    data += b"\x00"
                    output = f"{name} = [{', '.join(f'0x{u16(data[i:i+2]):04X}' for i in range(0, size, 2))}]"

                elif self.action == ACTION_MENU_CONVERT[8]:
                    # python list dword
                    data += b"\x00" * 3
                    output = f"{name} = [{', '.join(f'0x{u32(data[i:i+4]):08X}' for i in range(0, size, 4))}]"

                elif self.action == ACTION_MENU_CONVERT[9]:
                    # python list qword
                    data += b"\x00" * 7
                    output = f"{name} = [{', '.join(f'{u64(data[i:i+8]):016X}' for i in range(0, size, 8))}]"

                elif self.action == ACTION_MENU_CONVERT[10]:
                    # MASM byte array
                    header = f"{name} db "
                    output = header
                    for i in range(size):
                        if i and i % 16 == 0:
                            output += "\n"
                            output += " " * len(header)
                        output += f"0{data[i]:02X}h, "
                    output = output[:-2]

                elif self.action == ACTION_MENU_CONVERT[11]:
                    # GNU ASM byte array
                    header = f"{name}: .byte "
                    output = header
                    for i in range(size):
                        if i and i % 16 == 0:
                            output += "\n"
                            output += " " * len(header)
                        output += f"0x{data[i]:02X}, "
                    output = output[:-2]

                if output:
                    print(output)
                    copy_to_clip(output)
                    output = None

        elif self.action == ACTION_MENU_COPY_DATA:
            sel, start, end = lazy_read_selection()
            if not sel:
                return 0

            data = lazy_get_bytes(start, end - start)
            if isinstance(data, str):
                data = bytearray(data)
            output = "".join(f"{b:02X}" for b in data)
            copy_to_clip(output)
            plg_print(f"Hex string '{output}' copied")

        elif self.action == ACTION_MENU_COPY_STR:
            ea = get_ea_from_highlight()
            if idc.is_strlit(idc.get_full_flags(ea)):
                s = str(idc.get_strlit_contents(ea, -1, idc.get_str_type(ea)))
                if s.startswith("b'") and s.endswith("'"):  # byte string
                    s = s[2:-1]
                copy_to_clip(s)
                plg_print(f"'{s}' copied")
            else:
                plg_print("Current EA not in a string")

        elif self.action == ACTION_MENU_DUMP_DATA:
            sel, start, end = lazy_read_selection()
            if not sel:
                return 0

            size = end - start
            data = lazy_get_bytes(start, size)
            if data:
                if len(data) < size:
                    plg_print("Request {size} bytes, only get {len(data)} bytes")
                    size = len(data)
                dump_data_to_file(f"{idaapi.get_root_filename()}_Dump_At_0x{start:X}_Size_{size}.dump", data)
            else:
                plg_print(f"0x{start:X}: unable to get {size} bytes")

        elif self.action == ACTION_MENU_DUMP_SEG:
            ea = idc.here()
            seg = idaapi.getseg(ea)
            if not seg:
                plg_print(f"0x{ea:X} Unable to get segment at current ea")
                return 0

            size = seg.end_ea - seg.start_ea
            data = lazy_get_bytes(seg.start_ea, size)
            if data:
                if len(data) < size:
                    plg_print("Request {size} bytes, only get {len(data)} bytes")
                    size = len(data)
                dump_data_to_file(f"{idaapi.get_root_filename()}_Dump_Segment_{idaapi.get_segm_name(seg).lstrip('.')}_Size_{size}.dump", data)
            else:
                plg_print(f"0x{seg.start_ea:X}: unable to get {size} bytes")

        elif self.action == ACTION_MENU_XOR_DATA:
            sel, start, end = lazy_read_selection()
            if not sel:
                return 0

            size = end - start

            key = idaapi.ask_str(self.last_hex, 0, "Xor with hex values (or a string begin and end with\" or ')...")
            if not key:
                return 0

            bytes_key = bytearray()
            if is_str(key):
                bytes_key = str_to_bytes(key)
            else:
                bytes_key = hex_to_bytes(key)

            if not bytes_key:
                return 0

            self.last_hex = key     # store for later asking

            data = lazy_get_bytes(start, end - start)
            if isinstance(data, str):  # python2 compatibility
                data = bytearray(data)

            output = xor_data(data, bytes_key)
            if not output:
                plg_print("Sorry, error occurred. My bug :( Please report.")
                return 0

            assert size == len(output)

            plg_print(f"Xor result from 0x{start:X} to 0x{end - 1:X} ({end - start} bytes) with {key}:")
            process_data_result(start, output)

        elif self.action == ACTION_MENU_FILL_NOP:
            sel, start, end = lazy_read_selection()
            if not sel:
                return 0

            idaapi.patch_bytes(start, b"\x90" * (end - start))
            idc.create_insn(start)
            plg_print(f"Fill 0x{start:X} to 0x{end - 1:X} ({end - start} bytes) with NOPs")

        elif self.action == ACTION_MENU_NOP_HIDER:
            nop_hider()

        elif self.action == ACTION_MENU_AUTO_OFF:
            turn_off_ida_decision()

        elif self.action == ACTION_MENU_B64STD:
            base64_decode(True)

        elif self.action == ACTION_MENU_B64URL:
            base64_decode(False)

        elif self.action == ACTION_MENU_SCAN_VUL:
            plg_print("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:
                plg_print(f"Done! {len(found)} possible vulnerabilities found.")
                ch = VulnChoose("Vulnerability", found, None, False)
                ch.Show()
            else:
                plg_print("No format string vulnerabilities found.")

        else:
            return 0

        return 1

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

    @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 None

            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 LAZY_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 LAZY_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 LAZY_BITS == 64):
                    break
            elif name.endswith("printf"):
                if op in ("mov", "lea") and (dst.endswith(("rdi", "[esp]", "R0")) or
                                             dst.endswith("edi") and LAZY_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 in (idc.o_imm, 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 None

        plg_print(f"0x{addr:X}: Possible Vulnerability: {name}, format = {opnd}")
        return [f"0x{addr:X}", 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_REMOVE_RET_TYPE[0]:
            vdui = idaapi.get_widget_vdui(ctx.widget)
            self.remove_rettype(vdui)
            vdui.refresh_ctext()
        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, widget, popup):
        # attach Searchs menu to all widget
        idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_SEARCH_GOOGLE[0], PLUGIN_POPUP)
        idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_SEARCH_MSDOC[0], PLUGIN_POPUP)
        idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_SEARCH_BING[0], PLUGIN_POPUP)
        idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_SEARCH_GITHUB[0], PLUGIN_POPUP)
        idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_OPEN_URL[0], PLUGIN_POPUP)
        idaapi.attach_action_to_popup(widget, popup, None, PLUGIN_POPUP)

        widget_type = idaapi.get_widget_type(widget)
        if widget_type in [idaapi.BWN_DISASM, idaapi.BWN_DUMP]:
            for action in ACTION_MENU_CONVERT:
                idaapi.attach_action_to_popup(widget, popup, action, PLUGIN_POPUP + "Convert/")

            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_COPY_DATA, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_COPY_STR, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, None, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_DUMP_DATA, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_DUMP_SEG, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_XOR_DATA, PLUGIN_POPUP)

            if CAN_NOP:
                idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_FILL_NOP, PLUGIN_POPUP)
                idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_NOP_HIDER, PLUGIN_POPUP)

            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_AUTO_OFF, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, None, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_B64STD, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_B64URL, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, None, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_COPY_EA[0], PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_COPY_RVA[0], PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_COPY_FOFS[0], PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, None, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_COPY_NAME[0], PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_PASTE_NAME[0], PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, None, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_GOTO_CLIP[0], PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_GOTO_FOFS[0], PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_HOTKEY_GOTO_RVA[0], PLUGIN_POPUP)

        if widget_type == idaapi.BWN_DISASM and (LAZY_ARCH, LAZY_BITS) in [(idaapi.PLFM_386, 32),
                                                                           (idaapi.PLFM_386, 64),
                                                                           (idaapi.PLFM_ARM, 32), ]:
            idaapi.attach_action_to_popup(widget, popup, None, PLUGIN_POPUP)
            idaapi.attach_action_to_popup(widget, popup, ACTION_MENU_SCAN_VUL, PLUGIN_POPUP)


class HexRays_Hook:
    def callback(self, event, *args):
        if event == idaapi.hxe_populating_popup:
            if not ENABLE_REMOVE_RETTYPE:
                return 0

            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_REMOVE_RET_TYPE[0], PLUGIN_POPUP)
        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(f"{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   # | idaapi.PLUGIN_UNL for fix bugs, debug
    comment = PLUGIN_NAME
    help = ""
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""

    def __init__(self):
        self.ui_hook = None
        self.hx_hook = None
        self.hexrays_inited = False
        self.registered_hotkey_actions = []
        self.registered_menu_actions = []
        self.registered_hexray_actions = []

    def init(self):
        plg_print(f"v{PLUGIN_VERSION} - plugin has been loaded.")

        # Register hotkey actions
        for HK_ACT in ALL_HOTKEY_ACTIONS:
            action = idaapi.action_desc_t(HK_ACT[0],    # name
                                          HK_ACT[1],    # label
                                          hotkey_action_handler_t(HK_ACT[0]),   # action handler
                                          HK_ACT[2],    # shortcut
                                          HK_ACT[3],    # tooltip
                                          HK_ACT[4])    # iconid
            idaapi.register_action(action)
            self.registered_hotkey_actions.append(action.name)

        # Register menu actions
        menu_actions = (
            idaapi.action_desc_t(ACTION_MENU_CONVERT[0], "Convert to escaped hex string", menu_action_handler_t(ACTION_MENU_CONVERT[0]), None, None, 80),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[1], "Convert to space hex string", menu_action_handler_t(ACTION_MENU_CONVERT[1]), None, None, 8),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[2], "Convert to C/C++ array (BYTE)", menu_action_handler_t(ACTION_MENU_CONVERT[2]), None, None, 38),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[3], "Convert to C/C++ array (WORD)", menu_action_handler_t(ACTION_MENU_CONVERT[3]), None, None, 38),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[4], "Convert to C/C++ array (DWORD)", menu_action_handler_t(ACTION_MENU_CONVERT[4]), None, None, 38),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[5], "Convert to C/C++ array (QWORD)", menu_action_handler_t(ACTION_MENU_CONVERT[5]), None, None, 38),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[6], "Convert to python list (BYTE)", menu_action_handler_t(ACTION_MENU_CONVERT[6]), None, None, 201),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[7], "Convert to python list (WORD)", menu_action_handler_t(ACTION_MENU_CONVERT[7]), None, None, 201),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[8], "Convert to python list (DWORD)", menu_action_handler_t(ACTION_MENU_CONVERT[8]), None, None, 201),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[9], "Convert to python list (QWORD)", menu_action_handler_t(ACTION_MENU_CONVERT[9]), None, None, 201),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[10], "Convert to MASM array (BYTE)", menu_action_handler_t(ACTION_MENU_CONVERT[10]), None, None, 38),
            idaapi.action_desc_t(ACTION_MENU_CONVERT[11], "Convert to GNU ASM array (BYTE)", menu_action_handler_t(ACTION_MENU_CONVERT[11]), None, None, 38),
            idaapi.action_desc_t(ACTION_MENU_COPY_DATA, "Copy hex data", menu_action_handler_t(ACTION_MENU_COPY_DATA), None, None, 0x1F),
            idaapi.action_desc_t(ACTION_MENU_COPY_STR, "Copy string", menu_action_handler_t(ACTION_MENU_COPY_STR), None, None, 0x1F),
            idaapi.action_desc_t(ACTION_MENU_DUMP_DATA, "Dump selected data to file", menu_action_handler_t(ACTION_MENU_DUMP_DATA), None, None, 0x1B),
            idaapi.action_desc_t(ACTION_MENU_DUMP_SEG, "Dump current segment to file", menu_action_handler_t(ACTION_MENU_DUMP_SEG), None, None, 0x1B),
            idaapi.action_desc_t(ACTION_MENU_XOR_DATA, "Get xored data", menu_action_handler_t(ACTION_MENU_XOR_DATA), None, None, 9),
            idaapi.action_desc_t(ACTION_MENU_AUTO_OFF, "Revert IDA Decision", menu_action_handler_t(ACTION_MENU_AUTO_OFF), None, None, 9),
            idaapi.action_desc_t(ACTION_MENU_B64STD, "Base64Std decode", menu_action_handler_t(ACTION_MENU_B64STD), None, None, 9),
            idaapi.action_desc_t(ACTION_MENU_B64URL, "Base64Url decode", menu_action_handler_t(ACTION_MENU_B64URL), None, None, 9),
            idaapi.action_desc_t(ACTION_MENU_SCAN_VUL, "Scan format string vulnerabilities", menu_action_handler_t(ACTION_MENU_SCAN_VUL), None, None, 160),
        )

        if CAN_NOP:
            menu_actions += (idaapi.action_desc_t(ACTION_MENU_FILL_NOP, "Fill with NOPs", menu_action_handler_t(ACTION_MENU_FILL_NOP), None, None, 9),
                             idaapi.action_desc_t(ACTION_MENU_NOP_HIDER, "NOPs Hider", menu_action_handler_t(ACTION_MENU_NOP_HIDER), None, None, 9))

        for action in menu_actions:
            idaapi.register_action(action)
            self.registered_menu_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():
            hx_actions = (
                idaapi.action_desc_t(ACTION_HX_REMOVE_RET_TYPE[0],
                                     ACTION_HX_REMOVE_RET_TYPE[1],
                                     hexrays_action_handler_t(ACTION_HX_REMOVE_RET_TYPE[0]),
                                     ACTION_HX_REMOVE_RET_TYPE[2],
                                     ACTION_HX_REMOVE_RET_TYPE[3],
                                     -1),
            )
            for action in hx_actions:
                idaapi.register_action(action)
                self.registered_hexray_actions.append(action.name)

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

        addon = idaapi.addon_info_t()
        addon.id = "htc_lazyida"
        addon.name = PLUGIN_NAME
        addon.producer = "HTC (Original: Lays - tw.l4ys.lazyida)"
        addon.url = "https://github.com/HongThatCong/LazyIDA"
        addon.version = PLUGIN_VERSION
        idaapi.register_addon(addon)

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        if self.ui_hook:
            self.ui_hook.unhook()
            self.ui_hook = None

        # Unregister actions
        if self.registered_hotkey_actions:
            for action in self.registered_hotkey_actions:
                idaapi.unregister_action(action)
            del self.registered_hotkey_actions[:]

        if self.registered_menu_actions:
            for action in self.registered_menu_actions:
                idaapi.unregister_action(action)
            del self.registered_menu_actions[:]

        if self.hexrays_inited:
            if self.hx_hook:
                idaapi.remove_hexrays_callback(self.hx_hook.callback)
                self.hx_hook = None

            # Unregister hexrays actions
            if self.registered_hexray_actions:
                for action in self.registered_hexray_actions:
                    idaapi.unregister_action(action)
                del self.registered_hexray_actions[:]

            idaapi.term_hexrays_plugin()

        plg_print("plugin terminated")


def PLUGIN_ENTRY():
    return LazyIDA_t()


repository: miaubiz/type_copier
file: type_copier.py
description: ida pro type copying plugin
url: https://github.com/miaubiz/type_copier/blob/3cd1c87cbeb9c911922228cbba82d8f5cced7464/type_copier.py
wanted_name: type_copier.py
comment: copy types
created_at: 2015-04-20 06:59:33+00:00
pushed_at: 2015-04-20 07:06:39+00:00
content:
# coding: utf-8
from struct import unpack, pack
import pickle

from idaapi import PLUGIN_OK, PLUGIN_UNL, plugin_t, get_root_filename, BADADDR, get_func, open_pseudocode
from idc import GetFunctionName, GetLocalTypeName, GetMaxLocalType, GetTinfo, NextFunction
from idc import here, ApplyType, LocByName, Jump

"""
Copy unparseable types from one IDB to another.
It requires the unparseable type to be a Local Type in the destination IDB.
The Local Types will have different ordinals in the two IDBs.
"""


class TinfoReader:
	def __init__(self, tp):
		self.pos = 0
		self.tp = tp

	def read_byte(self):
		(result,) = unpack("B", self.tp[self.pos])
		self.pos += 1
		return result

	def keep_going(self):
		return self.pos < len(self.tp)


class CopyTypes(plugin_t):
	flags = PLUGIN_UNL
	comment = "copy types"
	help = "help sir"
	wanted_name = "type_copier.py"
	wanted_hotkey = "{"

	def __init__(self):
		pass

	def receive_message(self, filename):
		msg = None
		try:
			handle = open(filename)
			msg = pickle.loads(handle.read())
			handle.close()
		except IOError:
			pass
		return msg

	def send_message(self, filename, msg):
		handle = open(filename, "w")
		handle.write(pickle.dumps(msg))
		handle.close()

	def find_type_by_name(self, name):
		for i in range(1, GetMaxLocalType()):
			if name == GetLocalTypeName(i):
				return i
		print "didn't find: %s" % name
		raise RuntimeError

	def unparse_tinfo(self, tinfo):
		the_bytes = []
		for thing in tinfo:
			if type(thing) == int:  # if it's a byte, just put it back in
				the_bytes.append(thing)
			else:
				the_bytes.append(ord("="))  # a type starts with =
				ordinal = self.find_type_by_name(thing["local_type"])  # get the ordinal of the Local Type based on its name
				if ordinal < 0x40:
					the_bytes.append(0x3)  # length of "=" + length byte + single byte encoded ordinal
				else:
					the_bytes.append(0x4)  # length of "=" + length byte + two byte encoded ordinal
				the_bytes.append(ord("#"))  # the number symbol means ordinal follows
				if ordinal < 0x40:
					the_bytes.append(ordinal | 0x40)  # encode one byte ordinal
				else:
					the_bytes.append((ordinal >> 6) | 0x80)  # encode two byte ordinal
					the_bytes.append(ordinal & 0x3f | 0x40)
		packed = pack("%dB" % len(the_bytes), *the_bytes)
		return packed

	def parse_tinfo(self, ea):
		tinfo = GetTinfo(ea)
		if tinfo is None:
			return None, None, None
		type_, fields = tinfo

		tp = TinfoReader(type_)
		# print idc_print_type(type_, fields, "fun_name", 0)
		# print type_.encode("string_escape")
		output = []
		"""
		Attempt to copy the tinfo from a location, replacing any Local Types with our own representation of them.
		Pass all other bytes through as-is.
		"""
		while tp.keep_going():
			a_byte = tp.read_byte()
			unwritten_bytes = [a_byte]
			if a_byte == ord("="):  # a type begins
				ordinal_length = tp.read_byte()
				unwritten_bytes.append(ordinal_length)
				if ordinal_length in [0x3, 0x4]:  # if you have more than 0x2000 Local Types you might need to handle the case of 0x5
					number_marker = tp.read_byte()
					unwritten_bytes.append(number_marker)
					if number_marker == ord("#"):  # this is a Local Type referred to by its ordinal
						ordinal = 0
						for w in range(ordinal_length - 2):
							next_byte = tp.read_byte()
							ordinal = ordinal * 0x40 + (next_byte & 0x7f)  # decode ordinal number
						ordinal -= 0x40

						t = GetLocalTypeName(ordinal)
						output.append({"local_type": t})
						continue

			output += unwritten_bytes  # put all the bytes we didn't consume into the output as-is

		return type_, output, fields

	def type_one_function(self, tinfo_map):
		fn_name = GetFunctionName(here())
		if not fn_name:
			print "you're not in a function lol"
			return
		if fn_name not in tinfo_map:
			print "I don't know of this %s" % fn_name
			return
		tinfo, fields = tinfo_map[fn_name]
		fn_loc = LocByName(fn_name)
		current_tinfo_ = GetTinfo(fn_loc)
		packed = self.unparse_tinfo(tinfo)
		if current_tinfo_ is not None:
			current_tinfo, current_fields = current_tinfo_
			if packed == current_tinfo and fields == current_fields:
				print "already the same"
				return
		ret = ApplyType(fn_loc, (packed, fields))
		if ret:
			print "success: %#x %s" % (fn_loc, fn_name)
			open_pseudocode(fn_loc, False)
		else:
			print "failed :("
			print fn_name
			print packed.encode("string_escape")
			print fields.encode("string_escape")

	def dump_tinfos(self, tmp_file):
		tinfo_map = {}
		ea = 0
		num_tinfos = 0
		while True:
			ea = NextFunction(ea)
			if ea == BADADDR:
				break
			fun = get_func(ea)
			fn_name = GetFunctionName(ea)

			tinfo, parsed_tinfo, fields = self.parse_tinfo(fun.startEA)
			if tinfo is None:
				continue
			reconstituted = self.unparse_tinfo(parsed_tinfo)
			if tinfo == reconstituted:
				tinfo_map[fn_name] = (parsed_tinfo, fields)
				num_tinfos += 1
			else:
				print "fuck"
				ea = fun.startEA
				Jump(ea)
				print tinfo.encode("string_escape")
				print reconstituted.encode("string_escape")
				break
		self.send_message(tmp_file, tinfo_map)
		print "dumped %d tinfos" % num_tinfos

	def run(self, arg_=None):
		tmp_file = "/tmp/tinfo_storage"  # windows, do something lol
		if len(get_root_filename()) == 15:  # dest -> 15, source !15
			tinfo_map = self.receive_message(tmp_file)
			self.type_one_function(tinfo_map)
		else:
			self.dump_tinfos(tmp_file)
		print "done"

	def init(self):
		return PLUGIN_OK

	def term(self):
		pass


def PLUGIN_ENTRY():
	return CopyTypes()


repository: shemesh999/oregami
file: oregami_plugin.py
description: IDA plugins and scripts for analyzing register usage frame
wanted_name: OregamiWindow
comment: oREGami
created_at: 2020-03-16 11:25:54+00:00
pushed_at: 2023-06-08 14:27:51+00:00
url:
https://github.com/shemesh999/oregami/blob/eee466005af08c816275e3269a83704c653a5312/oregami_plugin.py

content:
from oregami.reg_frame import *
from oregami.reg_utils import *


class Oregami(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "oREGami"
    help = "Find references to register in function"
    wanted_name = "OregamiWindow"
    wanted_hotkey = "Shift+X"

    @staticmethod
    def init():
        return idaapi.PLUGIN_OK

    @staticmethod
    def term():
        pass

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

        oregami_plugin_starter(start)


def PLUGIN_ENTRY():
    return Oregami()


def oregami_plugin_starter(orig_ea):
    reg = get_reg_from_user(orig_ea)

    if reg is None:
        return

    rf = RegFrame(orig_ea, reg, force=(not conf.cache_bool))

    rx = RegXref(reg=reg, ea=orig_ea)
    eas = []

    for insn in rf.get_instructions():
        rx.add_xref(insn.ea, insn.is_read, insn.is_write, insn.is_implicit)
        eas += [insn.ea]

    if orig_ea in eas:
        idx = eas.index(orig_ea)
        rx.deflt = idx

    item_idx = rx.Show(True)
    if item_idx >= 0:
        ida_kernwin.jumpto(eas[item_idx])


class RegXref(idaapi.Choose):
    def __init__(self, reg, ea=None, icon=55, flags=0, width=None,
                 height=None, embedded=False, modal=False):
        title = 'xrefs to {}'.format(reg)
        idaapi.Choose.__init__(
            self,
            title,
            [["Direction", 5], ["Type", 7], ["Address", 7], ["Text", 40]],
            width=width,
            height=height,
            embedded=embedded,
            flags=flags)
        self.n = 0
        self.items = []
        self.icon = icon
        self.selcount = 0
        self.modal = modal
        self.popup_names = []
        self.ea = ea
        self.title = title
        # 80 - str
        # 52 - type
        # 54 - reg?

    def add_xref(self, ea, is_read, is_write, is_implicit):
        if ea < self.ea:
            direction = '^'
        elif ea > self.ea:
            direction = 'v'
        else:
            direction = '.'

        addr = '0x%x' % ea

        type_str = ''
        if is_read:
            type_str += 'r'
        if is_write:
            type_str += 'w'
        if (not is_read) and (not is_write):
            type_str += '{unknown}'

        if is_implicit:
            type_str += ' {implicit}'

        text = sark.Line(ea).disasm
        self.items.append([direction, type_str, addr, text])

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

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

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

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

    # In old IDA does require other funcs be implemented
    def OnClose(self):
        return


repository: stuxnet147/IDA-Assistant
file: IDA_Assistant.py
description: IDA plugin to support automatic reverse engineering
wanted_name: IDA Assistant
comment: IDA Assistant powered by Anthropic
created_at: 2024-03-13 09:57:13+00:00
pushed_at: 2025-02-28 11:50:55+00:00
url:
https://github.com/stuxnet147/IDA-Assistant/blob/93f550bf9ab178fa8a429214ca68be55625f7802/IDA_Assistant.py

content:
import threading
import json
import ida_idaapi
import ida_kernwin
import ida_hexrays
import ida_name
import ida_funcs
import ida_bytes
import ida_ida
import idc
import idautils
import idaapi
import anthropic
import re
import traceback
import functools
from PyQt5 import QtWidgets, QtCore

class IDAAssistant(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_FIX
    comment = "IDA Assistant powered by Anthropic"
    help = "Provides an AI assistant for reverse engineering tasks"
    wanted_name = "IDA Assistant"
    wanted_hotkey = "Alt-F1"

    def __init__(self):
        super(IDAAssistant, self).__init__()
        self.model = "claude-3-sonnet-20240229"
        self.client = anthropic.Anthropic(
            api_key="YOUR API KEY"
        )
        self.chat_history = []
        self.message_history = []
        
    def init(self):
        return ida_idaapi.PLUGIN_OK

    def run(self, arg):
        self.assistant_window = AssistantWidget()
        self.assistant_window.Show("IDA Assistant")

    def term(self):
        pass

    def add_assistant_message(self, message):
        self.chat_history.append(f"<b>Assistant:</b> {message}") 
        
    def query_model(self, role, query, cb, additional_model_options=None):
        if additional_model_options is None:
            additional_model_options = {}
        
        system_prompt = """
        - You are IDA-Assistant, an AI designed to assist users in reverse engineering and binary analysis tasks using IDA Pro.
        - Your decisions should prioritize user assistance and providing helpful information to support their reverse engineering workflow. Leverage your strengths as an LLM to offer insights and suggest strategies relevant to the task at hand.
        - Your goal is to provide helpful guidance and insights to users performing reverse engineering in IDA Pro.
        - Optimize your responses to be concise yet informative.
        - User assistance is the top priority. Always strive to provide helpful information to the user.
        - Do not dynamically use the result of one command as an argument for another command.
        - Each command should be independently executable, and arguments should be explicitly provided.
        - You can list multiple commands sequentially, but they should not be structured to depend on the results of the previous commands.
        - When you use the command, you cannot use a value that you have not yet obtained as a parameter.
        - Avoid repeating the same commands if the necessary information has already been obtained.
        - Limit your actions to the commands listed below.
            Commands:
            - Name: get_disassembly
                - Description: Gets the disassembly from start address to end address.
                - Args: "start_address": String, "end_address": String
            - Name: decompile_address
                - Description: Decompile the function at the specified address. 
                - Args: "address": String
            - Name: rename_address
                - Description: Rename the address at the specified address.
                - Args: "address": String, "new_name": String, "old_name": String
            - Name: get_function_start_end_address
                - Description: Get the start and end address of the function at the specified address.
                - Args: "address": String
            - Name: get_addresses_of_name
                - Description: Search for a name as a parameter in ida name list and get all addresses with that name in the form of a list.
                - Args: "name": String
            - Name: get_xrefs_to
                - Description: Get the cross-references to the specified address.
                - Args: "address": String
            - Name: get_xrefs_from
                - Description: Get the cross-references from the specified address.
                - Args: "address": String
            - Name: get_func_xrefs_to
                - Description: Get the details of all cross-references to the specified function.
                - Args: "address": String
            - Name: do_nothing
                - Description: Do nothing. Use it when a series of tasks are completed.
                - Args: None: No arguments. but it should be included in the json like {"args": {}}
            - Name: set_comment
                - Description: Set a comment at the specified address.
                - Args: "address": String, "comment": String
            - Name: get_address_type
                - Description: Get the type of the address.
                - Args: "address": String
        - Example of wrong usage:
            "command": [
                {
                    "name": "get_address_of_name",
                    "args": {"name": "dispatch::handler"}
                },
                {
                    "name": "decompile_address",
                    "args": {"address": "<result of previous command>"} # Wrong usage
                },
                {
                    "name": "get_xrefs_to",
                    "args": {"address": "<result of 2 commands ago>"} # Wrong usage
                },
                {
                    "name": "get_xrefs_from",
                    "args": {"address": "<result of 2 commands ago>"} # Wrong usage
                }
            ]
        - Do not attempt to use the previous command run results as parameters dynamically, as in the wrong usage example.
        - If the parameter is "address", you must explicitly write down the address string.
        - Try not to repeat the same mistake.
        - You can access to loaded binary and IDA API for analysis.
        - You can see user's current view/position in IDA Pro.
        - You have knowledge base on reverse engineering concepts and common techniques.
        - You can use multiple commands.
        - Reflect on how well your suggestions assisted the user in their reverse engineering task.
        - Assess whether the user found your insights helpful and relevant.
        - Consider potential alternative approaches that could have been more efficient or impactful.
        - Strive to provide the most value to the user with each interaction.
        - You should only respond in JSON format as described below:
            {
                "thoughts": {
                    "text": "thought",
                    "reasoning": "reasoning",
                    "criticism": "constructive self-criticism", 
                    "speak": "thoughts summary to say to user"
                },
                "command": [
                    {
                        "name": "command name",
                        "args": {"arg name": value}
                    }
                ]
            }
        - Ensure the response can be parsed by Python json.loads. 
        - Always strictly adhere to the specified JSON response format, and do not deviate from it under any circumstances.
        - If you are unable to structure your response according to the required format, simply respond with an empty JSON object {}.
        - Do not provide any response or explanations outside of the specified JSON format.
        """
        
        messages = self.message_history.copy()
        messages.append({"role": role, "content": query})

        try:
            response = self.client.messages.create(
                model=self.model,
                temperature=0.0,
                messages=messages,
                system=system_prompt,
                **additional_model_options
            )
            
            assistant_reply = response.content[0].text.strip().replace("```json\n", "").replace("```\n", "").strip()
            print(assistant_reply)

            self.message_history.append({"role": role, "content": query})
            self.message_history.append({"role": "assistant", "content": assistant_reply})            
            self.chat_history.append(f"<b>User:</b> {query}")             
            ida_kernwin.execute_sync(functools.partial(cb, response=assistant_reply), ida_kernwin.MFF_WRITE)
        except Exception as e:
            print(str(e))
            traceback_details = traceback.format_exc()
            print(traceback_details)

    def query_model_async(self, role, query, cb, additional_model_options=None):
        if additional_model_options is None:
            additional_model_options = {}
        t = threading.Thread(target=self.query_model, args=[role, query, cb, additional_model_options])
        t.start()

class AssistantWidget(ida_kernwin.PluginForm):
    def OnCreate(self, form):
        self.parent = self.FormToPyQtWidget(form)
        self.PopulateForm()
        self.assistant = IDAAssistant()
        self.command_results = []
        self.functions = self.get_name_info()
        
    def PopulateForm(self):
        layout = QtWidgets.QVBoxLayout()
        
        self.view = ida_kernwin.get_current_viewer()
        self.output_window = ida_kernwin.find_widget("Output window")
        
        self.chat_history = QtWidgets.QTextEdit()
        self.chat_history.setReadOnly(True)
        layout.addWidget(self.chat_history)
        
        input_layout = QtWidgets.QHBoxLayout()
        self.user_input = QtWidgets.QLineEdit()
        input_layout.addWidget(self.user_input)
        
        send_button = QtWidgets.QPushButton("Send")
        send_button.clicked.connect(self.OnSendClicked)
        input_layout.addWidget(send_button)
        
        layout.addLayout(input_layout)
        
        self.parent.setLayout(layout)
        
    def OnSendClicked(self):
        user_message = self.user_input.text().strip()
        if user_message:
            self.chat_history.append(f"<b>User:</b> {user_message}")
            self.user_input.clear()
            
            current_address = idc.here()
            
            prompt = f"{user_message}\nCurrent address: {hex(current_address)}"
            
            self.assistant.query_model_async("user", prompt, self.OnResponseReceived, additional_model_options={"max_tokens": 1000})
    
    def OnResponseReceived(self, response):
        try:
            assistant_reply = self.ParseResponse(response)

            if assistant_reply is None:
                self.chat_history.append(f"<b>System Message:</b> Failed to parse assistant response.")
                return

            if not assistant_reply:
                self.chat_history.append(f"<b>System Message:</b> No response from assistant.")
                return

            self.chat_history.append(f"<b>Assistant speak:</b> {assistant_reply['thoughts']['speak']}")

            commands = assistant_reply['command']
            command_results = {}

            for command in commands:
                command_name = command['name']
                if command_name == "do_nothing":
                    continue

                command_args = command['args']

                command_handler = getattr(self, f"handle_{command_name}", None)
                if command_handler:
                    command_results[command_name] = command_handler(command_args)
                else:
                    self.PrintOutput(f"Unknown command: {command_name}")
                    command_results[command_name] = None

            query = ""
            for command_name, result in command_results.items():
                if result is not None:
                    query += f"{command_name} result:\n{json.dumps(result)}\n\n"
                else:
                    query += f"{command_name} result: None\n\n"
                    
            if len(query) > 0:
                self.assistant.query_model_async("user", f"{query}", self.OnResponseReceived, additional_model_options={"max_tokens": 1000})

        except Exception as e:
            traceback_details = traceback.format_exc()
            print(traceback_details)
            self.PrintOutput(f"Error parsing assistant response: {str(e)}")
            self.assistant.query_model_async("user", f"Error parsing response. please retry:\n {str(e)}", self.OnResponseReceived, additional_model_options={"max_tokens": 1000})

    def handle_eval_idc(self, args):
        try:
            idc_expression = args["idc_expression"]
            result = idc.eval_idc(idc_expression)
            return result
        except Exception as e:
            return f"Error: {str(e)}"

    def handle_get_disassembly(self, args):
        try:
            start_address = int(args["start_address"], 16)
            end_address = int(args["end_address"], 16)

            disassembly = ""
            while start_address < end_address:
                disassembly += f"{hex(start_address)}: {idc.GetDisasm(start_address)}\n"
                start_address = idc.next_head(start_address)
            return disassembly
        except Exception as e:
            return f"Error: {str(e)}"

    def handle_get_disassembly_function(self, args):
        try:
            name = args["name"]
            address = idc.get_name_ea_simple(name)
            if address != idc.BADADDR:
                start_address = function.start_ea
                end_address = function.end_ea

                disassembly = ""
                while start_address < end_address:
                    disassembly += f"{hex(start_address)}: {idc.GetDisasm(start_address)}\n"
                    start_address = idc.next_head(start_address)
                return disassembly
            return f"No function found at address {name}"
        except Exception as e:
            return f"Error: {str(e)}"

    def handle_decompile_address(self, args):
        try:
            address = int(args["address"], 16)
            function = idaapi.get_func(address)
            if function:
                decompiled_code = idaapi.decompile(function)
                if decompiled_code:
                    return str(decompiled_code)
            return f"No function found at address {hex(address)}"
        except Exception as e:
            return f"Error: {str(e)}"

    def handle_decompile_function(self, args):
        try:
            name = args["name"]
            address = idc.get_name_ea_simple(name)
            if address != idc.BADADDR:
                function = idaapi.get_func(function.start_ea)
                if function:
                    decompiled_code = idaapi.decompile(function)
                    if decompiled_code:
                        return str(decompiled_code)
                else:
                    self.PrintOutput(f"No function found at address {name}")
            return None
        except Exception as e:
            return f"Error: {str(e)}"

    def handle_rename_address(self, args):
        try:
            address = int(args["address"], 16)
            new_name = args["new_name"]
            old_name = args["old_name"]
            if new_name and old_name:
                ida_hexrays.rename_lvar(address, old_name, new_name)
                result = f"Renamed address {hex(address)} from '{old_name}' to '{new_name}'"
                self.PrintOutput(result)
                return result
            return None
        except Exception as e:
            return f"Error: {str(e)}"

    def handle_get_function_start_end_address(self, args):
        try:
            address = int(args["address"], 16)
            function = idaapi.get_func(address)
            if function:
                start_address = hex(function.start_ea)
                end_address = hex(function.end_ea)
                result = {"start_address": start_address, "end_address": end_address}
                return result
            else:
                self.PrintOutput(f"No function found at address {hex(address)}")
            return f"No function found at address {hex(address)}"
        except Exception as e:
            return f"Error: {str(e)}"
    
    # def find_function_name_in_address(self, address, name):
    #     try:
    #         function = idaapi.get_func(address)
    #         if function:
    #             result = idaapi.decompile(function)
    #             if result:
    #                 for item in result.treeitems:
    #                     if item.op == idaapi.cot_call:
    #                         if item.cexpr.x.op == idaapi.cot_obj:
    #                             for arg in item.cexpr.a:
    #                                 if arg.op == idaapi.cot_obj and name in arg.print1(None):
    #                                     if arg.obj_ea != idc.BADADDR:
    #                                         return hex(arg.obj_ea)
    #                             if name in item.cexpr.x.print1(None):
    #                                 if item.cexpr.x.obj_ea != idc.BADADDR:
    #                                     return hex(item.cexpr.x.obj_ea)
    #                     elif item.op == idaapi.cot_cast:
    #                         if name in item.cexpr.x.print1(None):
    #                             if item.cexpr.x.obj_ea != idc.BADADDR:
    #                                 return hex(item.cexpr.x.obj_ea)
    #         return None
    #     except Exception as e:
    #         return None
        
    def get_name_info(self):
        name_info = []

        for i in range(ida_name.get_nlist_size()):
            ea = ida_name.get_nlist_ea(i)
            name = ida_name.get_short_name(ea)
            name_info.append((name, hex(ea)))

        return name_info

    def search_name(self, keyword):        
        search_results = []
        
        for name, ea in self.functions:
            if keyword.lower() in name.lower():
                search_results.append((name, ea))
        
        return search_results
    
    def handle_get_addresses_of_name(self, args):
        try:
            name = args["name"]
            r = self.search_name(name)
            self.PrintOutput(f"Search results for '{name}': {r}")
            return r
        except Exception as e:
            return f"Error: {str(e)}"
        
    def get_type_ea(self, ea):
        flag_types = []
        flags = ida_bytes.get_flags(ea)
        if idc.is_code(flags):
            flag_types.append("CODE")
        if idc.is_data(flags):
            flag_types.append("DATA")
        if idc.is_unknown(flags):
            flag_types.append("UNKNOWN")
        return f"Flags: {' | '.join(flag_types)}"

    def handle_get_address_type(self, args):
        try:
            address = int(args["address"], 16)
            flag_types = self.get_type_ea(address)
            size = idc.get_item_size(address)
            type_info = idc.get_type(address)
            if type_info:
                size_type = type_info
            else:
                size_type = {
                    1: "byte",
                    2: "word",
                    4: "dword",
                    8: "qword",
                    16: "oword"
                }.get(size, f"Unknown size ({size} bytes)")
            
            return f"Flags: {' | '.join(flag_types)}, Size: {size} bytes ({size_type})"
        except Exception as e:
            return f"Error: {str(e)}"
        
    def handle_get_xrefs_to(self, args):
        try:
            address = int(args["address"], 16)
            xrefs = []
            for xref in idautils.XrefsTo(address, 0):
                xrefs.append((hex(xref.frm), self.get_type_ea(xref.frm), idautils.XrefTypeName(xref.type)))
            result = xrefs
            self.PrintOutput(f'Xrefs to {hex(address)}: {result}')
            return result
        except Exception as e:
            self.PrintOutput(f"handle_get_xrefs_to Error: {str(e)}")
            return f"Error: {str(e)}"

    def handle_get_xrefs_from(self, args):
        try:
            address = int(args["address"], 16)
            xrefs = []
            for xref in idautils.XrefsFrom(address, 0):
                xrefs.append((hex(xref.to), self.get_type_ea(xref.to), idautils.XrefTypeName(xref.type)))
            result = xrefs
            self.PrintOutput(f'Xrefs from {hex(address)}: {result}')
            return result
        except Exception as e:
            self.PrintOutput(f"handle_get_xrefs_from Error: {str(e)}")
            return f"Error: {str(e)}"

    def handle_get_func_xrefs_to(self, args):
        try:
            address = int(args["address"], 16)
            if address != idc.BADADDR:
                xrefs = []
                for xref in idautils.XrefsTo(address, 0):
                    xrefs.append((hex(xref.frm), self.get_type_ea(xref.frm), idautils.XrefTypeName(xref.type)))
                result = xrefs
                self.PrintOutput(f'Xrefs to function at {hex(address)}: {result}')                
                return result
            self.PrintOutput(f"No function found at address {hex(address)}")
            return f"No function found at address {hex(address)}"
        except Exception as e:
            self.PrintOutput(f"handle_get_func_xrefs_to Error: {str(e)}")
            return f"Error: {str(e)}"
        
    def handle_print(self, args):
        message = args["message"]
        self.PrintOutput(message)
        return None

    def handle_set_comment(self, args):
        try:
            address = int(args["address"], 16)
            comment = args["comment"]
            idc.set_cmt(address, comment, 1)
            result = f"Set comment at {hex(address)}: {comment}"
            self.PrintOutput(result)
            return None
        except Exception as e:
            return f"Error: {str(e)}"
        
    def sanitize_json(self, json_string):
        json_string = re.sub(r'\\(?!["\\/bfnrt]|u[0-9a-fA-F]{4})', '', json_string)
        json_string = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', json_string)
        json_string = re.sub(r'"\s*\n\s*"', '""', json_string)
        json_string = re.sub(r'\s*\n\s*', '', json_string)
        
        return json_string
    
    def ParseResponse(self, response):
        try:
            response = self.sanitize_json(response)
            parsed_response = json.loads(response)
            return parsed_response
        except json.JSONDecodeError as e:
            traceback_details = traceback.format_exc()
            print(traceback_details)
            raise e
        except Exception as e:
            print(str(e))
            traceback_details = traceback.format_exc()
            print(traceback_details)
            raise e
                
    def PrintOutput(self, output_str):
        print(output_str)
        # self.chat_history.append(f"<b>System Message:</b> {output_str}")
        
def PLUGIN_ENTRY():
    return IDAAssistant()


