# plugins

repository: BinSigma/BinSourcerer-Dev
file: IDA Pro Side/IdaProTextExtractor_HRC.py
description: Development Branch of BinSourcerer
wanted_name: BinSourcerer Features Extractor
comment: Extract all the available features
created_at: 2014-09-30 15:37:23+00:00
pushed_at: 2014-09-30 15:38:44+00:00
url:
https://github.com/BinSigma/BinSourcerer-Dev/blob/d9f5f40d4ae4b8d14e8c8f04e9423cf25a3d0550/IDA%20Pro%20Side/IdaProTextExtractor_HRC.py

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




DIRECTION_FORW = 0x03

import idautils
import os

from idaapi import *
import hashlib
import idc

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

def PLUGIN_ENTRY():
    return TextExtract_t()



    


repository: wwkenwong/Deceiving-DNN-based-Binary-Matching
file: ida_binaryai_linux.py
description: Deceiving Deep Neural Networks-Based Binary Code Matching with Adversarial Programs
wanted_name: BinaryAI
created_at: 2022-06-30 07:39:54+00:00
pushed_at: 2023-02-10 15:46:38+00:00
url:
https://github.com/wwkenwong/Deceiving-DNN-based-Binary-Matching/blob/bdc04cd7a5d1448c6529f47d1ceb464e003b2be6/ida_binaryai_linux.py

content:
# coding: utf-8
import os
import platform

import idc
import ida_auto
import json
import idaapi
import idautils
import datetime
import binaryai as bai
#from PyQt5 import QtCore
from ida_hexrays import DecompilationFailure
#from PyQt5.QtWidgets import QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QWidget
from binaryai import BinaryAIException
import copy 
import random

def open_path(path):
    if not os.path.exists(path):
        os.makedirs(path)

POOL_OF_KEY = [<pasteme>]

class BinaryAILog(object):
    DEBUG = 0
    INFO = 1
    WARN = 2
    ERROR = 3
    level = INFO
    name = "BinaryAI"

    @staticmethod
    def log(level, msg, *args, **kwargs):
        if level >= BinaryAILog.level:
            if args:
                for v in args:
                    msg += str(args)
            if kwargs:
                msg += str(kwargs)

            print("[{}] {}".format(BinaryAILog.name, msg))

    @staticmethod
    def debug(msg, *args, **kwargs):
        BinaryAILog.log(BinaryAILog.DEBUG,
                        msg, *args, **kwargs)

    @staticmethod
    def skip(func_name, reason):
        BinaryAILog.log(BinaryAILog.INFO,
                        "{} is skipped because {}.".format(
                            func_name, reason))

    @staticmethod
    def fail(func_name, reason):
        BinaryAILog.log(BinaryAILog.WARN,
                        "{} failed because {}.".format(
                            func_name, reason))

    @staticmethod
    def success(func_name, status):
        BinaryAILog.log(BinaryAILog.INFO,
                        "{} successfully {}.".format(
                            func_name, status))

    @staticmethod
    def summary(succ, skip, fail, status):
        BinaryAILog.log(BinaryAILog.INFO,
                        "{} successfully {}, {} skipped, {} failed".format(
                            succ, status, skip, fail))

    @staticmethod
    def fatal(e):
        assert False, "[{}] {}".format(BinaryAILog.name, str(e))


class Config(dict):

    def __init__(self, path, default):
        self.path = path
        if not os.path.exists(path):
            json.dump(default, open(self.path, 'w'), indent=4)
        self.cfg = json.load(open(path))
        for k, v in default.items():
            if not (k in self.cfg and self.cfg[k] is not None):
                self.__setitem__(k, v)

    def __getitem__(self, key):
        return self.cfg[key]

    def __setitem__(self, key, val):
        if key in self.cfg and self.cfg[key] == val:
            return
        self.cfg[key] = val
        json.dump(self.cfg, open(self.path, 'w'), indent=4)


class BinaryAIConfig(Config):
    Default = {
        'token': '',
        'url': 'https://api.binaryai.tencent.com/v1/endpoint',
        'funcset': '',
        'usepublic': True,
        'topk': 10,
        'minsize': 3,
        'threshold': 0.9,
        'color': "0x817FFF",
        'first_use': True
    }

    def __init__(self, path=None, default=None):
        if path is None:
            path = os.path.join(idaapi.get_user_idadir(), idaapi.CFG_SUBDIR)
            if not os.path.exists(path):
                os.makedirs(path)
            path = os.path.join(path, "binaryai.cfg")
        if default is None:
            default = BinaryAIConfig.Default

        super(BinaryAIConfig, self).__init__(path, default)


class IDBStore(object):
    def __init__(self, name):
        self.netn = idaapi.netnode(str(name), 0, True)

    def __getitem__(self, item):
        val = self.netn.hashval(str(item))
        if val:
            val = json.loads(val)
        return val

    def __setitem__(self, key, value):
        self.netn.hashset(str(key), json.dumps(value).encode())

    def __delitem__(self, key):
        self.netn.hashdel(str(key))


class BinaryAIMark(object):

    def __init__(self):
        # record bai apply
        # {ea: {'name': name, 'score': score}
        self.record = IDBStore("BinaryAIMark")

    def add_record(self, ea, score):
        pfn = idaapi.get_func(ea)
        self.record[pfn.start_ea] = {
            'name': idaapi.get_ea_name(pfn.start_ea),
            'score': score
        }

    def is_bai_func(self, ea):
        v = self.record[ea]
        return v is not None

    @staticmethod
    def set_func_name(ea, name):
        if name.startswith("sub_"):
            idaapi.set_name(ea, "", idaapi.SN_AUTO)
        else:
            idaapi.set_name(ea, str(name), idaapi.SN_FORCE)

    def apply_bai_func(self, ea, name, score):
        if not self.is_bai_func(ea):
            self.add_record(ea, score)
        # in private lib
        if name.startswith("sub_"):
            name = "bai_" + name
        BinaryAIMark.set_func_name(ea, name)

    def apply_bai_high_score(self, ea, name, score):
        if self.is_bai_func(ea) \
                and score <= self.record[ea]['score']:
            return False
        self.apply_bai_func(ea, name, score)
        return True

    def revert_bai_func(self, ea):
        ea = idaapi.get_func(ea).start_ea
        if self.is_bai_func(ea):
            BinaryAIMark.set_func_name(
                ea,
                self.record[ea]['name']
            )
            del (self.record[ea])
            return True
        else:
            return False


"""global resource"""
bai_config = BinaryAIConfig()
bai_mark = BinaryAIMark()


class BinaryAIManager(object):
    def __init__(self):
        self.name = "BinaryAI"
        self._client = None
        self._funcset = None

    @property
    def client(self):
        if self._client is None:
            try:
                bai_config['token'] = random.choice(POOL_OF_KEY)
                self._client = bai.client.Client(bai_config['token'], bai_config['url'])
            except Exception:
                pass
        return self._client

    def update_token(self, token):
        self._client = None
        bai_config['token'] = token
        return self.client is not None

    @property
    def funcset(self):
        if bai_config['first_use']:
            self._funcset = bai.function.create_function_set(self.client)
            if self._funcset:
                bai_config['first_use'] = False

        if self._funcset is None:
            try:
                if bai.function.query_function_set(self.client, bai_config['funcset']):
                    self._funcset = bai_config['funcset']
            except BinaryAIException:
                pass
        return self._funcset

    def update_funcset(self, funcset):
        self._funcset = None
        bai_config['funcset'] = funcset
        return self.funcset is None

    def retrieve(self, ea, topk, funcset_ids, flag=1):
        func_id = self.upload(ea, None)
        if func_id:
            targets = bai.function.search_sim_funcs(self.client, func_id, funcset_ids=funcset_ids, topk=topk)
            if flag == 2:
                return targets, func_id
            return targets
        return None

    def upload(self, ea, funcset_id):
        func_feat = bai.ida.get_func_feature(ea)
        func_name = idaapi.get_func_name(ea)
        hf = idaapi.hexrays_failure_t()
        cfunc = idaapi.decompile(ea, hf, idaapi.DECOMP_NO_WAIT)
        if func_feat and func_name:
            func_id = bai.function.upload_function(
                self.client, func_name, func_feat, source_code=str(cfunc), funcset_id=funcset_id)
            return func_id
        return None


class SourceCodeViewer(object):
    class SourceCodeViewerUI(idaapi.simplecustviewer_t):
        def __init__(self, title):
            idaapi.simplecustviewer_t.__init__(self)
            self.ea = None
            self.query = None
            self.idx = None
            self.targets = None
            self.title = title
            self.Create(title)
            idaapi.set_code_viewer_is_source(idaapi.create_code_viewer(self.GetWidget(), 0x4))

        def set_user_data(self, ea, targets):
            self.idx = 0
            self.ea = idaapi.get_func(ea).start_ea
            self.query = idaapi.get_func_name(ea)
            self.targets = targets
            self._repaint()

        def _repaint(self):
            self.ClearLines()
            func = self.targets[self.idx]
            for line in SourceCodeViewer.source_code_comment(self.query, func, self.idx).split("\n"):
                self.AddLine(idaapi.COLSTR(line, idaapi.SCOLOR_RPTCMT))
            for line in SourceCodeViewer.source_code_body(func):
                self.AddLine(str(line))
            self.Refresh()

        def Show(self):
            widget = idaapi.get_current_widget()
            if idaapi.get_widget_title(widget) != self.title:
                if idaapi.get_widget_type(widget) != idaapi.BWN_PSEUDOCODE:
                    pseudo_view = idaapi.open_pseudocode(self.ea, 1)
                    pseudo_view.refresh_view(1)
                    widget = pseudo_view.toplevel
                pseudo_title = idaapi.get_widget_title(widget)

                idaapi.display_widget(self.GetWidget(),
                                      idaapi.PluginForm.WOPN_DP_TAB | idaapi.PluginForm.WOPN_RESTORE)
                idaapi.set_dock_pos(self.title, pseudo_title, idaapi.DP_RIGHT)

        def OnKeydown(self, vkey, shift):
            if shift == 0 and vkey == ord("K"):
                self.idx = (self.idx + len(self.targets) - 1) % len(self.targets)
                self._repaint()
            elif shift == 0 and vkey == ord("J"):
                self.idx = (self.idx + 1) % len(self.targets)
                self._repaint()

    @staticmethod
    def source_code_comment(query, func, idx=0):
        score = func["score"] if func["score"] < 1 else 1
        return """/*
    query:  {}
    target[{}]: {}
    target[{}] info: {}:{}
    score:  {:6f}
*/\n""".format(query,
               idx, func['function']['name'],
               idx, func['function']['sourceFile'], func['function']['sourceLine'],
               score)

    @staticmethod
    def source_code_body(func):
        body = func['function']['sourceCode'].split("\n")
        return filter(lambda l: not l.lstrip().startswith('#'), body)

    def __init__(self, title):
        self.view = None  # type: SourceCodeViewer.SourceCodeViewerUI
        self.title = title

    def is_visible(self):
        return self.view and self.view.GetWidget()

    def set_user_data(self, ea, targets):
        if not self.is_visible():
            self.view = SourceCodeViewer.SourceCodeViewerUI(self.title)
        self.view.set_user_data(ea, targets)
        self.view.Show()

    def get_current_info(self):
        return self.view.ea, \
               self.view.targets[self.view.idx]['function']['name'], \
               self.view.targets[self.view.idx]['score']


class BinaryAIOptionsForm(idaapi.Form):
    def __init__(self, mgr):
        self.mgr = mgr
        self.token = bai_config['token']
        self.funcset = bai_config['funcset']
        self.form_record = {}
        self.retrieve_list = ["Public", "Private"]
        super(BinaryAIOptionsForm, self).__init__(
            r'''STARTITEM 0
BUTTON YES* OK
BinaryAI Options
            {FormChangeCb}
            <Retrieve List  :{iretrieve_list}>
            <Topk           :{itopk}>
            <Threshold      :{ithreshold}>
            <Minsize        :{iminsize}>
            <Token          :{itoken}>
            <Function Set   :{ifuncset}>
            ''', {
                'iretrieve_list': self.DropdownListControl(
                    items=self.retrieve_list,
                    readonly=True,
                    selval=int(not bai_config['usepublic']),
                    width=32),
                'itopk': self.StringInput(value=str(bai_config["topk"])),
                'ithreshold': self.StringInput(value=str(bai_config["threshold"])),
                'iminsize': self.StringInput(value=str(bai_config["minsize"])),
                'itoken': self.StringInput(value=bai_config["token"]),
                'ifuncset': self.StringInput(value=bai_config["funcset"]),
                'FormChangeCb': self.FormChangeCb(self.on_form_change)
            }
        )
        self.Compile()

    def _get_float(self, ctl):
        try:
            return float(self.GetControlValue(ctl))
        except Exception:
            return -1

    def on_form_change(self, fid):
        if fid == self.iretrieve_list.id:
            v = self.GetControlValue(self.iretrieve_list)
            self.form_record['usepublic'] = False if v else True

        if fid == self.itopk.id:
            topk = int(self._get_float(self.itopk))
            if not (0 < topk <= 15):
                topk = BinaryAIManager.Default['topk']
            self.form_record['topk'] = topk

        if fid == self.ithreshold.id:
            threshold = self._get_float(self.ithreshold)
            if not (0 < threshold <= 1):
                threshold = BinaryAIManager.Default['threshold']
            self.form_record['threshold'] = threshold

        if fid == self.iminsize.id:
            minsize = int(self._get_float(self.iminsize))
            if not (1 <= minsize <= 5):
                minsize = BinaryAIManager.Default['minsize']
            self.form_record['minsize'] = minsize

        if fid == self.itoken.id:
            self.token = self.GetControlValue(self.itoken).strip()

        if fid == self.ifuncset.id:
            self.funcset = self.GetControlValue(self.ifuncset).strip()

        return 1

    @staticmethod
    def change_options(mgr, check_token=False, check_funcset=False):
        bai_options = BinaryAIOptionsForm(mgr)
        if bai_options.Execute():
            for k, v in bai_options.form_record.items():
                bai_config[k] = v

            if check_token or bai_options.token != bai_config['token']:
                mgr.update_token(bai_options.token)
                if not mgr.client:
                    idaapi.warning("Wrong token!")
                    BinaryAIOptionsForm.change_options(mgr, check_token=True)
            if check_funcset or bai_options.funcset != bai_config['funcset']:
                mgr.update_funcset(bai_options.funcset)
                if not mgr.funcset:
                    idaapi.warning("Wrong function set!")
                    BinaryAIOptionsForm.change_options(mgr, check_funcset=True)


# class CopyrightWindow(QWidget):
#     def __init__(self, ver, year, mgr):
#         super(CopyrightWindow, self).__init__()
#         self.mgr = mgr
#         self.setWindowFlags(QtCore.Qt.WindowMinimizeButtonHint)
#         self.setFixedSize(400, 200)
#         self.setWindowTitle("BinaryAI")
#         layoutCopyright = QVBoxLayout()
#         layoutCopyright.setSpacing(5)
#         layoutButtons = QHBoxLayout()
#         layoutButtons.setSpacing(10)
#         mainLayout = QVBoxLayout()

#         label1 = QLabel()
#         label2 = QLabel()
#         label3 = QLabel()

#         label1.setText("BinaryAI v{}".format(ver))
#         label2.setText("(c) Copyright {}, Tencent Security KEEN Lab".format(year))
#         label3.setText("<a href='https://binaryai.readthedocs.io/'>https://binaryai.readthedocs.io/</a>")
#         label3.setOpenExternalLinks(True)
#         label1.setAlignment(QtCore.Qt.AlignCenter)
#         label2.setAlignment(QtCore.Qt.AlignCenter)
#         label3.setAlignment(QtCore.Qt.AlignCenter)

#         btn1 = QPushButton()
#         btn2 = QPushButton()
#         btn1.setText("OK")
#         btn2.setText("Options")
#         btn1.setFixedSize(60, 20)
#         btn2.setFixedSize(60, 20)

#         btn1.clicked.connect(self.close)
#         btn2.clicked.connect(self.showOptions)

#         layoutCopyright.addWidget(label1)
#         layoutCopyright.addWidget(label2)
#         layoutCopyright.addWidget(label3)
#         layoutButtons.addWidget(btn1)
#         layoutButtons.addWidget(btn2)

#         mainLayout.addLayout(layoutCopyright)
#         mainLayout.addLayout(layoutButtons)

#         self.setLayout(mainLayout)

    def showOptions(self):
        BinaryAIOptionsForm.change_options(self.mgr)


class BinaryAIOperations(object):
    def __init__(self, mgr):
        self.mgr = mgr  # type: BinaryAIManager

    def check_before_use(self, check_funcset=False):
        if not self.mgr.client:
            idaapi.warning("Wrong token!")
            BinaryAIOptionsForm.change_options(self.mgr, check_token=True)
            return False
        if (not self.mgr.funcset) and \
                (check_funcset or not bai_config['usepublic']):
            idaapi.warning("Wrong function set!")
            BinaryAIOptionsForm.change_options(self.mgr, check_funcset=True)
            return False
        return True

    def retrieve(self, ea, cview):
        if not self.check_before_use():
            return
        func_name = idaapi.get_func_name(ea)
        funcset_ids = [self.mgr.funcset] if not bai_config['usepublic'] else None

        try:
            targets = self.mgr.retrieve(ea, bai_config['topk'], funcset_ids)
        except DecompilationFailure as e:
            BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
        except BinaryAIException as e:
            BinaryAILog.fatal(e)

        if targets is None:
            BinaryAILog.skip(func_name, "get function feature error")
            return

        cview.set_user_data(ea, targets)

    def _match_with_check(self, ea, topk, funcset_ids):
        fail, skip, succ = -1, 0, 1
        # < minsize
        pfn = idaapi.get_func(ea)
        if idaapi.FlowChart(pfn).size < bai_config['minsize']:
            return skip
        # do match
        try:
            targets = self.mgr.retrieve(ea, topk=bai_config['topk'], funcset_ids=funcset_ids)
        except DecompilationFailure as e:
            BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
            return fail
        except BinaryAIException as e:
            idaapi.hide_wait_box()
            BinaryAILog.fatal(e)
        if targets is None:
            return fail
        if not bai_mark.apply_bai_high_score(
                ea,
                targets[0]['function']['name'],
                targets[0]['score']):
            return skip
        return succ

    def match_funcs(self, funcs):
        if not self.check_before_use():
            return

        i, fail, skip, succ = 0, 0, 0, 0

        def stop():
            idaapi.hide_wait_box()
            BinaryAILog.summary(succ, skip, fail, "matched")

        funcs_len = len(funcs)
        idaapi.show_wait_box("Matching... (0/{})".format(funcs_len))
        funcset_ids = [self.mgr.funcset] if not bai_config['usepublic'] else None
        for ea in funcs:
            # refresh process status
            i += 1
            idaapi.replace_wait_box("Matching... ({}/{})".format(i, funcs_len))
            # check cancelled or not
            if idaapi.user_cancelled():
                stop()
                return
            status = self._match_with_check(ea, bai_config['topk'], funcset_ids)
            if status == 1:
                succ += 1
            elif status == 0:
                skip += 1
            else:
                fail += 1
        stop()

    def upload(self, ea):
        try:
            func_id = self.mgr.upload(ea, self.mgr.funcset)
        except DecompilationFailure as e:
            BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
        except BinaryAIException as e:
            BinaryAILog.fatal(e)

        if func_id:
            BinaryAILog.success(idaapi.get_func_name(ea), "uploaded")

    def upload_funcs(self, funcs):
        if not self.check_before_use(check_funcset=True):
            return

        i, succ, skip, fail = 0, 0, 0, 0

        def stop():
            idaapi.hide_wait_box()
            BinaryAILog.summary(succ, skip, fail, "uploaded")

        funcs_len = len(funcs)
        idaapi.show_wait_box("Uploading... (0/{})".format(funcs_len))
        for ea in funcs:
            i += 1
            idaapi.replace_wait_box("Uploading... ({}/{})".format(i, funcs_len))

            if idaapi.user_cancelled():
                stop()
                return
            # < minsize
            pfn = idaapi.get_func(ea)
            if idaapi.FlowChart(pfn).size < bai_config['minsize']:
                skip += 1
                continue
            # try upload
            func_id = None
            try:
                func_id = self.mgr.upload(ea, self.mgr.funcset)
            except DecompilationFailure as e:
                BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
                fail += 1
                continue
            except BinaryAIException as e:
                stop()
                BinaryAILog.fatal(e)
            # fail
            if not func_id:
                fail += 1
                continue
            succ += 1
        stop()

    def apply(self, cview):
        ea, name, score = cview.get_current_info()
        bai_mark.apply_bai_func(ea, name, score)

    def revert_funcs(self, funcs):
        i, succ, skip = 0, 0, 0

        def stop():
            idaapi.hide_wait_box()
            BinaryAILog.summary(succ, skip, 0, "reverted")

        funcs_len = len(funcs)
        idaapi.show_wait_box("Reverting... (0/{})".format(funcs_len))
        for ea in funcs:
            i += 1
            idaapi.replace_wait_box("Reverting... ({}/{})".format(i, funcs_len))

            if idaapi.user_cancelled():
                stop()
                return

            if bai_mark.revert_bai_func(ea):
                succ += 1
            else:
                skip += 1
        stop()


class UIManager:
    class UIHooks(idaapi.UI_Hooks):
        is_function_window_hooked = False

        def finish_populating_widget_popup(self, widget, popup, ctx=None):
            if idaapi.get_widget_type(widget) == idaapi.BWN_FUNCS:
                idaapi.attach_action_to_popup(widget, popup, "BinaryAI:MatchSelected", "BinaryAI/")
                idaapi.attach_action_to_popup(widget, popup, "BinaryAI:UploadSelected", "BinaryAI/")

                funcs = map(idaapi.getn_func, ctx.chooser_selection)
                funcs = map(lambda func: func.start_ea, funcs)
                for ea in funcs:
                    if bai_mark.is_bai_func(ea):
                        idaapi.attach_action_to_popup(widget, popup, "BinaryAI:RevertSelected", "BinaryAI/")
                        break

            if idaapi.get_widget_type(widget) == idaapi.BWN_CUSTVIEW:
                idaapi.attach_action_to_popup(widget, popup, "BinaryAI:Apply", "BinaryAI/")

        def get_chooser_item_attrs(self, chooser, n, attrs):
            func = idaapi.getn_func(n)
            if bai_mark.is_bai_func(func.start_ea):
                attrs.color = int(bai_config['color'], 16)

        def updating_actions(self, ctx):
            if not self.is_function_window_hooked:
                self.is_function_window_hooked = \
                    idaapi.enable_chooser_item_attrs("Functions window", True)

    class ActionHandler(idaapi.action_handler_t):
        def __init__(self, name, label, shortcut=None, tooltip=None, icon=-1, flags=0):
            idaapi.action_handler_t.__init__(self)
            self.name = name
            self.action_desc = idaapi.action_desc_t(name, label, self, shortcut, tooltip, icon, flags)

        def register_action(self, callback, toolbar_name=None, menupath=None):
            self.callback = callback
            if not idaapi.register_action(self.action_desc):
                return False
            if toolbar_name and not idaapi.attach_action_to_toolbar(toolbar_name, self.name):
                return False
            if menupath and not idaapi.attach_action_to_menu(menupath, self.name, idaapi.SETMENU_APP):
                return False
            return True

        def activate(self, ctx):
            self.callback(ctx)

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

    def __init__(self, name, bai_mgr):
        self.name = name
        self.bai_mgr = bai_mgr
        self.operations = BinaryAIOperations(bai_mgr)
        self.cview = SourceCodeViewer(self.name)
        self.hooks = UIManager.UIHooks()
        self.about_window = None

    def register_actions(self):
        toolbar_name, menupath = self.name, self.name
        idaapi.create_toolbar(toolbar_name, self.name)
        idaapi.create_menu(menupath, self.name, "Help")

        # UIManager.ActionHandler(self.name, self.name).register_action(self.binaryai_callback, toolbar_name)
        # action = UIManager.ActionHandler("BinaryAI:About", "About", "")
        # action.register_action(self.binaryai_callback, menupath=menupath)
        # action = UIManager.ActionHandler("BinaryAI:RetrieveFunction", "Retrieve function", "Ctrl+Shift+d", icon=99)
        # action.register_action(self.retrieve_callback, toolbar_name, menupath)
        # action = UIManager.ActionHandler("BinaryAI:UploadFunction", "Upload function", "", icon=97)
        # action.register_action(self.upload_callback, toolbar_name, menupath)
        # action = UIManager.ActionHandler("BinaryAI:MatchAll", "Match all functions", "", icon=188)
        # action.register_action(self.match_all_callback, toolbar_name, menupath)
        # action = UIManager.ActionHandler("BinaryAI:UploadAll", "Upload all functions", "", icon=88)
        # action.register_action(self.upload_all_callback, toolbar_name, menupath)

        # apply_action = UIManager.ActionHandler("BinaryAI:Apply", "Apply")
        # apply_action.register_action(self.apply_callback)

        # match_action = UIManager.ActionHandler("BinaryAI:MatchSelected", "Match")
        # upload_action = UIManager.ActionHandler("BinaryAI:UploadSelected", "Upload")
        # revert_action = UIManager.ActionHandler("BinaryAI:RevertSelected", "Revert")
        if match_action.register_action(self.selected_callback) and \
                upload_action.register_action(self.selected_callback) and \
                revert_action.register_action(self.selected_callback):
            self.hooks.hook()
            return True
        return False

    # def binaryai_callback(self, ctx):
    #     self.about_window = CopyrightWindow(bai.__version__, datetime.datetime.now().year, self.bai_mgr)
    #     self.about_window.show()

    def retrieve_callback(self, ctx, ea=None):
        func_ea = idaapi.get_screen_ea() if ea is None else ea
        self.operations.retrieve(func_ea, self.cview)

    def match_all_callback(self, ctx):
        self.operations.match_funcs(list(idautils.Functions()))

    def upload_callback(self, ctx):
        self.operations.upload(idaapi.get_screen_ea())

    def upload_all_callback(self, ctx):
        self.operations.upload_funcs(list(idautils.Functions()))

    def selected_callback(self, ctx):
        funcs = map(idaapi.getn_func, ctx.chooser_selection)
        funcs = [func.start_ea for func in funcs]
        if ctx.action == "BinaryAI:MatchSelected":
            self.operations.match_funcs(funcs)
        if ctx.action == "BinaryAI:UploadSelected":
            self.operations.upload_funcs(funcs)
        if ctx.action == "BinaryAI:RevertSelected":
            self.operations.revert_funcs(funcs)

    def apply_callback(self, ctx):
        self.operations.apply(self.cview)


class Plugin(idaapi.plugin_t):
    wanted_name = "BinaryAI"
    comment, help, wanted_hotkey = "", "", ""
    flags = idaapi.PLUGIN_FIX | idaapi.PLUGIN_HIDE

    def init(self):
        if idaapi.init_hexrays_plugin():
            bai_mgr = BinaryAIManager()
            ui_mgr = UIManager(Plugin.wanted_name, bai_mgr)
            if ui_mgr.register_actions():
                return idaapi.PLUGIN_OK
        return idaapi.PLUGIN_SKIP

    def run(self, ctx):
        return

    def term(self):
        return


def PLUGIN_ENTRY():
    BinaryAILog.level = BinaryAILog.DEBUG
    return Plugin()


def get_user_idadir():
    system = platform.system()
    if system == 'Windows':
        return os.path.join(os.getenv('APPDATA'), "Hex-Rays", "IDA Pro")
    elif system in ['Linux', 'Darwin']:
        return os.path.join(os.getenv('HOME'), ".idapro")
    else:
        return ""


def cmd_upload(funcset=bai_config['funcset']):
    succ = 0
    bai_mgr = BinaryAIManager()
    for ea in idautils.Functions():
        try:
            bai_mgr.upload(ea, funcset)
            succ += 1
        except Exception:
            continue
    return succ


def cmd_match(funcset_ids=None):
    bai_mgr = BinaryAIManager()
    output_json = {}
    for ea in idautils.Functions():
        try:
            #targets, func_id = bai_mgr.retrieve(ea, bai_config['topk'], funcset_ids, 2)
            targets, func_id = bai_mgr.retrieve(ea, 500, funcset_ids, 2)            
        except Exception as e:
            print(str(e))
            continue

        if targets and func_id:
            gg = copy.copy(targets)
            for iii in range(0,len(gg)):
                # dict_keys(['id', 'language', 'name', 'sourceCode', 'sourceFile', 'sourceLine'])
                tmp = gg[iii]['function'].pop('language')
                tmp = gg[iii]['function'].pop('sourceCode')
                tmp = gg[iii]['function'].pop('sourceLine')
            bai_mark.apply_bai_high_score(
                ea,
                targets[0]['function']['name'],
                targets[0]['score'])
        cur_json = {'id': func_id, 'score': targets[0]['score'],
                    'target': {'id': targets[0]['function']['id'], 'name': targets[0]['function']['name']},
                    'full' :gg }
        output_json[str(ea)] = cur_json
    a = ida_nalt.get_input_file_path()
    folder = "\\".join(a.split("\\"))#[:-1])
    #path = os.path.join(folder, "output.json")
    path = folder+"_output.json"
    #path = os.path.join(get_user_idadir(), "output.json")
    print(path)
    with open(path, "w") as f:
        f.write(str(json.dumps(output_json, sort_keys=True, indent=2)))
        f.close()


def match_by_addr(ea,funcset_ids=None):
    bai_mgr = BinaryAIManager()
    output_json = {}
    addr_list = list(idautils.Functions())
    addr_list = [int(x) for x in addr_list]
    if ea in addr_list:
        try:
            #targets, func_id = bai_mgr.retrieve(ea, bai_config['topk'], funcset_ids, 2)
            targets, func_id = bai_mgr.retrieve(ea, 500, funcset_ids, 2)            
        except Exception as e:
            print(str(e))
            return None 
        if targets and func_id:
            gg = copy.copy(targets)
            for iii in range(0,len(gg)):
                # dict_keys(['id', 'language', 'name', 'sourceCode', 'sourceFile', 'sourceLine'])
                tmp = gg[iii]['function'].pop('language')
                tmp = gg[iii]['function'].pop('sourceCode')
                tmp = gg[iii]['function'].pop('sourceLine')
            bai_mark.apply_bai_high_score(
                ea,
                targets[0]['function']['name'],
                targets[0]['score'])
        cur_json = {'id': func_id, 'score': targets[0]['score'],
                    'target': {'id': targets[0]['function']['id'], 'name': targets[0]['function']['name']},
                    'full' :gg }
        output_json[str(ea)] = cur_json
        a = ida_nalt.get_input_file_path()
        #folder = "\\".join(a.split("\\"))#[:-1])
        #path = os.path.join(folder, "output.json")
        path = a+"_output.json"
        #path = os.path.join(get_user_idadir(), "output.json")
        print(path)
        with open(path, "w") as f:
            f.write(str(json.dumps(output_json, sort_keys=True, indent=2)))
            f.close()


if __name__ == "__main__":
    addr = idc.ARGV[1]
    addr = int(addr,16)
    print(addr)
    ida_auto.auto_wait()
    match_by_addr(addr)
    ida_pro.qexit(0)


repository: opensrcsec/ida-rap-decode
file: rap_decode.py
description: IDA Pro Python plugin to decode RAP return hash sequences in the Linux kernel
wanted_name: RAP Decode
created_at: 2021-11-28 14:47:50+00:00
pushed_at: 2021-11-29 14:46:52+00:00
url:
https://github.com/opensrcsec/ida-rap-decode/blob/ca53f3edd26b6207c78b173e13ff71314dc849e8/rap_decode.py

content:
# BSD 3-Clause License
#
# Copyright (c) 2021, Open Source Security, Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Author: Pawel Wieczorkiewicz <wipawel@grsecurity.net>
#
from ida_allins import NN_jmp, NN_jmpshort, NN_jmpni, NN_jmpfi
from idautils import Functions, DecodePreviousInstruction
from ida_bytes import is_code, get_flags, del_items, get_item_size
from ida_auto import auto_wait
from ida_funcs import get_func
from ida_ua import create_insn
from idaapi import *

def is_jmp(insn):
  return insn and insn.itype in [NN_jmp, NN_jmpshort, NN_jmpni, NN_jmpfi]


class rap_decode_t(plugin_t):
    flags = PLUGIN_KEEP
    comment = ""
    help = "Run this plugin once and wait for it to finish"
    wanted_name = "RAP Decode"
    wanted_hotkey = "Alt-F8"

    MAX_ALIGN_LENGTH = 15

    def init(self):
        return PLUGIN_OK

    def run(self, arg):
        print("Running RAP Decode plugin...")
        self.rap_decode()

    def term(self):
        pass

    def rap_decode(self):
        for func_ea in Functions():
            func = get_func(func_ea)
            flowchart = FlowChart(func)

            for bb in flowchart:
                data_ea = bb.end_ea
                if is_code(get_flags(data_ea)):
                    continue

                last_insn = DecodePreviousInstruction(data_ea)
                if not is_jmp(last_insn):
                    continue

                # Delete all items from current EA up to maximal alignment length
                for ea in range(data_ea, data_ea + self.MAX_ALIGN_LENGTH):
                    del_items(ea)

                ea = data_ea
                while ea < data_ea + self.MAX_ALIGN_LENGTH:
                    size = create_insn(ea)
                    # Skip item if unable to create instruction for it
                    ea += size if size > 0 else get_item_size(ea)

        auto_wait()


def PLUGIN_ENTRY():
    return rap_decode_t()


repository: w4kfu/VG_Helper
file: ida_plugin/vtable_helper.py
description: Toolz for helping vuln research inside video game
wanted_name: vtable_helper
comment: Plugin for creating and commenting automatically vtable stuff
created_at: 2012-12-06 09:41:28+00:00
pushed_at: 2012-12-07 04:46:16+00:00
url:
https://github.com/w4kfu/VG_Helper/blob/08e8d7b8b9d24c8ef6d398f885fad80f2625b914/ida_plugin/vtable_helper.py

content:
import idc 
import idaapi
import idautils

class rename_idp_hook_t(idaapi.IDP_Hooks):
	def __init__(self):
		idaapi.IDP_Hooks.__init__(self)
		self.cmd = idaapi.cmd
	
	# http://code.google.com/p/idapython/source/browse/trunk/swig/idp.i?r=315#448
	# keyword "pass" is missing
	def custom_outop(self, op):
		pass

	# http://code.google.com/p/idapython/source/browse/trunk/swig/idp.i?r=315#459
	# keyword "pass" is missing
	def custom_mnem(self):
		pass

	def renamed(self, ea, name, local_name):
		print("Name %s" % name)


class VtableHelperForm(Form):
	def __init__(self):
		Form.__init__(self, r"""Vtable Helper
		<#Class Name :{txtClassName}>
		<#Class Size :{iClassSize}>
		<#Vtable Addr :{iVtableAddr}>

		""", {
		    	'txtClassName'	: 	Form.StringInput(),
			'iClassSize'	:	Form.NumericInput(),
			'iVtableAddr' 	:	Form.NumericInput(),
		})

	def OnFormChange(self, fid):
		return 1

	def Show(self):
        	if not self.Compiled():
            		_, args = self.Compile()
        	ok = self.Execute()
		if (ok != 0):
			NClass = TClass(self.txtClassName.value, 
					self.iClassSize.value, 
					self.iVtableAddr.value)
			NClass.printdbg()
			NClass.create_struct()
			NClass.create_vtable()
		return ok
		
class vtable_helper(idaapi.plugin_t):
	flags = 0
	comment = "Plugin for creating and commenting automatically vtable stuff"
	help = "RTFC"
	wanted_name = "vtable_helper"
	wanted_hotkey = "Alt-F5"

	def init(self):
		self.idphook = None
 		self.idphook = rename_idp_hook_t()
    		self.idphook.hook()
		self.segname = "vtablehelper"
		self.getinfo()
		self.seg = 0
		return idaapi.PLUGIN_KEEP

	def getinfo(self):
		for seg in idautils.Segments():
			if self.segname == idc.SegName(seg):
				print("[+] Segment vtablehelper found !")
				self.infostart = idc.SegStart(seg)
				self.infoend = idc.SegEnd(seg)
				self.seg = seg

	def saveinfo(self):
		if (self.seg == 0):
			i = AskYN(0, "Segment vtablehelper doesn't exist, do you want to create it for saving information ?")
			if (i == -1 or i == 0):
				return

	def run(self, args):
        	f = VtableHelperForm()
        	ok = f.Show()
        	if ok == 0:
            		f.Free()

	def term(self):
		self.saveinfo()
		if self.idphook:
			self.idphook.unhook()

class TClass():
	def __init__(self, name, size, addrvtable):
		self.name = name
		self.addr_vtable = addrvtable
		self.size = size
		self.vtable = 0

	def printdbg(self):
		print("ClassName : %s" % self.name)
		print("ClassSize : %X" % self.size)
		print("Addr_vtable : %X" % self.addr_vtable)

	def create_struct(self):
		if (self.size == 0):
			print("Invalid class size")
		struc_id = GetStrucIdByName(self.name)
		if (struc_id != -1):
			i = AskYN(0, "A class structure for %s already exists. Are you sur you want to remake it ?" % self.name)
			if (i == -1 or i == 0):
				self.struct = struc_id
				return
			DelStruc(struc_id)
		self.struct = AddStrucEx(-1, self.name, 0)
		for i in xrange(0, self.size / 4):
			AddStrucMember(self.struct, "field_" + str(i), i * 4, FF_DWRD | FF_DATA, -1, 4)
		if (self.size % 4) != 0:
			AddStrucMember(self.struct, "field_" + str(self.size / 4), (self.size / 4) * 4, FF_DATA, -1, self.size % 4)

	def create_vtable(self):
		SetMemberName(self.struct, 0, "vptr")
		self.vtable = Vtable(self.addr_vtable)
		self.vtable.getsize()

class Vtable():
	def __init__(self, addr):
		self.addr = addr
		self.size = 0
		self.entries = []

	def getsize(self):
		actual_ea = self.addr
		while (True):
			# first entry case
			f = idc.GetFlags(actual_ea)
			if (len(self.entries) == 0):
				if (not (idc.isRef(f) and (idc.hasName(f) or (f & FF_LABL)))):
					print("[-] Not an entry of vtable")
					return 0
			elif (idc.isRef(f) and (idc.hasName(f) or (f & FF_LABL))):
				# next vtable ?
				break
			if (not idc.hasValue(f) or not idc.isData(f)):
				break
			c = idc.Dword(actual_ea)
			if c:
				f = idc.GetFlags(c)
				if (not idc.hasValue(f) or not idc.isCode(f) or idc.Dword(c) == 0):
					break
			else:
				break
			self.entries.append(actual_ea)
			actual_ea += 4
		print("[+] Vtable %08X - %08X, methods : %d" % (self.addr, actual_ea, (actual_ea - self.addr) / 4))




def PLUGIN_ENTRY():
	return vtable_helper()


repository: sfinktah/sfida
file: idarest.py
description: idapro python toolkit
url: https://github.com/sfinktah/sfida/blob/b9a8307b5c3e3cbfeb2372757a0b3d0da84142ea/idarest.py
wanted_name: IDA Rest API (Terminate)
comment: Interface to IDA Rest API
created_at: 2021-03-01 10:05:12+00:00
pushed_at: 2024-10-23 03:13:14+00:00
content:
import sys
import os
import cgi
import errno
import itertools
import json, pickle
import re
import requests
import socket
import threading
import time
import traceback
import atexit
import inspect
import itertools
import urllib.request, urllib.error, urllib.parse as urlparse
from code import InteractiveInterpreter
from http.server import BaseHTTPRequestHandler, HTTPServer
from queue import Queue, Empty, Full
from socketserver import ThreadingMixIn

# pip install superglobal split_paren
from superglobals import setglobal, getglobal, superglobals
from split_paren import paren_multisplit

try:
    from .idarest_mixins import IdaRestConfiguration, BorrowStdOut, IdaRestLog
except:
    from idarest_mixins import IdaRestConfiguration, BorrowStdOut, IdaRestLog

# for testing outside ida
try:
    import idc
    import ida_idaapi
    import ida_kernwin
    import ida_loader
    import idaapi
    import idautils
    import ida_pro
    from PyQt5 import QtWidgets, QtCore
    delayed_exec_timer = QtCore.QTimer()
except:
    class idc:
        @staticmethod
        def msg(s):
            print(s)
        def get_idb_path():
            return '.'

    class ida_idaapi:
        plugin_t = object
        PLUGIN_UNL = PLUGIN_KEEP = 0

_API_IDB = idc.get_idb_path()
_API_FILE = __file__

def ir_static_vars(**kwargs):
    def decorate(func):
        for k in kwargs:
            setattr(func, k, kwargs[k])
        return func
    return decorate

def _classname(instance):
    return getattr(getattr(instance, '__class__', object), '__name__', '')

def _asBytes(s):
    if isinstance(s, str):
        return s.encode('utf-8')
    return s

def _asStringRaw(o):
    return o.decode('raw_unicode_escape') if isinstance(o, (bytes, bytearray)) else o

def _execute_sync(cmd, *args):
    if ida_pro.is_main_thread():
        return cmd()
    request = ida_kernwin.execute_sync(cmd, ida_kernwin.MFF_WRITE)
    #  request = ida_kernwin.execute_sync(cmd, ida_kernwin.MFF_NOWAIT)
    #  print("[_execute_sync] request: {}".format(request))
    #  HTTPRequestHandler.delayed_call(lambda: ida_kernwin.cancel_exec_request(request))

"""
IDA Pro Plugin Interface

Define an IDA Python plugin required class and function.
"""

# If used as a plugin everything just becomes harder, though it does clean-up used ports
# properly. 
MENU_PATH = 'Edit/Other'
class idarest_plugin_t(IdaRestConfiguration, IdaRestLog, ida_idaapi.plugin_t):
    count = 0
    flags = ida_idaapi.PLUGIN_UNL
    comment = "Interface to IDA Rest API"
    help = "IDA Rest API for basic RE tool interoperability"
    wanted_name = "IDA Rest API (Terminate)"
    wanted_hotkey = ""

    @staticmethod
    def test_bind_port(port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                # s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                s.bind((idarest_plugin_t.config['api_bind_ip'], port))
            except socket.error as e:
                if e.errno != errno.EADDRINUSE:
                    if idarest_plugin_t.config['api_debug']: idc.msg(e)
                else:
                    if idarest_plugin_t.config['api_debug']: idc.msg("[idarest port-in-use]: {} ".format(port))
                return False
        return True

    def init(self):
        if idarest_plugin_t.count:
            return ida_idaapi.PLUGIN_KEEP
        idarest_plugin_t.count += 1
        if idarest_plugin_t.config['api_verbose']: idc.msg("[idarest_plugin_t::init]\n")
        super(idarest_plugin_t, self).__init__()
        #  self.load_configuration()

        # 10 queues for standard operation, 1 queue for iter/generator usage
        
        self.state = None
        # I replaced the previous menu contexts with hotkey contexts, but then
        # decided they were pretty useless
        #  new_ctx1 = ida_kernwin.add_hotkey("Alt-7", lambda *a: self.start())
        #  new_ctx2 = ida_kernwin.add_hotkey("Alt-8", lambda *a: self.stop())
        #  new_ctx2 = ida_kernwin.add_hotkey("Alt-9", lambda *a: self.term())
        #  self.ctxs = [new_ctx1, new_ctx2]
        self.worker = None
        self.timer = None
        for port in range(idarest_plugin_t.config['api_port'], 65535):
            if idarest_plugin_t.test_bind_port(port):
                self.port = port
                break
        self.host = idarest_plugin_t.config['api_bind_ip']
        # IDA 7 used totally new menu contexts, and I cba translating them.
        #  ret = self._add_menus()
        if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::init] done\n")
        return self.start()

    def is_alive(self):
        return self.worker and self.worker.is_alive()

    def start(self, *args):
        # load_and_run_plugin
        if idarest_plugin_t.config['api_verbose']: idc.msg("[idarest_plugin_t::start]\n")
        if self.worker and self.worker.is_alive():
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::start] already running\n")
            return ida_idaapi.PLUGIN_SKIP

        worker = Worker(self.host, self.port)
        self.worker = worker
        self.worker.start()

        if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::start] worker started: port {}\n".format(self.port))

        timer = Timer(self.host, self.port)
        self.timer = timer
        self.timer.start()

        if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::start] timer started\n")

        def cleanup():
            self.register(self.host, self.port, unregister=True)
            self.log("**atexit** cleanup")
            if worker and worker.is_alive():
                self.log("[idarest_plugin_t::start::cleanup] stopping..\n")
                worker.stop()
                self.log("[idarest_plugin_t::start::cleanup] joining..\n")
                worker.join()
                self.log("[idarest_plugin_t::start::cleanup] stopped\n")

            if timer and timer.is_alive() and not timer.stopped():
                self.log("[idarest_plugin_t::start::cleanup] stopping..\n")
                timer.stop()
                self.log("[idarest_plugin_t::start::cleanup] joining..\n")
                timer.join()
                self.log("[idarest_plugin_t::start::cleanup] stopped\n")

        print('[idarest_plugin_t::start] registered atexit cleanup')
        atexit.register(cleanup)


        #  idarest_plugin_t.register(self.host, self.port)

        return ida_idaapi.PLUGIN_KEEP

    def run(self, *args):
        if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::run] {}\n".format(args))
        pass

    def stop(self):
        if not self.timer or not self.timer.is_alive():
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::stop] timer was not running\n")
        else:
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::stop] stopping master timer..\n")
            self.timer.stop()
            #  idc.msg("[idarest_plugin_t::stop] joining..\n")
            #  self.timer.join()
        if not self.worker or not self.worker.is_alive():
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::stop] worker was not running\n")
        else:
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::stop] stopping RESTful service..\n")
            self.worker.stop()
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::stop] joining..\n")
            self.worker.join()
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::stop] stopped\n")
        idarest_plugin_t.register(self.host, self.port, unregister=True)
        idarest_main.instance = None
        if 'ir' in globals():
            globals().pop('ir')
        if idarest_plugin_t.count > 0:
            idarest_plugin_t.count -= 1


    def term(self):
        try:
            self.stop()
        except Exception as e:
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::term] {}: {}\n".format(e.__class__.__name__, str(e)))
        #  for ctx in self.ctxs:
            #  ida_kernwin.del_hotkey(ctx)
            #
            #

    @staticmethod
    def register(host, port, unregister=False):
        if unregister:
            url = 'http://{}:{}{}/unregister'.format(idarest_plugin_t.config['master_host'], idarest_plugin_t.config['master_port'], idarest_plugin_t.config['api_prefix'])
        else:
            url = 'http://{}:{}{}/register'.format(idarest_plugin_t.config['master_host'], idarest_plugin_t.config['master_port'], idarest_plugin_t.config['api_prefix'])
        if idarest_plugin_t.config['api_verbose']: idc.msg("[idarest_plugin_t::register] trying to connect to master at {}\n".format(url))
        master_plugin = _API_FILE.replace('idarest.py', 'idarest_master.py')
        connect_timeout = 1
        read_timeout = 1
        try:
            p = _API_IDB
            idb = os.path.split(os.path.split(p)[0])[1] + '/' + os.path.splitext(os.path.split(os.path.split(p)[1])[1])[0]
            
            r = requests.get(
                    url, 
                    params={ 
                        'host': host,
                        'port': port,
                        'idb': idb,
                    }, 
                    timeout=(connect_timeout, read_timeout))
            if r.status_code == 200:
                if r.content:
                    if idarest_plugin_t.config['api_verbose']: idc.msg("[idarest_plugin_t::register] master responded correctly: {}, {}\n".format(url, r.content))
                else:
                    if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::register] master failed to return data: {}, {}\n".format(url, r.content))
                return
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::register] master returned status_code: {}\n".format(r.status_code))
        except Exception as e:
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::register] failed to connect to master at {}: {}\n".format(url, e.__class__.__name__))
            if not unregister:
                if e.__class__.__name__ in ('ConnectTimeout', 'ConnectionError'):
                    if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::register] launching new master\n")
                    _execute_sync(lambda: ida_loader.load_plugin(master_plugin), ida_kernwin.MFF_WRITE)
                else:
                    if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::register] exception not of a type to trigger loading new master\n")

                
        except requests.exceptions.ConnectionError as e:
            if idarest_plugin_t.config['api_info']: 
                idc.msg("[idarest_plugin_t::register] failed to connect to master: {}: {}\n".format(e.__class__.__name__, str(e)))
            if not unregister:
                if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::register] launching new master\n")
                _execute_sync(lambda: ida_loader.load_plugin(master_plugin), ida_kernwin.MFF_WRITE)
        # ConnectionAbortedError
            if idarest_plugin_t.config['api_info']: idc.msg("[idarest_plugin_t::register] failed to connect to master: {}: {}\n".format(e.__class__.__name__, str(e)))

    @property
    def handler(self):
        return self.worker.httpd.RequestHandlerClass

    def add_route(self, route_pattern, f):
        self.handler.add_route(route_pattern, f)

    def remove_route(self, route_pattern):
        self.handler.remove_route(self.handler, route_pattern)

idarest_plugin_t.load_configuration()


class HTTPRequestError(BaseException):
    def __init__(self, msg, code):
        self.msg = msg
        self.code = code

class UnknownApiError(HTTPRequestError):
    pass

class HTTPRequestHandler(BaseHTTPRequestHandler):
    routes = {}
    docs = {}
    prefns = {}
    postfns = {}
    idarest_queue = [Queue() for x in range(11)]
    uid_iterator = itertools.count()

    @property
    def uid(self):
        # is probably guaranteed to be atomic https://stackoverflow.com/a/27062830/912236
        value = next(HTTPRequestHandler.uid_iterator)
        return value % 10

    def log_message(self, format, *args):
        return

    @staticmethod
    def _get_params(f):
        if idarest_plugin_t.config['api_debug']: idc.msg(inspect.getsource(ir.handler.routes['call'][1]))

    @staticmethod
    def set_result(uid, value):
        if idarest_plugin_t.config['api_verbose']: idc.msg("[set_result] {}: {}\n".format(uid, value))
        HTTPRequestHandler.idarest_queue[uid].put(value)
        return uid

    @staticmethod
    def synced_next(value):
        try:
            v = next(value)
            if idarest_plugin_t.config['api_debug']: idc.msg("next value: {}".format(v))
            HTTPRequestHandler.set_result(10, v)
        except StopIteration as e:
            HTTPRequestHandler.set_result(10, e)

    @staticmethod
    def wrapped_iter(value):
        while True:
            _execute_sync(lambda: HTTPRequestHandler.synced_next(value), ida_kernwin.MFF_WRITE)
            yield HTTPRequestHandler.get_result(10)
        yield None

    @staticmethod
    def get_result(uid):
        try:
            value = HTTPRequestHandler.idarest_queue[uid].get(timeout=2)
            if str(type(value)) == "<class 'generator'>":
                if idarest_plugin_t.config['api_debug']: idc.msg("[get_result] return wrapped_iter")
                return HTTPRequestHandler.wrapped_iter(value)
        except Empty:
            value = {'code': 400, 'msg': 'No response: timeout',
                    "traceback": traceback.format_exc()}
        except Exception as e:
            value = {'code': 400, 'msg': 'Unhandled Exception: ({}) {}'.format(type(e), str(e)),
                    "traceback": traceback.format_exc()}
        except:
            value = "timeout"
        if idarest_plugin_t.config['api_verbose']: idc.msg("[get_result] {}: {}\n".format(uid, value))
        return value

    @staticmethod
    def build_route_pattern(route):
        return re.compile("^{0}$".format(route))

    @staticmethod
    def route(route_str):
        def decorator(f):
            #  if idarest_plugin_t.config['api_debug']: print("[route] {}: {}".format(route_str, f.__doc__))
            route_path = idarest_plugin_t.config['api_prefix'] + '/' + route_str + '/?'
            route_pattern = HTTPRequestHandler.build_route_pattern(route_path)
            HTTPRequestHandler.routes[route_str] = (route_pattern, f)
            HTTPRequestHandler.docs[route_str] = f.__doc__
            # HTTPRequestHandler.params[route_str] = HTTPRequestHandler._get_params(f)
            return f
        return decorator

    @classmethod
    def add_route(cls, route_str, f):
        #  print("[HTTPRequestHandler::add_route] cls: {}".format(cls))
        route_path = idarest_plugin_t.config['api_prefix'] + '/' + route_str + '/?'
        route_pattern = cls.build_route_pattern(route_path)
        cls.routes[route_str] = (route_pattern, f)
        cls.docs[route_str] = f.__doc__
        return f

    def remove_route(self, route_str):
        if route_str in self.routes:
            self.routes.pop(route_str)
            return True
        return False


    @staticmethod
    def prefn(route_str):
        def decorator(f):
            HTTPRequestHandler.prefns.setdefault(route_str, []).append(f)
            return f
        return decorator

    @staticmethod
    def postfn(route_str):
        def decorator(f):
            HTTPRequestHandler.postfns.setdefault(route_str, []).append(f)
            return f
        return decorator

    def _get_route_match(self, path):
        for (key, (route_pattern, view_function)) in self.routes.items():
            if False: print("[debug] route_pattern:{}, path:{}".format(route_pattern, path))

            m = route_pattern.match(path)
            if m:
                return key, view_function
        return None

    def _get_route_prefn(self, key):
        try:
            return self.prefns[key]
        except:
            return []

    def _get_route_postfn(self, key):
        try:
            return self.postfns[key]
        except:
            return []

    def _write(self, s):
        self.wfile.write(_asBytes(s))

    def _write_chunk(self, r):
        l = len(r)
        self._write('{:X}\r\n{}\r\n'.format(l, r))


    def _serve_route(self, args, opts):
        # Python>urlparse.urlparse('http://localhost.local:1234/ida/api/v1.0/get_type.json?idiot=phil&a=b')
        # ('http', 'localhost.local:1234', '/ida/api/v1.0/get_type.json', '', 'idiot=phil&a=b', '')
        # Python>pp(read_everything(u))
        #  [   ('fragment', '',                            "<class 'str'>"),
        #      ('hostname', 'localhost.local',             "<class 'str'>"),
        #      ('netloc',   'localhost.local:1234',        "<class 'str'>"),
        #      ('params',   '',                            "<class 'str'>"),
        #      ('password', None,                          "<class 'NoneType'>"),
        #      ('path',     '/ida/api/v1.0/get_type.json', "<class 'str'>"),
        #      ('port',     '0x4d2',                       "<class 'int'>"),
        #      ('query',    'idiot=phil&a=b',              "<class 'str'>"),
        #      ('scheme',   'http',                        "<class 'str'>"),
        #      ('username', None,                          "<class 'NoneType'>")]
        path = urlparse.urlparse(self.path).path
        # Python>os.path.splitext(u.path)
        # ('/ida/api/v1.0/get_type', '.json')
        path, response_encoding = os.path.splitext(path)
        if response_encoding:
            response_encoding = response_encoding.lstrip('.')
        else:
            response_encoding = 'json'

        opts['response_encoding'] = response_encoding

        route_match = self._get_route_match(path)
        if route_match:
            key, view_function = route_match
            # these won't run in the main thread, so could cause issues if they try to interact with the idb
            for prefn in self._get_route_prefn(key):
                args = prefn(self, args)

            uid = self.uid
            #  while not HTTPRequestHandler.idarest_queue[uid].empty():
                #  print("[_serve_route] flushing queue...")
                #  HTTPRequestHandler.idarest_queue[uid].get()
                #

            def _exec():
                try:
                    HTTPRequestHandler.set_result(uid, view_function(self, args))
                except Exception as e:
                    HTTPRequestHandler.set_result(uid, e)

            # ida_kernwin.execute_sync(lambda: HTTPRequestHandler.set_result(uid, view_function(self, args)), ida_kernwin.MFF_WRITE)
            _execute_sync(_exec, ida_kernwin.MFF_WRITE)
            results = HTTPRequestHandler.get_result(uid)

            # these won't run in the main thread, so could cause issues if they try to interact with the idb
            for postfn in self._get_route_postfn(key):
                results = postfn(self, results)

            return results
        else:
            raise UnknownApiError('Route "{0}" has not been registered'.format(path), 404)

    def _serve_queue(self, q):
        response = {
            'code':  200,
            'msg':   'CONTROL',
            'control': {'queue': 'stop'},
            'data': False,
        }

        jsonp_callback = self._extract_callback()
        if jsonp_callback:
            content_type = 'application/javascript'
            response_fmt = jsonp_callback + '({0});'
        else:
            content_type = 'application/json'
            response_fmt = '{0}'

        self.send_response(200)
        self.send_header('Content-Type', content_type)
        self.send_origin_headers()
        self.send_header('Transfer-Encoding', 'chunked')
        self.end_headers()

        try:
            for i in itertools.count(start=0):
                self._write_chunk(response_fmt.format(json.dumps(response)))

                if idarest_plugin_t.config['api_debug']: idc.msg("[HTTPRequestHandler::_serve_queue] wrote: {}\n".format(response))
                data = q.get(timeout=idarest_plugin_t.config['api_queue_result_qget_timeout'])
                if data is None:
                    if idarest_plugin_t.config['api_debug']: idc.msg("[HTTPRequestHandler::_serve_queue] Queue returned None\n")
                    break
                response = {
                    'code': 200,
                    'msg':  'OK',
                    'oob':  {'count': i},
                    'data': data,
                }
                if isinstance(response['data'], dict):
                    if 'error' in response['data']:
                        response['msg'] = 'FAIL'
        except Empty:
            pass

        response = {
            'code':    200,
            'msg':     'CONTROL',
            'control': {'queue': 'stop'},
            'data':    False,
        }
        self._write_chunk(response_fmt.format(json.dumps(response)))
        self._write_chunk('')
        # self._write('0\r\n\r\n')

    def _serve_generator(self, it):
        iterable = True
        exception = True
        queue = True

        response = {
            'code' : 200,
            'msg'  : 'CONTROL',
            'data' : False,
            'control': {'iterable': 'start'},
        }

        jsonp_callback = self._extract_callback()
        if jsonp_callback:
            content_type = 'application/javascript'
            response_fmt = jsonp_callback + '({0});'
        else:
            content_type = 'application/json'
            response_fmt = '{0}'

        self.send_response(200)
        self.send_header('Content-Type', content_type)
        self.send_origin_headers()
        self.send_header('Transfer-Encoding', 'chunked')
        self.end_headers()

        try:
            for i in itertools.count(start=0):
                self._write_chunk(response_fmt.format(json.dumps(response)))

                if idarest_plugin_t.config['api_debug']: idc.msg("[HTTPRequestHandler::_serve_generator] wrote: {}".format(r))
                data = next(it)
                if data is None:
                    break
                if isinstance(data, StopIteration):
                    break
                response = {
                    'code' : 200,
                    'msg'  : 'OK',
                    'oob': {'count': i},
                    'data' : data,
                }
                if isinstance(response['data'], dict):
                    if 'error' in response['data']:
                        response['msg'] = 'FAIL'
        except StopIteration:
            pass

        response = {
            'code' : 200,
            'msg'  : 'CONTROL',
            'control': {'iterable': 'stop'},
            'data' : False,
        }
        self._write_chunk(response_fmt.format(json.dumps(response)))
        self._write_chunk('')
        #  self._write('0\r\n\r\n')

    def send_origin_headers(self):
        if self.headers.get('Origin', '') == 'null':
            self.send_header('Access-Control-Allow-Origin', self.headers.get('Origin'))
        self.send_header('Vary', 'Origin')

    def _serve(self, args):
        error = False
        try:
            opts = {}
            it = self._serve_route(args, opts)
            # it = sleeping_generator_test({}, {})
            if issubclass(it.__class__, Exception):
                response = {'code': 400, 'msg': 'returned exception {}: {}'.format(it.__class__, str(it))}
                error = True
                print('returned exception {}: {}'.format(it.__class__, str(it)))
            elif str(type(it)) == "<class 'generator'>":
                #  if idarest_plugin_t.config['api_debug']: print("[_serve] Generator!")
                return self._serve_generator(it)
            elif it.__class__.__name__ == "Queue":
                #  if idarest_plugin_t.config['api_debug']: print("[_serve] Queue!")
                return self._serve_queue(it)
            else:
                if isinstance(it, dict) and it.get('code', None):
                    response = it
                else:
                    response = {
                        'code' : 200,
                        'msg'  : 'OK',
                        'data' : it,
                    }

                    if isinstance(response['data'], dict):
                        if 'error' in response['data']:
                            response['code'] = 400
                            response['msg'] = response['data']['error']
                            if 'error_trace' in response['data']:
                                response['traceback'] = response['data']['error_trace']
                            response.pop('data')

        except UnknownApiError as e:
            self.send_error(e.code, e.msg)
            return
        except HTTPRequestError as e:
            response = {'code': e.code, 'msg' : e.msg}
        except ValueError as e:
            response = {'code': 400, 'msg': 'ValueError: ' + str(e)}
        except KeyError as e:
            response = {'code': 400, 'msg': 'KeyError: ' + str(e)}
        except StopIteration as e:
            response = {'code': 400, 'msg': 'StopIteration: ' + str(e)}
        except Exception as e:
            response = {
                    'code': 400, 
                    'msg': '{}: {}'.format(_classname(e), str(e)),
                    "traceback": traceback.format_exc()
            }

        

        encoder = json.dumps
        response_encoding = self._extract_response_encoding()
        if response_encoding == 'pickle':
            encoder = pickle.dumps

        jsonp_callback = self._extract_callback()
        if jsonp_callback:
            content_type = 'application/javascript'
            response_fmt = jsonp_callback + '({0});'
        else:
            content_type = 'application/json'
            response_fmt = '{0}'

        self.send_response(200)
        self.send_header('Content-Type', content_type)
        self.send_origin_headers()
        self.end_headers()

        self._write(response_fmt.format(encoder(response)))
        return

    def _extract_post_map(self):
        content_type, _t = cgi.parse_header(self.headers.get('content-type'))
        #  if content_type != 'application/json':
            #  raise HTTPRequestError(
                    #  'Bad content-type, use application/json',
                    #  400)
        length = int(self.headers.get('content-length'))
        try:
            return json.loads(self.rfile.read(length))
        except ValueError as e:
            raise HTTPRequestError(
                    'Bad or malformed json content',
                    400)

    def _extract_query_map(self):
        query = urlparse.urlparse(self.path).query
        qd = urlparse.parse_qs(query)
        args = {}
        for k, v in qd.items():
            if len(v) != 1:
                raise HTTPRequestError(
                    "Query param specified multiple times : " + k,
                    400)
            args[k.lower()] = v[0]
        return args

    def _extract_callback(self):
        try:
            args = self._extract_query_map()
            return args['callback']
        except:
            return ''

    def _extract_response_encoding(self):
        try:
            path = urlparse.urlparse(self.path).path
            # Python>os.path.splitext(u.path)
            # ('/ida/api/v1.0/get_type', '.json')
            path, response_encoding = os.path.splitext(path)
            if response_encoding:
                response_encoding = response_encoding.lstrip('.')
            else:
                response_encoding = 'json'

            return response_encoding
        except:
            return ''

    def process_headers(self):
        self.header_dict = dict()
        for a, c in self.headers.items():
            if a and c:
                self.header_dict[a.strip()] = c.strip()

    def do_POST(self):
        self.process_headers()
        try:
            args = self._extract_post_map()
        except TypeError as e:
            # thrown on no content, just continue on
            args = '{}'
        except HTTPRequestError as e:
            self.send_error(e.code, e.msg)
            return
        self._serve(args)

    def do_GET(self):
        self.process_headers()
        try:
            HTTPRequestHandler._request = self
            args = self._extract_query_map()
        except HTTPRequestError as e:
            self.send_error(e.code, e.msg)
            return
        try:
            self._serve(args)
        except ConnectionAbortedError as e:
            if idarest_plugin_t.config['api_debug']: idc.msg("[HTTPRequestHandler::do_GET] {}: {}\n".format(e.__class__.__name__, str(e)))

    @staticmethod
    def delayed_call(f):
        """
        call function (delayed)

        allows us to start a function that will fill a Queue, while still being
        able to return that Queue before the function blocks
        """
        with ida_kernwin.disabled_script_timeout_t():

            def delayed_exec():
                f()

            # delayed_exec_timer = QtCore.QTimer()
            delayed_exec_timer.singleShot(0, delayed_exec)

    @staticmethod
    def fake_cli(text):
        """
        fake text input into cli

        TODO: copy ida stdout perhaps?  see ipyida's Zmq console tee.
        """
        with ida_kernwin.disabled_script_timeout_t():

            # We'll now have to schedule a call to the standard
            # 'execute' action. We can't call it right away, because
            # the "Output window" doesn't have focus, and thus
            # the action will fail to execute since it requires
            # the "Output window" as context.

            def delayed_exec(*args):
                output_window_title = "Output window"
                tw = ida_kernwin.find_widget(output_window_title)
                if not tw:
                    raise Exception("Couldn't find widget '%s'" % output_window_title)

                # convert from a SWiG 'TWidget*' facade,
                # into an object that PyQt will understand
                w = ida_kernwin.PluginForm.TWidgetToPyQtWidget(tw)

                line_edit = w.findChild(QtWidgets.QLineEdit)
                if not line_edit:
                    raise Exception("Couldn't find input")
                line_edit.setFocus() # ensure it has focus
                QtWidgets.QApplication.instance().processEvents() # and that it received the focus event

                # inject text into widget
                line_edit.setText(text)

                # and execute the standard 'execute' action
                ida_kernwin.process_ui_action("cli:Execute")

            # delayed_exec_timer = QtCore.QTimer()
            delayed_exec_timer.singleShot(0, delayed_exec)


"""
API handlers for IDA

"""
def check_ea(f):
    def wrapper(self, args):
        if 'ea' in args:
            try:
                ea = int(args['ea'], 16)
            except ValueError:
                raise IDARequestError(
                        'ea parameter malformed - must be 0xABCD', 400)
            if ea > idc.MaxEA():
                raise IDARequestError(
                        'ea out of range - MaxEA is 0x%x' % idc.MaxEA(), 400)
            args['ea'] = ea
        return f(self, args)
    return wrapper

def check_color(f):
    def wrapper(self, args):
        if 'color' in args:
            color = args['color']
            try:
                color = color.lower().lstrip('#').rstrip('h')
                if color.startswith('0x'):
                    color = color[2:]
                # IDA Color is BBGGRR, we need to convert from RRGGBB
                color = color[-2:] + color[2:4] + color[:2]
                color = int(color, 16)
            except:
                raise IDARequestError(
                        'color parameter malformed - must be RRGGBB form', 400)
            args['color'] = color
        return f(self, args)
    return wrapper

# this doesn't presently work, as it ends up being executed in IDA's main
# working thread. but let's leave it in so we can still use the decorator
# and maybe it will be fixed.
def require_params(*params):
    def decorator(f):
        def require_params_wrapper(self, args):
            for x in params:
                if ':' not in x and ' ' not in x and x not in args:
                    raise IDARequestError('missing parameter {0}'.format(x), 400)
            return f(self, args)
        require_params_wrapper.__doc__ = f.__doc__
        require_params_wrapper._params = getattr(f, '_params', [])
        require_params_wrapper._params.append(params)
        return require_params_wrapper
    return decorator

class IDARequestError(HTTPRequestError):
    pass

class EvalInterpreter(InteractiveInterpreter):
    def __init__(self, locals=None):
        super(EvalInterpreter, self).__init__(locals)
        self.buffer  = []
        self.stdout_list  = []
        self.stderr_list  = []

    def write(self, data):
        self.buffer.append(data)

    def eval(self, cmd):
        is_help = cmd.startswith('help(')
        result = None
        r = None
        with BorrowStdOut(self.stdout_list, self.stderr_list, is_help):
            r = self.runsource(cmd)
            if isinstance(self.stderr_list, list): 
                if self.stderr_list:
                    start = False
                    result = []
                    for line in self.stderr_list:
                        if 'File "<input>"' in line:
                            start = True
                        if start:
                            result.append(line)
                    if result:
                        self.stderr_list = [self.stderr_list[0]] + result
                        return {
                                'code': 400,
                                'msg': ''.join(self.stderr_list),
                        }

                result = ''.join(self.stdout_list + self.buffer)
            
        if r != False:
            raise SyntaxError(cmd)
        if isinstance(self.stdout_list, list): 
            return result

class EvalInterpreterQueue(EvalInterpreter):
    def __init__(self, locals, queue):
        super(EvalInterpreterQueue, self).__init__(locals)
        self.buffer  = queue
        self.stdout_list  = self.buffer
        self.stderr_list  = self.buffer

        self.buffer.append = self.buffer.put

class IDARequestHandler(HTTPRequestHandler):
    @staticmethod
    def _hex(v):
        return hex(v).rstrip('L')

    @staticmethod
    def _from_hex(v):
        if re.match(r'0x[0-9a-fA-F]+$', v):
            return int(v, 16)
        if re.match(r'\d+$', v):
            return int(v, 10)
        if re.match(r'".*"$', v):
            return v[1:-1]
        if re.match(r"'.*'$", v):
            return v[1:-1]
        return v

    def paren_split(s, delim):
        """ split string, respecting quotes """
        result = paren_multisplit(s, delim, '\'"', '\'"')
        if idarest_plugin_t.config['api_debug']: print("[paren_split] in: {}".format(s))
        if idarest_plugin_t.config['api_debug']: print("[paren_split] out: {}".format(result))
        return result


    @staticmethod
    def _dotted(key):
        pieces = key.split('.')
        return pieces

    @staticmethod
    def _ensure_path(_dict, path):
        if not path:
            if idarest_plugin_t.config['api_debug']: idc.msg("[_ensure_path] empty path\n")
            return None
        for piece in path:
            try:
                if piece in _dict:
                    _dict = _dict[piece]
            except TypeError:
                if hasattr(_dict, piece):
                    _dict = getattr(_dict, piece)
                else:
                    return None
        return _dict


    @staticmethod
    def _getplus(key):
        return getglobal(key)
        #        _globals = superglobals()
        #
        #        if isinstance(key, list):
        #            path = key
        #        else:
        #            path = IDARequestHandler._dotted(key)
        #
        #        base = IDARequestHandler._ensure_path(_globals, path)
        #        return base

    @staticmethod
    def error(e):
        if issubclass(e.__class__, Exception):
            _class = _classname(e)
            _message = str(e)
            result = {
                    "error": "{}: {}".format(_class, _message),
                    "error_trace": traceback.format_exc(),
            }
        else:
            result = {
                    "error": e,
            }
        return result


    @HTTPRequestHandler.route('info')
    def info(self, args):
        # No args, Return everything we can meta-wise about the ida session
        # file crcs
        result = {
                'md5' : _asStringRaw(idc.GetInputMD5()),
                'idb_path' : idc.GetIdbPath(),
                'file_path' : idc.GetInputFilePath(),
                'ida_dir' : idc.GetIdaDirectory(),
                'min_ea' : self._hex(idc.MinEA()),
                'max_ea' : self._hex(idc.MaxEA()),
                'segments' : self.segments({})['segments'],
                # idaapi.cvar.inf
                'procname' : idc.GetLongPrm(idc.INF_PROCNAME),
            }
        return result

    @HTTPRequestHandler.route('query')
    def query(self, args):
        # multiple modes
        # with address return everything about that address
        # with name, return everything about that name
        if 'ea' in args:
            idc.jumpto(args['ea'])
            return idc.here()
        if 'name' in args:
            return idc.get_name_ea_simple(args['name'])



    @HTTPRequestHandler.route('cursor')
    @check_ea
    def cursor(self, args):
        # XXX - Doesn't work
        #if 'window' in args:
        #    tform = idaapi.find_tform(args['window'])
        #    if tform:
        #        idaapi.switchto_tform(tform, 1)
        #    else:
        #        raise IDARequestError(
        #            'invalid window - {0}'.format(args['window']), 400)
        result = {}
        if idarest_plugin_t.config['api_debug']: idc.msg(args)
        if 'ea' in args:
            success = idc.jumpto(args['ea'])
            result['moved'] = success
        else:
            result['error'] = "missing argument: ea"
        result['ea'] = self._hex(idc.here())
        return result

    @HTTPRequestHandler.route('color')
    @check_color
    @check_ea
    @require_params('ea')
    def color(self, args):
        ea = args['ea']
        if 'color' in args:
            color = args['color']
            def f():
                idc.SetColor(ea, idc.CIC_ITEM, color)
            # idaapi.execute_sync(f, idaapi.MFF_WRITE)
            f()
            idc.Refresh()
            return {}
        else:
            return {'color' : str(GetColor(ea, idc.CIC_ITEM))}


    @HTTPRequestHandler.route('get_type')
    @require_params('type', 'comma seperated list of types')
    def get_type(self, args):
        """
        get type definition(s)

        :param type: comma seperated list of types
        :return

        $ echo -e 'GET /ida/api/v1.0/get_type?type=Vehicle,Entity HTTP/1.1\n\n' | nc localhost 2245
        {
            'code': 200,
            'msg': 'OK',
            'data': [{
                'name': 'Vehicle',
                'msg': 'OK',
                'data': '/* 11472 */\ntypedef int Vehicle;\n\n'
            }, {
                'name': 'Entity',
                'msg': 'OK',
                'data': '/* 1124 */\ntypedef int Entity;\n\n'
            }]
        }
        """
        def get_tinfo_by_parse(name):
            result = idc.parse_decl(name, idc.PT_SILENT)
            if result is None:
                return
            _, tp, fld = result
            tinfo = idaapi.tinfo_t()
            tinfo.deserialize(idaapi.cvar.idati, tp, fld, None)
            return tinfo

        def get_tinfo_magic(name):
            tif = get_tinfo_by_parse(name + '*')
            if not tif:
                return
            return tif.get_pointed_object()


        def my_print_decls(name, flags=0):
            names = name if isinstance(name, list) else [name]
            ordinals = []
            for name in names:
                ti = get_tinfo_magic(name)
                if ti:
                    ordinal = ti.get_ordinal()
                    if ordinal:
                        ordinals.append(ordinal)
                        continue
                if idarest_plugin_t.config['api_debug']: print("[warn] couldn't get ordinal for type '{}'".format(name))
            if not ordinals:
                if idarest_plugin_t.config['api_debug']: print("[warn] couldn't get ordinals for types '{}'".format(name))
                return ''
            else:
                if idarest_plugin_t.config['api_debug']: print("[info] ordinals: {}".format(ordinals))

            result = ''
            if ordinals:
                result = idc.print_decls(','.join([str(x) for x in ordinals if x > -1]), flags)
            if idarest_plugin_t.config['api_debug']: print(result)
            return result

        types = IDARequestHandler.paren_split(args['type'], ',')
        #  flags = idc.PDF_INCL_DEPS | idc.PDF_DEF_FWD
        flags = 3
        if 'flags' in args:
            flags = int(args['flags'], 0)
        if idarest_plugin_t.config['api_debug']: print("request for type definitions: {}".format(types))
        result = []
        if types:
            for t in set(types):
                if t != 'void':
                    et = my_print_decls(t, flags=flags)
                    response = {
                        'name' : t,
                        'msg'  : 'OK',
                        'data' : et,
                    }
                    if not et:
                        if idarest_plugin_t.config['api_debug']: print("**** Type Error ****\n{}\n".format(t))
                        response['msg'] = 'FAIL'
                    result.append(response)
        return result


    def _get_segment_info(self, s):
        return {
            'name' : idaapi.get_true_segm_name(s),
            'ida_name' : idaapi.get_segm_name(s),
            'start' : self._hex(s.startEA),
            'end' : self._hex(s.end_ea),
            'size' : self._hex(s.size())
        }

    @HTTPRequestHandler.route('segments')
    @check_ea
    def segments(self, args):
        if 'ea' in args:
            s = idaapi.getseg(args['ea'])
            if not s:
                raise IDARequestError('Invalid address', 400)
            return {'segment': self._get_segment_info(s)}
        else:
            m = {'segments': []}
            for i in range(idaapi.get_segm_qty()):
                s = idaapi.getnseg(i)
                m['segments'].append(self._get_segment_info(s))
            return m

    @HTTPRequestHandler.route('get')
    @require_params('var', 'name of variable to retrieve')
    def get_var(self, args):
        """get global variable

        :param var: var name

        """
        if idarest_plugin_t.config['api_debug']: idc.msg("[get_var]\n")
        try:
            if not 'var' in args:
                return IDARequestHandler.error('missing parameter \'var\'')
            var = args.pop('var')
            value = IDARequestHandler._getplus(var)
            if idarest_plugin_t.config['api_debug']:
                print({
                    'var': var,
                    'value': value,
                })
            if value is None:
                return IDARequestHandler.error(NameError("name '{}' is not defined or is None".format(var)))
            result = value
            if idarest_plugin_t.config['api_debug']: idc.msg('result: {}\n'.format(result))
            return result

        except Exception as e:
            return IDARequestHandler.error(e)
        except:
            return IDARequestHandler.error("Unknown Exception")

    @HTTPRequestHandler.route('call')
    @require_params('cmd',             'name of callable')
    @require_params('args:optional',   'comma seperated list of positional arguments')
    @require_params('kwargs:optional', 'object of keyword arguments')
    def call(self, args):
        """run callable and return result

        :param cmd: callable
        :param args: [optional] comma seperated list of positional arguments
        :param *: [optional] keyword arguments

        $ wget 'http://127.0.0.1:2000/ida/api/v1.0/call?cmd=_type=idc.GetType(0x1412E9E98)&return=_type' -O - -q
        {
            'code': 200,
            'msg': 'OK',
            'data': 'void __fastcall(uint8_t *buffer, uint32_t data, uint32_t bits, int32_t offset)'
        }

        """
        if idarest_plugin_t.config['api_debug']: idc.msg("[call]\n")
        try:
            if not 'cmd' in args:
                return IDARequestHandler.error('missing parameter \'cmd\'')
            cmd = args.pop('cmd')
            _args = []
            _kwargs = {}
            if 'args' in args:
                _args = IDARequestHandler.paren_split(args.pop('args'), ',')
                _args = [ IDARequestHandler._from_hex(x) for x in _args ]
                if idarest_plugin_t.config['api_debug']: idc.msg('_args: {}\n'.format(_args))
            for k, v in args.items():
                _kwargs[k] = IDARequestHandler._from_hex(v)
                if idarest_plugin_t.config['api_debug']: idc.msg('_kwarg: {}: {}\n'.format(k, v))
            fn = IDARequestHandler._getplus(cmd)
            if idarest_plugin_t.config['api_debug']:
                print({
                    'cmd': cmd,
                    'fn': fn,
                    'args': _args,
                    'kwargs': _kwargs,
                })
            if fn is None:
                return IDARequestHandler.error(NameError("name '{}' is not defined".format(cmd)))
            if not callable(fn):
                return IDARequestHandler.error(NameError("name '{}' is not callable".format(cmd)))
            result = fn(*_args, **_kwargs)
            if idarest_plugin_t.config['api_debug']: idc.msg('result: {}\n'.format(result))
            return result

        except Exception as e:
            return IDARequestHandler.error(e)
        except:
            return IDARequestHandler.error("Unknown Exception")

    @HTTPRequestHandler.route('eval')
    @HTTPRequestHandler.route('exec')
    @require_params('cmd',             'string to evaluate')
    @require_params('return:optional', 'variable to return, else return stdout')
    def eval(self, args):
        """evaluate expression via python exec() and return stdout/stderr (or variable)

        :param cmd: string to evaluate
        :param return: [optional] name of variable to return

        $ wget 'http://127.0.0.1:2000/ida/api/v1.0/eval?cmd=idc.GetType(0x1412E9E98)' -O - -q
        {
            'code': 200,
            'msg': 'OK',
            'data': 'void __fastcall(uint8_t *buffer, uint32_t data, uint32_t bits, int32_t offset)'
        }

        """
        if idarest_plugin_t.config['api_debug']: idc.msg("Hello Eval\n")
        try:
            if not 'cmd' in args:
                return IDARequestHandler.error('missing parameter \'cmd\'')
            cmd = args['cmd']
            if idarest_plugin_t.config['api_debug']: idc.msg('cmd: {}\n'.format(cmd))
            i = EvalInterpreter(superglobals())
            _result = i.eval(cmd)

            if idarest_plugin_t.config['api_debug']: idc.msg('result: {}\n'.format(_result))
            if 'return' in args:
                return getglobal(args['return'], None)
            else:
                return _result
        except Exception as e:
            return IDARequestHandler.error(e)

    @HTTPRequestHandler.route('evalq')
    @require_params('cmd',             'string to evaluate')
    def evalq(self, args):
        """evaluate expression via python exec() and return stdout via chunked queue

        :param cmd: string to evaluate

        $ wget 'http://127.0.0.1:2000/ida/api/v1.0/evalq?cmd=idc.GetType(0x1412E9E98)' -O - -q
        {
            'code': 200,
            'msg': 'OK',
            'data': 'void __fastcall(uint8_t *buffer, uint32_t data, uint32_t bits, int32_t offset)'
        }

        """
        if idarest_plugin_t.config['api_debug']: idc.msg("Hello EvalQ\n")
        try:
            if not 'cmd' in args:
                return IDARequestHandler.error('missing parameter \'cmd\'')
            cmd = args['cmd']
            if idarest_plugin_t.config['api_debug']: idc.msg('cmd: {}\n'.format(cmd))
            q = Queue(maxsize=1)
            i = EvalInterpreterQueue(superglobals(), q)
            def delayed():
                i.eval(cmd)
            HTTPRequestHandler.delayed_call(delayed)

            #  _result = i.eval(cmd)

            return q
        except Exception as e:
            return IDARequestHandler.error(e)

    @HTTPRequestHandler.route('cli')
    @require_params('cmd', 'string to evaluate')
    def cli(self, args):
        """fake input to CLI

        :param cmd: string to evaluate
        """
        if idarest_plugin_t.config['api_debug']: idc.msg("Hello CLI\n")
        try:
            if not 'cmd' in args:
                return IDARequestHandler.error('missing parameter \'cmd\'')
            cmd = args['cmd']
            if idarest_plugin_t.config['api_debug']: idc.msg('cmd: {}\n'.format(cmd))
            HTTPRequestHandler.fake_cli(cmd)
        except Exception as e:
            return IDARequestHandler.error(e)

class _IDARestStdOut:
    """
    Dummy file-like class that receives stout and stderr
    """
    buffer = ''
    def __init__(self):
        pass
        #  self.buffer = ''

    def write(self, text):
        # NB: in case 'text' is Unicode, msg() will decode it
        # and call msg() to print it
        self.buffer += text

    def flush(self):
        pass

    def isatty(self):
        return False

"""
Threaded HTTP Server and Worker

Use a worker thread to manage the server so that we can run inside of
IDA Pro without blocking execution.

"""
class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
    allow_reuse_address = True


class Worker(threading.Thread):
    def __init__(self, host=idarest_plugin_t.config['api_bind_ip'], port=idarest_plugin_t.config['api_port']):
        threading.Thread.__init__(self)
        self.httpd = ThreadedHTTPServer((host, port), IDARequestHandler)
        self.host = host
        self.port = port

    def run(self):
        self.httpd.serve_forever()

    def stop(self):
        self.httpd.shutdown()
        self.httpd.server_close()
        if idarest_plugin_t.config['api_info']: idc.msg("httpd shutdown...\n")
        if idarest_plugin_t.config['api_info']: idc.msg("httpd server_close...\n")

class Timer(threading.Thread):
    def __init__(self, host, port):
        super(Timer, self).__init__()
        self._stop_event = threading.Event()
        self.host = host
        self.port = port

    def run(self):
        if idarest_plugin_t.config['api_info']: print("[idarest::Timer::run] started\n")
        while True:
            result = idarest_plugin_t.register(self.host, self.port)
            if self._stop_event.wait(60.0):
                break
        if idarest_plugin_t.config['api_info']: print("[idarest::Timer::run] stopped")

        #  if not self.running:
            #  self.running = True
            #  while self.running:
                #  time.sleep(60.0 - ((time.time() - self.starttime) % 60.0))
                #  if idarest_plugin_t.config['api_debug']: print(Handler.get_json(Handler.hosts, {'ping': time.time()}))
            #  if idarest_plugin_t.config['api_info']: print("[idarest::Timer::run] stopped")

    def stop(self):
        if self.is_alive():
            if self.stopped():
                if idarest_plugin_t.config['api_info']: print("[idarest::Timer::stop] already stopping...")
            else:
                if idarest_plugin_t.config['api_info']: print("[idarest::Timer::stop] stopping...")
                self._stop_event.set()
        else:
            if idarest_plugin_t.config['api_info']: print("[idarest::Timer::stop] not running")

    def stopped(self):
        return self._stop_event.is_set()

# This has been redesigned **not** to require running as a plugin, but if you
# really insist on trying, it does work.  To dynamically add routes if loaded
# as a plugin, you will have to do something like this:
#
# >>> ir = sys.modules['__plugins__idarest'].instance
# >>> ir.add_route(...)
#
# Or the startup.py script should create an instance of `ir`
def PLUGIN_ENTRY():
    globals()['instance'] = idarest_main() # idarest_plugin_t()
    return globals()['instance']

def idarest_main(*args):
    # pretend to be a plugin
    if idarest_main.instance is None:
        idarest_main.instance = ir = idarest_plugin_t()
        idarest_main.instances.append(ir)
        ir.init()
        #  globals()['ir'] = ir
        #  def cleanup():
            #  ir.term()
        #  atexit.register(cleanup)
    else:
        idc.msg('idarest_main.instance was not None!!!!!\n')
        return idarest_main.instance

    ### example dnamic routes
    def name_generator(self, args):
        """return all extant names and addresses via generator"""
        m = {'names' : []}
        for n in idautils.Names():
            yield {n[1]: self._hex(n[0])}

    def names(self, args):
        """return all extant names and addresses"""
        m = {'names' : []}
        for n in idautils.Names():
            m['names'].append([self._hex(n[0]), n[1]])
        return m

    def functions(self, args):
        """return all extant names and addresses"""
        m = {'functions' : []}
        for n in idautils.Functions():
            m['functions'].append([n, idc.get_name(n)])
        return m


    def sleeping_generator_test(self, args):
        for r in range(5):
            if idarest_plugin_t.config['api_debug']: idc.msg("[sleeping_generator_test] {}".format(r))
            yield r
            time.sleep(1)

    def queue_test(self, args):
        _q = Queue()
        for r in range(10):
            _q.put(r)
        return _q

    def sleeping_queue_test(self, args):
        q = Queue()
        def test():
            for r in range(60):
                q.put(r)
                time.sleep(1)
        HTTPRequestHandler.delayed_call(test)
        return q

    @ir_static_vars(q=Queue())
    def tap(self, args):
        if BorrowStdOut.is_default_stdout(sys.stdout):
            sys.stdout = BorrowStdOut.IDARestStdOutTee(sys.stdout, BorrowStdOut.IDARestStdOutQueue(tap.q))
            # sys.stderr = BorrowStdOut.IDARestStdOutTee(sys.stderr, BorrowStdOut.IDARestStdOutQueue(tap.q))
        else:
            pass
            # return IDARequestHandler.error('stdout is already tapped')
        return tap.q


    ir.add_route('sleep', sleeping_generator_test)
    ir.add_route('name_generator', name_generator)
    ir.add_route('names', names)
    ir.add_route('functions', functions)
    ir.add_route('ea', lambda o, a: idc.here())
    ir.add_route('echo', lambda o, a: {'args': a})
    ir.add_route('name', lambda o, a: idc.get_name(eax(a['ea'])))
    ir.add_route('q', queue_test)
    ir.add_route('q2', sleeping_queue_test)
    ir.add_route('tap', tap)
    ir.add_route('make_patch', lambda o, a: "\n".join(make_native_patchfile(a['ea'])))
    ir.add_route('getglobal', lambda o, a: getglobal(a['name'], None))
    ir.add_route('setglobal', lambda o, a: setglobal(a['name'], a['value']))
    ir.add_route('patches', lambda o, a: get_patches_idarest())
    ### end example routes


    ### some stuff I use that won't work for you
    def relist_iter(self, args):
        return iter_retrace_list(q, once=1)

    def relist_queue(self, args):
        setglobal("_q", Queue())
        HTTPRequestHandler.fake_cli("retrace_list(q, output=_q)")
        return getglobal("_q")

    ir.add_route('list1', relist_iter)
    ir.add_route('list2', relist_queue)
    ### end "my stuff"

    return ir


if not hasattr(idarest_main, 'instance'):
    idarest_main.instance = None
    idarest_main.instances = []

if idarest_plugin_t.config['api_verbose']:
    print("[idarest]: __name__: ", __name__)
    print("[idarest]: __file__: ", __file__)


_load_method = None
if __name__ == "__main__":
    # loaded directly
    _load_method = 'direct'
elif __name__.startswith("__plugins__"):
    _load_method = 'plugin'
    # loaded as a plugin
elif __name__ == "idarest" or __name__ == "_idarest":
    _load_method = 'module'
elif __name__ == "idarest.idarest":
    _load_method = 'package'
else:
    # unknown load method (filename could be changed?)
    _load_method = 'unknown'
    print("[idarest]: unknown load method '{}'".format(__name__))

# find existing instance of idarest (unless we're loading as an ida plugin)
def get_ir():
    if idarest_plugin_t.config['api_debug']: idc.msg("[get_ir] attempting to find existing idarest instance...\n")
    # check if idarest is loaded as a plugin
    ir = None
    if idarest_plugin_t.config['api_verbose']: idc.msg('[get_ir] sys.modules.__plugins__\n')
    for k in [x for x in sys.modules if x.startswith('__plugins__')]:
        if idarest_plugin_t.config['api_verbose']: idc.msg('[get_ir]     ' + k + '\n')
    ir = ir or getglobal('sys.modules.__plugins__idarest_plugin.instance', None)
    if ir:
        if idarest_plugin_t.config['api_info']: idc.msg("[get_ir] got reference to idarest from sys.modules.__plugins__idarest_plugin.instance\n")
    else:
        ir = ir or getglobal('sys.modules.__plugins__idarest.instance', None)
        if ir:
            if idarest_plugin_t.config['api_info']: idc.msg("[get_ir] got reference to idarest from sys.modules.__plugins__idarest.instance\n")
        # check if idarest is loaded as a module
        else:
            ir = ir or getglobal('sys.modules.idarest.instance', None)
            if ir:
                if idarest_plugin_t.config['api_info']: idc.msg("[get_ir] got reference to idarest from sys.modules.idarest.instance\n")
            else:
                # check if idarest is loaded in global context
                ir = ir or getglobal('idarest_main.instance', None)
                if ir:
                    if idarest_plugin_t.config['api_info']: idc.msg("[get_ir] got reference to idarest from idarest_main.instance\n") 
                    # else start a new idarest instance
                    if ir and _load_method == "direct":
                        if idarest_plugin_t.config['api_info']: idc.msg("[get_ir] will (re)start idarest\n")
                        ir.term()
                        ir = None

    if not ir or not ir.is_alive():
        if idarest_plugin_t.config['api_info']: idc.msg("[get_ir] starting\n")
        ir = idarest_main()

    return ir

def get_ir_plugin():
    if '__plugins__idarest_plugin' in sys.modules:
        return sys.modules['__plugins__idarest_plugin']

def find_plugin(pattern):
    for name in sys.modules.copy():
        if name.startswith('__plugins__') and re.match(pattern, name[11:]):
            yield name[11:]

def get_plugin(name):
    return sys.modules['__plugins__' + name]

def unload_module(pattern):
    for x in [x for x in sys.modules.keys() if re.match(pattern, x)]:
        print("unloading {}".format(x))
        del sys.modules[x]

def unload_plugin(pattern, reload=False):
    for name in find_plugin(pattern):
        l = ida_loader.load_plugin(get_plugin(name).__file__)
        ida_loader.run_plugin(l, 0)
        unload_module(pattern)
        if reload:
            plugin = get_ir_plugin()
            if plugin:
                l = ida_loader.load_plugin(plugin.__file__)

def reload_idarest():
    plugin = get_ir_plugin()
    if plugin:
        l = ida_loader.load_plugin(plugin.__file__)
        ida_loader.run_plugin(l, 0)
    unload_module('idarest')
    if plugin:
        l = ida_loader.load_plugin(plugin.__file__)
    else:
        print("idarest was not loaded as a plugin, cannot reload")

def test_eval(cmd):
    ir = getglobal('ir', None)
    if ir:
        ir.term()
    # unload_module('idarest')
    # _from('idarest.idarest import *')
    try:
        ida_idaapi.require('idarest.idarest')
        ida_idaapi.require('idarest.idarest_mixins')
        from idarest.idarest import get_ir
    except ModuleNotFoundError:
        ida_idaapi.require('idarest')
        ida_idaapi.require('idarest_mixins')
        from idarest import get_ir
    setglobal('ir', get_ir())
    e = EvalInterpreter(superglobals())
    r = e.eval(cmd)
    return r


repository: secrary/IDA-scripts
file: idenLib/idenLib.py
description: IDAPro scripts/plugins
comment: idenLib - Library Function Identification
created_at: 2018-01-11 07:05:10+00:00
pushed_at: 2019-02-26 13:58:04+00:00
url:
https://github.com/secrary/IDA-scripts/blob/0a646594225fa784021d0cc2aa79741abc6dd04c/idenLib/idenLib.py

content:
# -------------------------------------------------------------------------------
#
# Copyright (c) 2019
# Lasha Khasaia @_qaz_qaz
#
# -------------------------------------------------------------------------------


from __future__ import print_function
import ida_name
import ida_bytes
import ida_funcs
import idaapi
import ida_idaapi
import ida_diskio
import ida_ua
import idautils
import ida_name
import ida_nalt
import idc
import os
import zstd
import capstone
import cPickle as pickle

MIN_FUNC_SIZE = 0x20
MAX_FUNC_SIZE = 0x100
PLUGIN_NAME = "idenLib"
PLUGIN_VERSION = "v0.4"
IDENLIB_CONTACT = "Contact: Lasha Khasaia (@_qaz_qaz)"

local_appdata = os.getenv('LOCALAPPDATA')

# __EA64__ is set if IDA is running in 64-bit mode
__EA64__ = ida_idaapi.BADADDR == 0xFFFFFFFFFFFFFFFFL

if __EA64__ :
    CAPSTONE_MODE = capstone.CS_MODE_64
    SIG_EXT = ".sig64"
    idenLibCache = local_appdata + os.sep + PLUGIN_NAME + os.sep + "idenLibCache64"
    idenLibCacheMain = local_appdata + os.sep + PLUGIN_NAME + os.sep + "idenLibCacheMain64"
else:
    CAPSTONE_MODE = capstone.CS_MODE_32
    SIG_EXT = ".sig"
    idenLibCache = local_appdata + os.sep + PLUGIN_NAME + os.sep + "idenLibCache"
    idenLibCacheMain = local_appdata + os.sep + PLUGIN_NAME + os.sep + "idenLibCacheMain"
idenLib_appdata = local_appdata + os.sep + PLUGIN_NAME

func_sigs = {}
mainSigs = {}

ida_dir = ida_diskio.idadir("")
symEx_dir = ida_dir + os.sep + "SymEx"

def getNames():
    for ea, name in idautils.Names():
        yield name

def getFiles(path):  
    for file in os.listdir(path):
        if os.path.isfile(os.path.join(path, file)):
            yield path + "\\" + file

# return (start_ea, size)
def getFuncRanges():
    funcs_addr = []
    start = 0
    next_func =  ida_funcs.get_next_func(start)
    while next_func:
        size = next_func.size()
        if (size) < MIN_FUNC_SIZE:
           next_func = ida_funcs.get_next_func(next_func.start_ea)
           continue
        elif size > MAX_FUNC_SIZE:
            size = MAX_FUNC_SIZE
        yield (next_func.start_ea, size)
        funcs_addr.append(next_func.start_ea - start)
        next_func = ida_funcs.get_next_func(next_func.start_ea)

def getOpcodes(addr, size):
    md = capstone.Cs(capstone.CS_ARCH_X86, CAPSTONE_MODE)
    md.detail = True
    instr_bytes = ida_bytes.get_bytes(addr, size)
    opcodes_buf = b''
    for i in md.disasm(instr_bytes, size):
        # get last opcode
        if (i.opcode[3] != 0):
            opcodes_buf += "%02x" % (i.opcode[3])
        elif (i.opcode[2] != 0):
            opcodes_buf += "%02x" % (i.opcode[2])
        elif(i.opcode[1] != 0):
            opcodes_buf += "%02x" % (i.opcode[1])
        else:
            opcodes_buf += "%02x" % (i.opcode[0])
    return opcodes_buf

def idenLibProcessSignatures():
    global func_sigs
    global mainSigs
    for file in getFiles(symEx_dir):
        if not file.endswith(SIG_EXT):
            continue
        with open(file, 'rb') as ifile:
            sig = ifile.read()
            sig = zstd.decompress(sig).strip()
            sig = sig.split(b"\n")
            for line in sig:
                sig_opcodes, name = line.split(" ")
                if '_' in sig_opcodes: # "main" signatures
                    opcodeMain, mainIndexes = sig_opcodes.split('_')
                    fromFunc, fromBase = mainIndexes.split("!")
                    mainSigs[opcodeMain] = (name.strip(), int(fromFunc), int(fromBase))
                elif '+' in sig_opcodes:
                    opcodes, strBranches = sig_opcodes.split('+')
                    nBranches = int(strBranches)
                    func_sigs[opcodes.strip()] = (name.strip(), nBranches)
                else:
                    func_sigs[sig_opcodes.strip()] = (name.strip(), 0)
    if not os.path.isdir(idenLib_appdata):
        os.mkdir(idenLib_appdata)
    pickle.dump(func_sigs, open( idenLibCache, "wb" ))
    pickle.dump(mainSigs, open( idenLibCacheMain, "wb" ))
    print("[idenLib] Signatures refreshed...\n")

def idenLib():
    global func_sigs
    global mainSigs
    # function sigs from the current binary
    func_bytes_addr = {}
    for addr, size in getFuncRanges():
        f_bytes = getOpcodes(addr, size)
        func_bytes_addr[f_bytes] = addr
        
    # load sigs
    if not os.path.isdir(symEx_dir):
        print("[idenLib - FAILED] There is no {} directory".format(symEx_dir))
        return

    if os.path.isfile(idenLibCache):
        func_sigs = pickle.load( open( idenLibCache, "rb" ) )
        if os.path.isfile(idenLibCacheMain):
            mainSigs = pickle.load( open( idenLibCacheMain, "rb" ) )
    else:
        idenLibProcessSignatures()
    # apply sigs
    counter = 0
    mainDetected = False
    for sig_opcodes, addr in func_bytes_addr.items():
        if func_sigs.has_key(sig_opcodes):
            func_name = func_sigs[sig_opcodes][0]
            current_name = ida_funcs.get_func_name(addr)
            if (current_name != func_name):
                digit = 1
                while func_name in getNames():
                    func_name = func_name + str(digit)
                    digit = digit + 1
                ida_name.set_name(addr, func_name, ida_name.SN_NOCHECK)
                print("{}: {}".format(hex(addr), func_name))
                counter = counter + 1
        if mainSigs.has_key(sig_opcodes): # "main" sig
            callInstr = mainSigs[sig_opcodes][1] + addr
            if ida_ua.print_insn_mnem(callInstr) == "call":
                call_target = idc.get_operand_value(callInstr, 0)
                current_name = ida_funcs.get_func_name(call_target)
                func_name = mainSigs[sig_opcodes][0]
                if (current_name != func_name):
                    ida_name.set_name(call_target, func_name, ida_name.SN_NOCHECK)
                    print("{}: {}".format(hex(call_target), func_name))
                    counter = counter + 1
                    mainDetected = True
    if not mainDetected:
        for entry in idautils.Entries():
            for sig_opcodes, name_funcRva_EntryRva in mainSigs.items():
                callInstr = name_funcRva_EntryRva[2] + entry[2] # from EP
                if ida_ua.print_insn_mnem(callInstr) == "call":
                    fromFunc = name_funcRva_EntryRva[1]
                    func_start = callInstr - fromFunc
                    func_opcodes = getOpcodes(func_start, MAX_FUNC_SIZE)
                    if func_opcodes.startswith(sig_opcodes):
                        call_target = idc.get_operand_value(callInstr, 0)
                        current_name = ida_funcs.get_func_name(call_target)
                        func_name = mainSigs[sig_opcodes][0]
                        if (current_name != func_name):
                            ida_name.set_name(call_target, func_name, ida_name.SN_NOCHECK)
                            print("{}: {}".format(hex(call_target), func_name))
                            counter = counter + 1
                            mainDetected = True
                            break
            
    print("[idenLib] Applied to {} function(s)".format(counter))


class idenLib_class(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
    def activate(self, ctx):
        idenLib()
        return 1
    def update(self, ctx):
        return idaapi.AST_ENABLE_FOR_WIDGET if ctx.widget_type == idaapi.BWN_DISASM else idaapi.AST_DISABLE_FOR_WIDGET

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

    def activate(self, ctx):
        print(PLUGIN_NAME + " " + PLUGIN_VERSION)
        print(IDENLIB_CONTACT)
        return 1

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

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

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

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

# icon author: https://www.flaticon.com/authors/freepik
icon_data = "".join([
    "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A\x00\x00\x00\x0D\x49\x48\x44\x52\x00\x00\x00\x18\x00\x00\x00\x18\x08\x03\x00\x00\x00\xD7\xA9\xCD\xCA\x00\x00\x00\x4E\x50\x4C\x54\x45\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xC4\xA2\xA6\x59\x00\x00\x00\x19\x74\x52\x4E\x53\x00\x20\xEE\x4F\xC9\x64\xD3\xB3\x32\x99\x88\x17\x0C\xC1\x5C\x28\xF6\x7F\xE6\xDD\xBB\xA2\x47\x41\x90\xCE\x19\x07\xA1\x00\x00\x00\xC8\x49\x44\x41\x54\x28\xCF\x75\xD1\xDB\xAE\x83\x20\x10\x85\xE1\x35\x08\x0E\xCA\x16\x3C\xDB\xF5\xFE\x2F\xBA\xC7\x58\xDB\xB4\xA1\xFF\x8D\xC8\x27\x48\x02\x7E\x26\xD6\xDF\xE7\x58\x70\x46\xAB\x79\x82\x23\x19\xD4\x31\x55\xC1\x93\x47\x75\xAB\xFD\x10\xA9\xAE\x38\x16\xEA\x0B\x36\x6F\x6D\x88\x56\x8A\xE4\xFC\x02\xA5\xA5\x58\x9C\x73\x19\x23\x99\x6E\x88\x12\xA3\x94\x6B\x2B\x78\x9B\xB8\xA1\xA5\x9B\xE9\x9F\xF0\x20\xA7\x37\x58\x37\x64\x52\xAB\x50\x48\x57\x85\xF3\x21\x55\x18\x6C\xA6\x0A\x3D\xD9\x1B\x68\x37\x7E\x41\xD3\x4E\x0A\x2C\x40\xF7\x05\x12\x60\x2B\x5C\xC2\x70\x43\x0E\x21\x14\xD8\x97\xD0\x02\x8E\xB3\xFD\xA3\x1D\xD4\x0F\xD0\x75\x5D\x77\x03\x1D\x99\xD1\x5B\x25\xED\x21\x34\x09\x93\x8D\xA3\x41\x9E\xEC\xA5\xB3\xA2\xBF\xB6\x7A\xD8\xF8\x04\xD9\xDA\xA1\x76\x5C\x24\x3A\xBD\x6E\x4D\xCE\xD2\xFB\x36\x05\xBF\xFB\x07\x19\xFC\x16\xA4\x38\xC6\x08\x3D\x00\x00\x00\x00\x49\x45\x4E\x44\xAE\x42\x60\x82"
])

class idenLibMain(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "idenLib - Library Function Identification"
    help = IDENLIB_CONTACT
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ''

    def run(self, arg):
        idenLib()
        pass

    def term(self):
        pass

    def init(self):
        act_icon = idaapi.load_custom_icon(data=icon_data, format="png")
        act_name = "idenLib:action"
        idaapi.register_action(idaapi.action_desc_t(
                act_name,
                "idenLib - Function Identification",
                idenLib_class(),
                None,
                "idenLib",
                act_icon))
        # Insert the action in a toolbar
        idaapi.attach_action_to_toolbar("DebugToolBar", act_name)
        idaapi.attach_action_to_menu(
                    'Edit/idenLib/',
                    act_name,
                    idaapi.SETMENU_APP)

        # refresh signatures
        act_name = "idenLib:refresh"
        idaapi.register_action(idaapi.action_desc_t(
                act_name,
                "Refresh Signatures",
                RefreshHandler(),
                None,
                "idenLib - Refresh"))
        idaapi.attach_action_to_menu(
                    'Edit/idenLib/',
                    act_name,
                    idaapi.SETMENU_APP)

        # about
        act_name = "idenLib:about"
        idaapi.register_action(idaapi.action_desc_t(
                act_name,
                "About",
                AboutHandler(),
                None,
                "idenLib - About"))
        idaapi.attach_action_to_menu(
                    'Edit/idenLib/',
                    act_name,
                    idaapi.SETMENU_APP)

        return idaapi.PLUGIN_OK



def PLUGIN_ENTRY():
    return idenLibMain()

if __name__ == "__main__":
    PLUGIN_ENTRY()


repository: mahmoudimus/.idapro
file: scripts/xray.py
created_at: 2025-02-16 21:33:20+00:00
pushed_at: 2025-04-09 18:15:30+00:00
url:
https://github.com/mahmoudimus/.idapro/blob/e0bf0d58fb6de4eb6af0b87ffbd869d8c7deef08/scripts/xray.py

content:
import shutil, re, os, configparser, errno
from PyQt5.QtWidgets import QWidget, QSplitter

import ida_hexrays
import ida_idaapi
import ida_kernwin as kw
import ida_lines as il
import ida_diskio
from ida_pro import IDA_SDK_VERSION

__author__ = "Dennis Elser"

PLUGIN_NAME = "xray"

XRAY_FILTER_ACTION_ID = "%s:filter" % PLUGIN_NAME
XRAY_LOADCFG_ACTION_ID = "%s:loadcfg" % PLUGIN_NAME
XRAY_QUERY_ACTION_ID = "%s:query" % PLUGIN_NAME
PATTERN_LIST = []
HIGH_CONTRAST = False

IS_ACTIVATED = False
TEXT_INPUT_FORMS = {}

CFG_FILENAME = "%s.cfg" % PLUGIN_NAME
DEFAULT_CFG = """# configuration file for xray.py

[global]
# set to 1 for better contrast
high_contrast=0
# enable/disable xray when loading a database 
auto_enable=0

# each group contains a list of regular
# expressions, a background color in
# RRGGBB format and an optional hint field.
# priority is determined by order of
# appearance, first group gets assigned
# lowest priority.
# check out https://regex101.com/r and
# https://www.debuggex.com/

[group_01]
hint=loop
bgcolor=4c0037

expr_01=^while\(
expr_02=^for\(

[group_02]
hint=function name
bgcolor=00374c

expr_01=recv\(
expr_02=malloc\(
expr_03=realloc\(
expr_04=free\(
expr_05=memcpy\(
expr_06=memmove\(
expr_07=strcpy\(
expr_08=sscanf\(
expr_09=sprintf\(
expr_10=recvfrom\(

[group_03]
hint=format strings
bgcolor=4c1500

expr_01=sscanf\(.*,.*%s.*,.*\)
expr_02=sprintf\(.*,.*%s.*,.*\)

[group_04]
hint=arithmetic
bgcolor=4c1500

expr_01=malloc\(.*[\*\+\-\/%][^>].*?\)
expr_02=realloc\(([^,]+,){1}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_03=memcpy\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_04=memmove\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_05=recv\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_06=recvfrom\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])"""

# -----------------------------------------------------------------------------
def is_plugin():
    """returns True if this script is executed from within an IDA plugins
    directory, False otherwise."""
    return "__plugins__" in __name__

# -----------------------------------------------------------------------------
def get_dest_filename():
    """returns destination path for plugin installation."""
    return os.path.join(
        ida_diskio.get_user_idadir(),
        "plugins",
        "%s%s" % (PLUGIN_NAME, ".py"))

# -----------------------------------------------------------------------------
def is_installed():
    """checks whether script is present in designated plugins directory."""
    return os.path.isfile(get_dest_filename())

# -----------------------------------------------------------------------------
def get_cfg_filename():
    """returns full path for config file."""
    return os.path.join(
        ida_diskio.get_user_idadir(),
        "plugins",
        "%s" % CFG_FILENAME)

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

# -----------------------------------------------------------------------------
def is_compatible():
    """Checks whether script is compatible with current IDA and
    decompiler versions."""
    min_ida_ver = 720
    return is_ida_version(min_ida_ver) and ida_hexrays.init_hexrays_plugin()

# -----------------------------------------------------------------------------
SELF = __file__
def install_plugin():
    """Installs script to IDA userdir as a plugin."""
    dst = get_dest_filename()
    src = SELF
    if is_installed():
        btnid = kw.ask_yn(kw.ASKBTN_NO,
            "File exists:\n\n%s\n\nReplace?" % dst)
        if btnid is not kw.ASKBTN_YES:
            return False
    else:
        btnid = kw.ask_yn(kw.ASKBTN_NO,
            "This plugin is about to be installed to:\n\n%s\n\nInstall now?" % dst)
        if btnid is not kw.ASKBTN_YES:
            return False

    usrdir = os.path.dirname(dst)
    kw.msg("%s: copying script from \"%s\" to \"%s\" ..." % (PLUGIN_NAME, src, usrdir))
    if not os.path.exists(usrdir):
        try:
            os.makedirs(usrdir)
        except OSError as e:
            if e.errno != errno.EEXIST:
                kw.msg("failed (mkdir)!\n")
                return False
    try:
        shutil.copy(src, dst)
    except:
        kw.msg("failed (copy)!\n")
        return False
    kw.msg(("done\n"
        "Plugin installed - please restart this instance of IDA.\n"))
    return True

# -----------------------------------------------------------------------------
def swapcol(x):
    """converts between RRGGBB and BBGGRR color encodings."""
    return (((x & 0x000000FF) << 16) |
             (x & 0x0000FF00) |
            ((x & 0x00FF0000) >> 16))

# -----------------------------------------------------------------------------
def load_cfg(reload=False):
    """loads xray configuration from file. Creates and loads default config
    if none is present."""
    global PATTERN_LIST
    global HIGH_CONTRAST
    global IS_ACTIVATED

    cfg_file = get_cfg_filename()
    kw.msg("%s: %sloading %s...\n" % (PLUGIN_NAME,
        "re" if reload else "",
        cfg_file))
    if not os.path.isfile(cfg_file):
        kw.msg("%s: %s does not exist! creating default config... " % (PLUGIN_NAME, cfg_file))
        try:
            with open(cfg_file, "w") as f:
                f.write(DEFAULT_CFG)
                kw.msg("success!\n")
        except:
            kw.msg("failed!\n")
            return False
        return load_cfg(reload=True)

    PATTERN_LIST = []

    config = configparser.RawConfigParser()
    config.readfp(open(cfg_file))

    # read all sections
    for section in config.sections():
        expr_list = []
        if section.startswith("group_"):
            for k,v in config.items(section):
                if k.startswith("expr_"):
                    expr_list.append(v)
            try:
                bgcolor = swapcol(int(config.get(section, "bgcolor"), 16))
            except:
                bgcolor = swapcol(0x000000)
            try:
                hint = config.get(section, "hint")
            except:
                hint = None
            PATTERN_LIST.append(ConfigGroupSection(expr_list, bgcolor, hint))
        elif section == "global":
            try:
                HIGH_CONTRAST = config.getboolean(section, "high_contrast")
            except:
                HIGH_CONTRAST = False
            if not reload:
                try:
                    IS_ACTIVATED = config.getboolean(section, "auto_enable")
                except:
                    IS_ACTIVATED = False

    if not len(PATTERN_LIST):
        kw.warning("Config file does not contain any regular expressions.")
    return True

# -----------------------------------------------------------------------------
class TextInputForm(kw.Form):
    """Input form for regex search queries."""

    # flags
    SO_FIND_TEXT = 1
    SO_FIND_REGEX = 2
    SO_FILTER_TEXT = 4
    SO_FILTER_COLOR = 8
    SO_FIND_CASE = 16

    def __init__(self, parent_widget):
        self.parent_widget = parent_widget
        self.parent_title = kw.get_widget_title(self.parent_widget)
        i=1
        while kw.find_widget("%s-%d" % (PLUGIN_NAME, i)):
            i+=1
        self.idx = i
        __title = "%s-%s" % (PLUGIN_NAME, self.idx)
        self.options = (TEXT_INPUT_FORMS[self.parent_title].options
            if self.parent_title in TEXT_INPUT_FORMS.keys()
            else TextInputForm.SO_FILTER_TEXT | TextInputForm.SO_FIND_TEXT)
        self.query = (TEXT_INPUT_FORMS[self.parent_title].query
            if self.parent_title in TEXT_INPUT_FORMS.keys()
            else "")
        kw.Form.__init__(self,
("BUTTON YES NONE\n"
"BUTTON NO NONE\n"
"BUTTON CANCEL NONE\n"
"%s\n\n"
"{FormChangeCb}\n"
"<##Enter text##Filter:{cbEditable}>"
"|<##Search##Text:{rText}><Regex:{rRegex}>{cSearchMethod}>"
"|<##Filter##Lines:{rFilterLines}><Colors:{rFilterColors}>{cFilterType}>"
"|<##Options##Case sensitive:{rCase}>{cOptions}>\n") % (__title), {
    'FormChangeCb': kw.Form.FormChangeCb(self.OnFormChange),
    'cbEditable': kw.Form.StringInput(value = self.query),
    'cSearchMethod': kw.Form.RadGroupControl(("rText", "rRegex")),
    'cFilterType': kw.Form.RadGroupControl(("rFilterLines", "rFilterColors")),
    'cOptions': kw.Form.ChkGroupControl(("rCase",))})

    def init_controls(self):
        self.SetControlValue(self.cbEditable, self.query)
        self.SetControlValue(self.cSearchMethod, 0 if self.options & TextInputForm.SO_FIND_TEXT else 1)
        self.SetControlValue(self.cFilterType, 0 if self.options & TextInputForm.SO_FILTER_TEXT else 1)
        self.SetControlValue(self.cOptions, 1 if self.options & TextInputForm.SO_FIND_CASE else 0)
        self.SetFocusedField(self.cbEditable)
        return

    def _commit_changes(self):
        vu = ida_hexrays.get_widget_vdui(self.parent_widget)
        if vu:
            vu.refresh_ctext()
            # "refresh_ctext()" took away the focus, take it back
            kw.activate_widget(kw.find_widget(self.title), True)
            self.SetFocusedField(self.cbEditable)
            return True
        return False

    def OnFormChange(self, fid):
        if fid == self.cbEditable.id:
            self.query = self.GetControlValue(self.cbEditable)
        elif fid == self.rCase.id:
            if self.GetControlValue(self.cOptions):
                self.options |= TextInputForm.SO_FIND_CASE
            else:
                self.options &= ~TextInputForm.SO_FIND_CASE & 0xFFFFFFFF
        elif fid in [self.rFilterLines.id, self.rFilterColors.id]:
            filter_text = fid == self.rFilterLines.id
            filter_color = fid == self.rFilterColors.id

            if filter_text:
                self.options |= TextInputForm.SO_FILTER_TEXT
            else:
                self.options &= ~TextInputForm.SO_FILTER_TEXT & 0xFFFFFFFF

            if filter_color:
                self.options |= TextInputForm.SO_FILTER_COLOR
            else:
                self.options &= ~TextInputForm.SO_FILTER_COLOR & 0xFFFFFFFF
        elif fid in [self.rText.id, self.rRegex.id]:
            find_text = fid == self.rText.id
            find_regex = fid == self.rRegex.id

            if find_text:
                self.options |= TextInputForm.SO_FIND_TEXT
            else:
                self.options &= ~TextInputForm.SO_FIND_TEXT & 0xFFFFFFFF

            if find_regex:
                self.options |= TextInputForm.SO_FIND_REGEX
            else:
                self.options &= ~TextInputForm.SO_FIND_REGEX & 0xFFFFFFFF
        
        self._commit_changes()
        return 1

# -----------------------------------------------------------------------------
class ConfigGroupSection():
    """class that represents a config file's "group" section."""
    def __init__(self, expr_list, bgcolor, hint):
        self.expr_list = expr_list
        self.bgcolor = bgcolor
        self.hint = hint

# -----------------------------------------------------------------------------
class xray_hooks_t(ida_hexrays.Hexrays_Hooks):
    """class for handling decompiler events."""

    def _remove_color_tags(self, l):
        """removes all color tags from a tagged simple_line_t object
        but preserves COLOR_ADDR tags."""
        line = ""
        i = 0
        while i<len(l):
            if l[i] is il.COLOR_ON:
                n = il.tag_skipcode(l[i:])
                if l[i:].find(chr(il.COLOR_ADDR)) == 1:
                    line += l[i:i+n]
                i += n
            elif l[i] in [il.COLOR_OFF, il.COLOR_ESC, il.COLOR_INV]:
                n = il.tag_skipcode(l[i:])
                i += n
            else:
                line += l[i]
                i += 1
        return line

    def _search(self, regexp, sl, case_sensitive = False):
        line = il.tag_remove(sl.line).lstrip().rstrip()
        return re.search(regexp, line, flags=re.I if not case_sensitive else 0) is not None

    def _apply_xray_filter(self, vu, pc):
        if IS_ACTIVATED and pc:
            #col = il.calc_bg_color(ida_idaapi.get_inf_structure().min_ea)
            #col = pc[0].bgcolor
            for sl in pc:
                match=False
                for group in PATTERN_LIST:
                    for expr in group.expr_list:
                        if self._search(expr, sl):
                            #sl.bgcolor = (col & 0xfefefe) >> 1
                            sl.bgcolor = group.bgcolor
                            match=True
                            break
                if not match and HIGH_CONTRAST:
                    sl.line = self._remove_color_tags(sl.line)
        return

    def _apply_query_filter(self, vu, pc):
        new_pc = []
        title = kw.get_widget_title(vu.ct)
        if title in TEXT_INPUT_FORMS.keys() and pc:
            sq = TEXT_INPUT_FORMS[title]
            query = sq.query
            options = sq.options
            case_sensitive = options & TextInputForm.SO_FIND_CASE

            # TODO
            if options & TextInputForm.SO_FIND_TEXT:
                kw.set_highlight(vu.ct, query, HL_FLAGS)
                tmpquery = query.lower() if not case_sensitive else query
                for sl in pc:
                    haystack = il.tag_remove(sl.line).lstrip().rstrip()
                    haystack = haystack.lower() if not case_sensitive else haystack
                    if tmpquery in haystack:
                        new_pc.append(sl.line)
                    else:
                        if options & TextInputForm.SO_FILTER_COLOR:
                            # add line but remove color
                            new_pc.append(self._remove_color_tags(sl.line))
                        elif options & TextInputForm.SO_FILTER_TEXT:
                            # do not add non-matching text
                            pass
            elif options & TextInputForm.SO_FIND_REGEX:
                kw.set_highlight(vu.ct, None, 0)
                for sl in pc:
                    try:
                        if self._search(query, sl, case_sensitive):
                            new_pc.append(sl.line)
                        else:
                            if options & TextInputForm.SO_FILTER_COLOR:
                                new_pc.append(self._remove_color_tags(sl.line))
                            elif options & TextInputForm.SO_FILTER_TEXT:
                                # do not add non-matching text
                                pass
                    except re.error as error:
                        kw.msg("%s: %s: \"%s\"" %
                            (PLUGIN_NAME, error, query))
                        return
            pc.clear()
            sl = kw.simpleline_t()
            for line in new_pc:
                sl.line = line
                pc.push_back(sl)
        return

    def _build_hint(self, vu):
        if IS_ACTIVATED and vu.refresh_cpos(ida_hexrays.USE_MOUSE):
            sl = vu.cfunc.get_pseudocode()[vu.cpos.lnnum]
            hint_lines = ["%s pattern(s):" % PLUGIN_NAME]
            delim_s = "%s" % "="*len(hint_lines[0])
            delim_e = "\n%s\n" % ("-"*len(hint_lines[0]))
            hint_lines.append(delim_s)
            hint = ""
            hint_created = False
            for group in PATTERN_LIST:
                for expr in group.expr_list:
                    if self._search(expr, sl):
                        tmp = (" (%s)" % group.hint) if group.hint else ""
                        hint_lines.append("> \"%s\"%s" % (expr, tmp))
                        hint_created = True
            hint_lines.append(delim_e)
            hint = "\n".join(hint_lines)
            if hint_created:
                return (hint, len(hint_lines)+1)
        return None

    def text_ready(self, vu):
        pc = vu.cfunc.get_pseudocode()
        self._apply_query_filter(vu, pc)
        self._apply_xray_filter(vu, pc)
        return 0

    def populating_popup(self, widget, phandle, vu):
        kw.attach_action_to_popup(vu.ct, None, XRAY_FILTER_ACTION_ID, PLUGIN_NAME+"/")
        kw.attach_action_to_popup(vu.ct, None, XRAY_LOADCFG_ACTION_ID, PLUGIN_NAME+"/")
        kw.attach_action_to_popup(vu.ct, None, XRAY_QUERY_ACTION_ID, PLUGIN_NAME+"/")
        return 0

    def create_hint(self, vu):
        result = self._build_hint(vu)
        if result:
            hint, count = result
            return (2, hint, count)
        return (0, None)

# -----------------------------------------------------------------------------
class xray_action_handler_t(kw.action_handler_t):
    """action handler for turning xray on and off."""
    def __init__(self):
        kw.action_handler_t.__init__(self)

    def activate(self, ctx):
        global IS_ACTIVATED
        IS_ACTIVATED = not IS_ACTIVATED
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        if vu:
            vu.refresh_ctext()
        kw.msg("[%s] %sactivated.\n" % (PLUGIN_NAME, "" if IS_ACTIVATED else "de"))
        return 1

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

# -----------------------------------------------------------------------------
class loadcfg_action_handler_t(kw.action_handler_t):
    """action handler for reloading xray cfg file."""
    def __init__(self):
        kw.action_handler_t.__init__(self)

    def activate(self, ctx):
        if load_cfg(reload=True):
            vu = ida_hexrays.get_widget_vdui(ctx.widget)
            if vu:
                vu.refresh_ctext()
        return 1

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

# -----------------------------------------------------------------------------
class regexfilter_action_handler_t(kw.action_handler_t):
    """action handler for search queries."""
    def __init__(self):
        kw.action_handler_t.__init__(self)

    # TODO
    def _dirty_resize_hack(self, w, form):
        title = form.title
        widget = kw.find_widget(title)
        if not widget:
            return

        w1 = kw.PluginForm.TWidgetToPyQtWidget(widget)
        w2 = kw.PluginForm.TWidgetToPyQtWidget(w)
        if not w1 or not w2:
            return

        p1 = w1.parentWidget()
        p2 = w2.parentWidget()
        if not p1 or not p2:
            return

        splitter = p1.parentWidget()
        hr = p2.parentWidget()
        if not splitter or not hr:
            return

        if not type(splitter) is QSplitter or not type(hr) is QWidget:
            return

        sizes = splitter.sizes()
        if len(sizes) != 2:
            return

        idx = splitter.indexOf(p1)
        _min, _max = splitter.getRange(idx)
        sizes[idx] = _min

        idx = splitter.indexOf(hr)
        _min, _max = splitter.getRange(idx)
        sizes[idx] = _max

        splitter.setSizes(sizes)
        return

    def _open_search_form(self, widget):
        global TEXT_INPUT_FORMS

        title = kw.get_widget_title(widget)
        if title not in TEXT_INPUT_FORMS.keys():
            search_form = TextInputForm(widget)
            search_form.modal = False
            search_form.openform_flags = (kw.PluginForm.WOPN_DP_BOTTOM |
                kw.PluginForm.WOPN_PERSIST)
            search_form, _ = search_form.Compile()
            search_form.Open()
            TEXT_INPUT_FORMS[title] = search_form
            self._dirty_resize_hack(widget, search_form)
        else:
            search_form = TEXT_INPUT_FORMS[title]
            search_form.Open()
            search_form.init_controls()
            self._dirty_resize_hack(widget, search_form)
        return

    def activate(self, ctx):
        self._open_search_form(ctx.widget)
        return 1

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

# -----------------------------------------------------------------------------
class xray_plugin_t(ida_idaapi.plugin_t):
    """plugin class."""
    flags = ida_idaapi.PLUGIN_HIDE
    comment = PLUGIN_NAME
    help = PLUGIN_NAME
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""

    def init(self):
        self.xray_hooks = None
        if not is_compatible():
            kw.msg("%s: decompiler not available, skipping." % PLUGIN_NAME)
            return ida_idaapi.PLUGIN_SKIP

        load_cfg()

        kw.register_action(
            kw.action_desc_t(
                XRAY_LOADCFG_ACTION_ID,
                "%s: reload config" % PLUGIN_NAME,
                loadcfg_action_handler_t(),
                "Ctrl-R"))

        kw.register_action(
            kw.action_desc_t(
                XRAY_FILTER_ACTION_ID,
                "%s: toggle" % PLUGIN_NAME,
                xray_action_handler_t(),
                "F3"))

        kw.register_action(
            kw.action_desc_t(
                XRAY_QUERY_ACTION_ID,
                "%s: search" % PLUGIN_NAME,
                regexfilter_action_handler_t(),
                "Ctrl-F"))

        self.xray_hooks = xray_hooks_t()
        self.xray_hooks.hook()
        kw.msg("[%s] F3: toggle xray, Ctrl-F: filter, Ctrl-R: reload config." % PLUGIN_NAME)
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        return

    def term(self):
        if self.xray_hooks:
            self.xray_hooks.unhook()
            kw.unregister_action(XRAY_FILTER_ACTION_ID)
            kw.unregister_action(XRAY_LOADCFG_ACTION_ID)
            kw.unregister_action(XRAY_QUERY_ACTION_ID)
        return

# -----------------------------------------------------------------------------
def PLUGIN_ENTRY():
    """plugin entry point."""
    return xray_plugin_t()

# -----------------------------------------------------------------------------
def SCRIPT_ENTRY():
    """script entry point."""
    if not is_plugin():
        (kw.info("Success!") if install_plugin() else
            kw.warning("Error! Plugin could not be installed!"))
    return

# -----------------------------------------------------------------------------
HL_FLAGS = kw.HIF_LOCKED | kw.HIF_NOCASE if is_ida_version(740) else kw.HIF_LOCKED
SCRIPT_ENTRY()


repository: The-Itach1/Audit
file: A_M_Audit/A_M_Audit.py
description: 使用idapython写的一个辅助二进制固件漏洞挖掘的代码审计脚本
wanted_name: A_M_Audit
comment: A_M_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/A_M_Audit/A_M_Audit.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



class CArmAudit:
    def getArgAddr(self,start_addr, regNum):
        armcondition = []
        scan_deep = 50
        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(self,start_addr, regNum):
        arm = ['LDR', 'CMP', 'STR', 'MOV']
        arg_addr = self.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 = self.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(self,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(self.getArg(call_addr, num))
        ret_list.append(local_buf_size)
        return ret_list

    def getFormatString(self,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:
            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(self,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(self.getArg(call_addr, num))
        mid_arg_addr = self.getArgAddr(call_addr, format_function_offset_dict[func_name])

        Rx_arg = self.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 = self.getArgAddr(mid_arg_addr, int(Rx_arg[-1]))
                true_arg_addr = self.getArgAddr(true_arg_addr, int(Rx_arg[-1]))
            else:
                true_arg_addr = self.getArgAddr(mid_arg_addr, int(Rx_arg[1:]))
                true_arg_addr = self.getArgAddr(true_arg_addr, int(Rx_arg[1:]))
        else:
            true_arg_addr = self.getArgAddr(mid_arg_addr, format_function_offset_dict[func_name])
        string_and_addr = self.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(self.getArg(call_addr, num))
            ret_list.append(format_and_value)
        ret_list.append(local_buf_size)
        return ret_list

    def audit(self,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 = self.auditFormat(call_addr, func_name, arg_num)
                else:
                    info = self.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(self):
        # the word create with figlet
        print("Auditing dangerous functions ......")
        for func_name in dangerous_functions:
            self.audit(func_name)

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

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

        print("Finished! Enjoy the result ~")

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        print("Finished! Enjoy the result ~")

m_initialized = False


class A_M_Audit_Plugin_t(idaapi.plugin_t):
    comment = "A_M_Audit plugin for IDA Pro"  # 描述信息
    help = "todo"  # 帮助信息
    wanted_name = "A_M_Audit"  # 菜单中显示的名字
    wanted_hotkey = "Ctrl-Alt-T"  # 希望注册的快捷键
    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()
        print(info.procName)
        if 'mips' in info.procName:
            m_mips=CMipsAudit()
            m_mips.MipsAudit()
        elif 'ARM' in info.procName:
            m_arm=CArmAudit()
            m_arm.ArmAudit()
        else:
            print('A_M_Audit is not supported on the current arch')

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


repository: Donaldduck8/IDA-Scripts-and-Tools
file: donald_ida_utils.py
description: A collection of plug-ins I've developed for IDA Pro.
wanted_name: Dummy Plugin
comment: Dummy Plugin
created_at: 2023-08-02 10:06:40+00:00
pushed_at: 2024-04-09 16:47:04+00:00
url:
https://github.com/Donaldduck8/IDA-Scripts-and-Tools/blob/4ec79fdcd2a2a43348282ee427944ed8a7593a57/donald_ida_utils.py

content:
import re
import envi
import idc
import idaapi
import ida_ida
import idautils
import ida_name
import ida_xref
import ida_idaapi
import ida_search
import ida_hexrays
import ida_kernwin

from typing import List


DEFAULT_PREFIX = "Decrypted: "


import ida_bytes
import ida_name
import ida_ua

def define_and_rename_global(addr, name, data_type=ida_ua.dt_dword):
    """
    Check if a global variable exists at the given address. If not, define one and then rename it.

    :param addr: Address to check and define the global variable at.
    :param name: New name for the global variable.
    :param data_type: The data type to define at the address (default is dword).
    """
    # Check if the address is unknown (i.e., not defined)
    if ida_bytes.is_unknown(ida_bytes.get_flags(addr)):
        # Define the global variable at the address with the specified data type
        if not ida_bytes.create_data(addr, data_type, ida_bytes.get_data_elsize(addr, data_type), ida_idaapi.BADADDR):
            print(f"Failed to create data at address {hex(addr)}.")
            return False
        
    # Rename the address to the specified name
    if not ida_name.set_name(addr, name, ida_name.SN_CHECK):
        print(f"Failed to rename address {hex(addr)} to {name}.")
        return False

    print(f"Successfully defined and renamed global at {hex(addr)} to {name}.")
    return True


def get_all_instructions_in_line(ida_cfunc, ea) -> List[int]:
    """Return the address of every instruction that corresponds to a desired line of pseudocode."""
    insn_eas = []
    for ida_item in ida_cfunc.get_boundaries().items():
        range_set = ida_item[1]
        num_ranges = range_set.nranges()

        for i in range(num_ranges):
            boundary_start = range_set.getrange(i).start_ea
            boundary_end = range_set.getrange(i).end_ea
            range_size = boundary_end - boundary_start

            # If the boundary is hilariously large, ignore it
            if range_size > 0x100:
                continue

            if boundary_start <= ea < boundary_end:
                while boundary_end > boundary_start:
                    boundary_end = idaapi.prev_head(boundary_end, boundary_start)
                    insn_eas.append(boundary_end)

    return insn_eas


def get_name_for_address(ea, ref_addr=None):
    if ref_addr == None:
        ref_addr = ida_idaapi.BADADDR
    name = ida_name.get_ea_name(ea, ida_name.calc_gtn_flags(ref_addr, ea))
    if not name:
        name = hex(ea)

    return name


def add_pseudocode_comment(ea, comment, prefix=DEFAULT_PREFIX, quoted=True, sanitize=True, add_to_existing=True) -> None:
    """Add comment to the line of pseudocode that corresponds to the provided address."""
    ida_func = idaapi.get_func(ea)

    # Check if the function exists and has a decompiled representation
    if ida_func is None or not idaapi.init_hexrays_plugin():
        print("Error: Unable to find the function or decompiled view for address:", hex(ea))
        return

    ida_cfunc = idaapi.decompile(ida_func)

    if sanitize:
        # Carriage Return and Line Feed
        def newlines_escape(match):
            return match.group().replace("\r", "\\r").replace("\n", "\\n")

        trailing_newlines = re.compile(r"[\r\n]+?$")
        comment = trailing_newlines.sub(newlines_escape, comment)

        starting_newlines = re.compile(r"^[\r\n]+?")
        comment = starting_newlines.sub(newlines_escape, comment)

        # Control Characters
        def control_chars_to_hex(match):
            return r"\x{0:02x}".format(ord(match.group()))

        control_chars_class = re.compile(r"[\x00-\x09\x0B\x0C\x0E-\x1F]")
        comment = control_chars_class.sub(control_chars_to_hex, comment)

    if quoted:
        comment = '"' + comment + '"'

    if prefix:
        comment = prefix + comment

    insn_ea = ida_cfunc.eamap[ea][0].ea
    treeloc = idaapi.treeloc_t()
    treeloc.ea = insn_ea
    treeloc.itp = idaapi.ITP_BLOCK1 # BLOCK1 is the only reliable ITP

    if add_to_existing:
        existing_cmt = ida_cfunc.get_user_cmt(treeloc, ida_hexrays.RETRIEVE_ALWAYS)

        if existing_cmt:
            comment = existing_cmt + "\n" + comment

    print("Adding comment to", hex(insn_ea), comment)
    ida_cfunc.set_user_cmt(treeloc, comment)
    ida_cfunc.save_user_cmts()


def add_disassembly_comment(ea, text):
    ## Set in dissassembly
    idc.set_cmt(ea, text,0)


def find_references_to(ea) -> List[int]:
    """Find Xrefs and immediate references to an address."""
    refs = []

    # Xrefs
    refs += [x.frm for x in list(idautils.XrefsTo(ea, ida_xref.XREF_ALL))]

    # References to immediate value
    found_eas = [0]
    while True:
        result_ea, result_code = ida_search.find_imm(found_eas[-1], ida_search.SEARCH_DOWN, ea)

        if (result_code == -1 or
                result_ea in found_eas or
                result_ea == 0xffffffffffffffff or
                result_ea in refs):
            break

        found_eas.append(result_ea)
        refs.append(result_ea)

    return refs


def open_synced_disassembly_view():
    # Get active view title
    pseudocode_view = ida_kernwin.get_current_viewer()
    pseudocode_view_title = ida_kernwin.get_widget_title(pseudocode_view)

    # Open disassembly view
    disasm_view_title = f"Synced Disasm ({pseudocode_view_title})"
    disasm_view = ida_kernwin.open_disasm_window(disasm_view_title)

    # Set disassembly view to text view
    ida_kernwin.set_view_renderer_type(disasm_view, ida_kernwin.TCCRT_FLAT)

    # Sync the disassembly view to the pseudocode view
    what = ida_kernwin.sync_source_t(disasm_view)
    _with = ida_kernwin.sync_source_t(pseudocode_view)
    ida_kernwin.sync_sources(what, _with, True)

    # Dock the disassembly view to the right of the pseudocode view
    ida_kernwin.set_dock_pos(pseudocode_view_title, disasm_view_title, ida_kernwin.WOPN_DP_RIGHT)


class DummyPlugin(ida_idaapi.plugin_t):
    """Dummy plugin to make IDA stop complaining when this file is in the plugins folder."""
    
    # These fields are necessary for whatever reason
    flags = ida_idaapi.PLUGIN_UNL
    comment = "Dummy Plugin"
    help = ""
    wanted_name = "Dummy Plugin"
    wanted_hotkey = ""

    def init(self):
        return ida_idaapi.PLUGIN_UNL

    def run(self, args):
        pass

    def term(self):
        pass

def PLUGIN_ENTRY():
    return DummyPlugin()


repository: hugsy/ida-headless
file: ida_rpyc_server.py
description: IDA (sort of) headless
created_at: 2021-07-07 19:45:13+00:00
pushed_at: 2024-02-17 19:30:59+00:00
url:
https://github.com/hugsy/ida-headless/blob/2f2d6ec20fc0f5a824f3aedcf57cd812e52fa15f/ida_rpyc_server.py

content:
"""

> "Because IDA sucks"


Embed RPyc in IDA to expose IDA's API externally, by a background thread that runs the TCP server. Also in iPython this
provides autocomplete.

Props to https://github.com/vrtadmin/FIRST-plugin-ida/blob/master/first_plugin_ida/first.py#L87
for the workaround on the threading issue, for IDA Pro >= 7.2

Quick start
```
>>> import rpyc
>>> c = rpyc.connect("ida.rpyc.server", 18812)
#
# IDA namespace will be in `c.root`
#
>>> c.root.idaapi.get_root_filename()
'ntoskrnl.exe'
>>> hex( c.root.idc.here() )
0x140088194
>>> c.root.idaapi.jumpto( 0x1400881EE )
True
```

For more facility, you can alias it:
```
>>> idc = c.root.idc
```

Then, it becomes super readable
```
>>> idc.jumpto( idc.get_name_ea_simple("DriverEntry") )
True
>>> idc.set_cmt( idc.here(), "@hugsy was here", 1)
True
```

For generator objects, you now need to use the wrapper `c.root.iterate()`.

Example:
```
>>> idc = c.root.idc
>>> idautils = c.root.idautils
>>> for ea in c.root.iterate( idautils.Functions() ):
...    print( idc.get_func_name(ea) )
```

Blame HexRays for making their API more confusing at every release.

Ref:
- https://www.hex-rays.com/products/ida/support/ida74_idapython_no_bc695_porting_guide.shtml

Demo:
- https://youtu.be/obX2GreSsFU

"""

import sys
import random
import threading

import rpyc
import idc
import idaapi
import idautils
import ida_bytes
import ida_funcs
import ida_hexrays

# exported modules must be imported first
import sys
import os

PLUGIN_NAME = "RunRpycServer"
PLUGIN_HOTKEY = "Ctrl-Alt-K"
PLUGIN_VERSION = "0.3"
PLUGIN_AUTHOR = "@_hugsy_"

HOST, PORT = "0.0.0.0", 18812
DEBUG = False


EXPOSED_MODULES = [idaapi, idc, idautils, ida_hexrays, os, sys]

def xlog(x):
    sys.stderr.write("{} - {}\n".format(threading.current_thread().name, x)) and sys.stderr.flush()

def err(msg):
    xlog(f"[!] {msg}")

def ok(msg):
    xlog(f"[+] {msg}")

def dbg(msg):
    if DEBUG:
        xlog(f"[*] {msg}")



class IdaWrapper:
    def __getattribute__(self, name):
        default = "IDoNotExistButNoReallyISeriouslyDoNotAndCannotExist"

        dbg("trying to get {}".format(name,))

        if name.startswith("exposed_"):
            name = name.replace("exposed_", "")
            dbg("changed to get {}".format(name,))

        for mod in EXPOSED_MODULES:
            val = getattr(mod, name, default)
            if val != default:
                break

        if val == default:
            raise AttributeError("unknown {}".format(name,))

        if hasattr(val, '__call__'):
            dbg("{} is callable".format(val,))
            def call(*args, **kwargs):
                holder = [None] # need a holder, because 'global' sucks

                def trampoline():
                    holder[0] = val(*args, **kwargs)
                    return 1

                idaapi.execute_sync(trampoline, idaapi.MFF_WRITE)
                return holder[0]
            return call
        else:
            return val



g_IdaWrapper     = IdaWrapper()


class IdaRpycService(rpyc.Service):
    ALIASES = ["ida", ]


    def on_connect(self, conn):
        ok("connect open: {}".format(conn,))
        for mod in EXPOSED_MODULES:
            setattr(self, f"exposed_{mod.__name__}", g_IdaWrapper)
        return


    def on_disconnect(self, conn):
        ok(f"connection closed: {str(conn)}")
        return


    def exposed_eval(self, cmd):
        return eval(cmd)


    def exposed_exec(self, cmd):
        return exec(cmd)


    def exposed_iterate(self, iterator):
        default = f"IDoNotExistButNoReallyISeriouslyDoNotAndCannotExist {random.randint(0, 65535)}"
        holder = [default]
        def trampoline():
            try:
                holder[0] = next(iterator)
            except StopIteration:
                holder[0] = default
            return 1
        while True:
            idaapi.execute_sync(trampoline, idaapi.MFF_WRITE)
            if holder[0] == default:
                return
            yield holder[0]



g_IdaServer = IdaRpycService()



def start():
    global g_IdaServer
    srv = None

    for i in range(1):
        port = PORT + i
        try:
            srv = rpyc.utils.server.OneShotServer(g_IdaServer, hostname=HOST, port=port) if DEBUG \
                else rpyc.utils.server.ThreadedServer(g_IdaServer, hostname=HOST, port=port)
            break
        except OSError:
            srv = None

    if not srv:
        err("failed to start server...")
        return

    ok("starting server...")
    srv.start()
    srv.close()
    ok("server closed")
    return


t = None

def main():
    global t
    if t is not None:
        err(f"thread is already running as {str(t)}")
        return

    t = threading.Thread(target=start)
    t.daemon = True
    t.start()
    ok(f"service listening on {HOST}:{PORT}...")


class dummy(idaapi.plugin_t):
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_UNL
    comment = ""
    help = ""

    def init(self): return idaapi.PLUGIN_OK
    def run(self, arg): pass
    def term(self): pass


def PLUGIN_ENTRY():
    main()
    return dummy()


if __name__ == "__main__":
    PLUGIN_ENTRY()


repository: sunyymq/windows.cracking.tools
file: IDA_Pro_v6.8/plugins/keypatch.py
description: windows crackings tools
wanted_name: Keypatch Patcher
comment: Keypatch plugin for IDA Pro (using Keystone framework)
created_at: 2018-03-06 11:32:00+00:00
pushed_at: 2017-11-22 02:33:01+00:00
url:
https://github.com/sunyymq/windows.cracking.tools/blob/4cfafb4bf9c82ad8a8975d3bedbe49f4df1207ca/IDA_Pro_v6.8/plugins/keypatch.py

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

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

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

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

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

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

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


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


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

X86_NOP = "\x90"

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

# save all the info on patching
patch_info = []


def to_hexstr(buf, sep=' '):
    return sep.join("{0:02x}".format(ord(c)) for c in buf).upper()

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

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

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

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

# download a file from @url, then return (result, file-content)
# return (0, content) on success, or ({1|2}, None) on download failure
def url_download(url):
    from urllib2 import Request, urlopen, URLError, HTTPError

    # create the url and the request
    req = Request(url)

    # Open the url
    try:
        # download this URL
        f = urlopen(req)
        content = f.read()
        return (0, content)

    # handle errors
    except HTTPError, e:
        # print "HTTP Error:", e.code , url
        # fail to download this file
        return (1, None)
    except URLError, e:
        # print "URL Error:", e.reason , url
        # fail to download this file
        return (1, None)
    except Exception as e:
        # fail to save the downloaded file
        # print("Error:", e)
        return (2, None)


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

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

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

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

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

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

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

        # heuristically detect hardware setup
        info = idaapi.get_inf_structure()
        if idaapi.IDA_SDK_VERSION >= 700:
            # IDA >= 7.0 case change
            cpuname = info.procname.lower()
        else:
            cpuname = info.procName.lower()
        #print("Keypatch MF = %s" %idaapi.cvar.inf.mf)
        if cpuname == "metapc":
            arch = KS_ARCH_X86
            if info.is_64bit():
                mode = KS_MODE_64
            elif info.is_32bit():
                mode = KS_MODE_32
            else:
                mode = KS_MODE_16
        elif cpuname.startswith("arm"):
            # ARM or ARM64
            if info.is_64bit():
                arch = KS_ARCH_ARM64
                if idaapi.cvar.inf.mf:
                    mode = KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_LITTLE_ENDIAN
            else:
                arch = KS_ARCH_ARM
                # either big-endian or little-endian
                if idaapi.cvar.inf.mf:
                    mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("sparc"):
            arch = KS_ARCH_SPARC
            if info.is_64bit():
                mode = KS_MODE_SPARC64
            else:
                mode = KS_MODE_SPARC32
            if idaapi.cvar.inf.mf:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("ppc"):
            arch = KS_ARCH_PPC
            if info.is_64bit():
                mode = KS_MODE_PPC64
            else:
                mode = KS_MODE_PPC32
            if cpuname == "ppc":
                # do not support Little Endian mode for PPC
                mode += KS_MODE_BIG_ENDIAN
        elif cpuname.startswith("mips"):
            arch = KS_ARCH_MIPS
            if info.is_64bit():
                mode = KS_MODE_MIPS64
            else:
                mode = KS_MODE_MIPS32
            if idaapi.cvar.inf.mf:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
            arch = KS_ARCH_SYSTEMZ
            mode = KS_MODE_BIG_ENDIAN

        return (arch, mode)

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

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

        return text.strip()

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

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

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

                sym = name

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

                (t, v) = idaapi.get_name_value(address, sym)

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

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

            return _op

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

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

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

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

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

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

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

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

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

        len = idc.ItemSize(address)
        item = idaapi.get_many_bytes(address, len)

        if item is None:
            return (None, 0)

        if hex_output:
            item = to_hexstr(item)

        return (item, len)

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

        dtype = idaapi.cmd.Operands[op_idx].dtyp
        dtyp_size = idaapi.get_dtyp_size(dtype)
        if dtype == idaapi.dt_tbyte:
            if dtyp_size == 10:
                return 'xword'

        dtyp_name = dtyp_lists.get(idaapi.cmd.Operands[op_idx].dtyp, None)

        return dtyp_name

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

        return codes

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

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

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

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

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

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

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

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

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

                o = ''.join(parts)

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

            opers.append(o)
            i += 1

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

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

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

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

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

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

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

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

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

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

                return assembly

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

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

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

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

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

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

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

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

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

        def is_thumb(address):
            return idc.GetReg(address, 'T') == 1

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

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

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

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

        return (encoding, count)


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

        while ea < (address + len):

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

            orig_byte = idc.Byte(ea)
            orig_data += chr(orig_byte)
            patch_byte = ord(patch_data[ea - address])

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

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

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

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

            return (None, None)

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

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

        return (patched_len, orig_data)

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

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

        orig_comment = idc.Comment(address)
        if orig_comment == None:
            orig_comment = ''

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

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

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

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

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

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

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

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

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

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

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

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

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

        return plen

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

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

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

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

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

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

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

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

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

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

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

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

        return plen


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

        return (keys, values)

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

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

        return val

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

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

        return idx

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

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

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

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


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

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

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

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

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

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

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

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

    def __init__(self, kp_asm, address, assembly=None, patch_mode=False, opts=0):
        pass

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

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

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

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

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

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

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

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

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

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

        return 1

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

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

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

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

        return 1

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

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

        return opts


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

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

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

        self.Compile()

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

        return self.update_patchform(fid)


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

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

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

        self.Compile()

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

        return self.update_patchform(fid)


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

    def OnClose(self):
        return

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

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

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

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


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

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

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

        self.Compile()

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

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

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

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

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

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

        return 1


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

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

            Keypatch is released under the GPL v2.
            Find more info at http://www.keystone-engine.org/keypatch
            """ %(version, keystone.__version__), {
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(KP_HOMEPAGE, new = 2)

        return 1


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

            {FormChangeCb}
            Your Keypatch is v%s
            %s
            """ %(version, message), {
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(KP_HOMEPAGE, new = 2)

        return 1


try:
    # adapted from pull request #7 by @quangnh89
    class Kp_Menu_Context(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

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

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

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

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

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

        @classmethod
        def update(self, ctx):
            try:
                if ctx.form_type == idaapi.BWN_DISASM:
                    return idaapi.AST_ENABLE_FOR_FORM
                else:
                    return idaapi.AST_DISABLE_FOR_FORM
            except:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS
            
    # context menu for Patcher
    class Kp_MC_Patcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.patcher()
            return 1

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

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

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

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

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


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


# check if we already initialized Keypatch
kp_initialized = False

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


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

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

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

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

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

    def init(self):
        global kp_initialized

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

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

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

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

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

            self.load_configuration()

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

        return idaapi.PLUGIN_KEEP

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

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

    # null handler
    def menu_null(self):
        pass

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

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

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

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

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

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

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

        if idaapi.IDA_SDK_VERSION >= 700:
            # IDA >= 7.0 read selection change
            p0 = idaapi.twinpos_t()
            p1 = idaapi.twinpos_t()
            view = idaapi.get_current_viewer()
            selection = idaapi.read_selection(view, p0, p1)
            if selection:
                place0  = p0.place(view)
                place1 = p1.place(view)

                addr_begin  = place0.toea()
                addr_end = place1.toea()
        else:
            selection, addr_begin, addr_end = idaapi.read_selection()

        if selection:
            # call Fill Range function on this selected code
            return self.fill_range()

        address = idc.ScreenEA()

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

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

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

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

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

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

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

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

        if idaapi.IDA_SDK_VERSION >= 700:
            # IDA >= 7.0 read selection change
            p0 = idaapi.twinpos_t()
            p1 = idaapi.twinpos_t()
            view = idaapi.get_current_viewer()
            selection = idaapi.read_selection(view, p0, p1)

            place0  = p0.place(view)
            place1 = p1.place(view)

            if selection:
                place0  = p0.place(view)
                place1 = p1.place(view)
    
                addr_begin  = place0.toea()
                addr_end = place1.toea()
        else:
            selection, addr_begin, addr_end = idaapi.read_selection()

        if not selection:
            idc.Warning("ERROR: Keypatch requires a range to be selected for fill in, try again")
            return

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

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

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

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

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

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

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

        # free this form
        f.Free()

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


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


repository: youshi01/ctf-tools
file: IDA7.0/plugins/keypatch.py
wanted_name: Keypatch Patcher
comment: Keypatch plugin for IDA Pro (using Keystone framework)
created_at: 2023-11-07 02:13:17+00:00
pushed_at: 2023-11-06 14:27:14+00:00
url:
https://github.com/youshi01/ctf-tools/blob/567850d908e287536411ea4ba6b751b51e7503d0/IDA7.0/plugins/keypatch.py

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

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

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

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

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

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

import os
import re
import json
from keystone import *
import idc
import idaapi
from idc import GetOpType, GetOpnd, ItemEnd

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


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

X86_NOP = "\x90"

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

# save all the info on patching
patch_info = []


def to_hexstr(buf, sep=' '):
    return sep.join("{0:02x}".format(ord(c)) for c in buf).upper()


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

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

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

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


# download a file from @url, then return (result, file-content)
# return (0, content) on success, or ({1|2}, None) on download failure
def url_download(url):
    from urllib2 import Request, urlopen, URLError, HTTPError

    # create the url and the request
    req = Request(url)

    # Open the url
    try:
        # download this URL
        f = urlopen(req)
        content = f.read()
        return (0, content)

    # handle errors
    except HTTPError, e:
        # print "HTTP Error:", e.code , url
        # fail to download this file
        return (1, None)
    except URLError, e:
        # print "URL Error:", e.reason , url
        # fail to download this file
        return (1, None)
    except Exception as e:
        # fail to save the downloaded file
        # print("Error:", e)
        return (2, None)


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

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

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

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

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

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

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

        # heuristically detect hardware setup
        info = idaapi.get_inf_structure()
        
        try:
            cpuname = info.procname.lower()
        except:
            cpuname = info.procName.lower()

        try:
            # since IDA7 beta 3 (170724) renamed inf.mf -> is_be()/set_be()
            is_be = idaapi.cvar.inf.is_be()
        except:
            # older IDA versions
            is_be = idaapi.cvar.inf.mf
        # print("Keypatch BIG_ENDIAN = %s" %is_be)
        
        if cpuname == "metapc":
            arch = KS_ARCH_X86
            if info.is_64bit():
                mode = KS_MODE_64
            elif info.is_32bit():
                mode = KS_MODE_32
            else:
                mode = KS_MODE_16
        elif cpuname.startswith("arm"):
            # ARM or ARM64
            if info.is_64bit():
                arch = KS_ARCH_ARM64
                if is_be:
                    mode = KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_LITTLE_ENDIAN
            else:
                arch = KS_ARCH_ARM
                # either big-endian or little-endian
                if is_be:
                    mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("sparc"):
            arch = KS_ARCH_SPARC
            if info.is_64bit():
                mode = KS_MODE_SPARC64
            else:
                mode = KS_MODE_SPARC32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("ppc"):
            arch = KS_ARCH_PPC
            if info.is_64bit():
                mode = KS_MODE_PPC64
            else:
                mode = KS_MODE_PPC32
            if cpuname == "ppc":
                # do not support Little Endian mode for PPC
                mode += KS_MODE_BIG_ENDIAN
        elif cpuname.startswith("mips"):
            arch = KS_ARCH_MIPS
            if info.is_64bit():
                mode = KS_MODE_MIPS64
            else:
                mode = KS_MODE_MIPS32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
            arch = KS_ARCH_SYSTEMZ
            mode = KS_MODE_BIG_ENDIAN

        return (arch, mode)

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

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

        return text.strip()

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

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

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

                sym = name

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

                (t, v) = idaapi.get_name_value(address, sym)

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

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

            return _op

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

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

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

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

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

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

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

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

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

        len = idc.ItemSize(address)
        item = idaapi.get_many_bytes(address, len)

        if item is None:
            return (None, 0)

        if hex_output:
            item = to_hexstr(item)

        return (item, len)

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

        dtype = idaapi.cmd.Operands[op_idx].dtyp
        dtyp_size = idaapi.get_dtyp_size(dtype)
        if dtype == idaapi.dt_tbyte:
            if dtyp_size == 10:
                return 'xword'

        dtyp_name = dtyp_lists.get(idaapi.cmd.Operands[op_idx].dtyp, None)

        return dtyp_name

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

        return codes

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

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

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

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

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

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

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

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

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

                o = ''.join(parts)

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

            opers.append(o)
            i += 1

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

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

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

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

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

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

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

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

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

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

                return assembly

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

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

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

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

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

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

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

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

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

        def is_thumb(address):
            return idc.GetReg(address, 'T') == 1

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

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

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

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

        return (encoding, count)


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

        while ea < (address + len):

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

            orig_byte = idc.Byte(ea)
            orig_data += chr(orig_byte)
            patch_byte = ord(patch_data[ea - address])

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

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

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

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

            return (None, None)

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

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

        return (patched_len, orig_data)

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

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

        orig_comment = idc.Comment(address)
        if orig_comment is None:
            orig_comment = ''

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

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

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

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

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

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

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

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

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

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

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

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

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

        return plen

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

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

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

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

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

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

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

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

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

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

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

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

        return plen


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

        return (keys, values)

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

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

        return val

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

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

        return idx

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return 1

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

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

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

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

        return 1

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

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

        return opts


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

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

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

        self.Compile()

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

        return self.update_patchform(fid)


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

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

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

        self.Compile()

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

        return self.update_patchform(fid)


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

    def OnClose(self):
        return

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

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

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

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


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

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

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

        self.Compile()

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

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

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

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

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

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

        return 1


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

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

            Keypatch is released under the GPL v2.
            Find more info at http://www.keystone-engine.org/keypatch
            """ %(version, keystone.__version__), {
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(KP_HOMEPAGE, new = 2)

        return 1


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

            {FormChangeCb}
            Your Keypatch is v%s
            %s
            """ %(version, message), {
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })
        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(KP_HOMEPAGE, new = 2)

        return 1


try:
    # adapted from pull request #7 by @quangnh89
    class Kp_Menu_Context(idaapi.action_handler_t):

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

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

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

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

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

        @classmethod
        def update(self, ctx):
            try:
                if ctx.form_type == idaapi.BWN_DISASM:
                    return idaapi.AST_ENABLE_FOR_FORM
                else:
                    return idaapi.AST_DISABLE_FOR_FORM
            except:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS
            
    # context menu for Patcher
    class Kp_MC_Patcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.patcher()
            return 1

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

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

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

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

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


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


# check if we already initialized Keypatch
kp_initialized = False

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

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

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

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

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

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

    def init(self):
        global kp_initialized

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

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

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

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

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

            self.load_configuration()

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

        return idaapi.PLUGIN_KEEP

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

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

    # null handler
    def menu_null(self):
        pass

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

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

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

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

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

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

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

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

        address = idc.ScreenEA()

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

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

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

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

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

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

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

    # handler for Fill Range menu
    def fill_range(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            idc.Warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return
               
        selection, addr_begin, addr_end = idaapi.read_selection()
        if not selection:
            idc.Warning("ERROR: Keypatch requires a range to be selected for fill in, try again")
            return

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

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

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

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

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

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

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

        # free this form
        f.Free()

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


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


repository: Ga-ryo/IDAObjectViewer-deprecated-
file: ida_object_viewer.py
description: Graphical object viewer for IDA Pro.(Unmaintained and not completed. Use watch-view in IDA)
wanted_name: IDAObjectViewer
comment: Graphical object viewer for IDA
created_at: 2018-04-16 09:51:29+00:00
pushed_at: 2018-04-25 17:19:55+00:00
url:
https://github.com/Ga-ryo/IDAObjectViewer-deprecated-/blob/ddf82da1326757e52f9f2eb443df7135b4014df7/ida_object_viewer.py

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

nodz = None
max_depth = 2 #TEMP
#TODO hook debugger activate event?
dbg_active = False
"""
struct struc_1 -> struc_1
struc_1 ->  struc_1
"""
pattern = '^[ ]*(struct[ ]+)?'
pat = re.compile(pattern)
bits = 0
endian = 'little'


############################################################################################################################
##############################################   nodz_util.py   ############################################################
############################################################################################################################

def convertDataToColor(data=None, alternate=False, av=20):
    """
    Convert a list of 3 (rgb) or 4(rgba) values from the configuration
    file into a QColor.

    :param data: Input color.
    :type  data: List.

    :param alternate: Whether or not this is an alternate color.
    :type  alternate: Bool.

    :param av: Alternate value.
    :type  av: Int.

    """
    # rgb
    if len(data) == 3:
        color = QtGui.QColor(data[0], data[1], data[2])
        if alternate:
            mult = generateAlternateColorMultiplier(color, av)


            color = QtGui.QColor(data[0]-(av*mult), data[1]-(av*mult), data[2]-(av*mult))
        return color

    # rgba
    elif len(data) == 4:
        color = QtGui.QColor(data[0], data[1], data[2], data[3])
        if alternate:
            mult = generateAlternateColorMultiplier(color, av)
            color = QtGui.QColor(data[0]-(av*mult), data[1]-(av*mult), data[2]-(av*mult), data[3])
        return color

    # wrong
    else:
        print 'Color from configuration is not recognized : ', data
        print 'Can only be [R, G, B] or [R, G, B, A]'
        print 'Using default color !'
        color = QtGui.QColor(120, 120, 120)
        if alternate:
            color = QtGui.QColor(120-av, 120-av, 120-av)
        return color

def generateAlternateColorMultiplier(color, av):
    """
    Generate a multiplier based on the input color lighness to increase
    the alternate value for dark color or reduce it for bright colors.

    :param color: Input color.
    :type  color: QColor.

    :param av: Alternate value.
    :type  av: Int.

    """
    lightness = color.lightness()
    mult = float(lightness)/255

    return mult

def createPointerBoundingBox(pointerPos, bbSize):
    """
    generate a bounding box around the pointer.

    :param pointerPos: Pointer position.
    :type  pointerPos: QPoint.

    :param bbSize: Width and Height of the bounding box.
    :type  bbSize: Int.

    """
    # Create pointer's bounding box.
    point = pointerPos

    mbbPos = point
    point.setX(point.x() - bbSize / 2)
    point.setY(point.y() - bbSize / 2)

    size = QtCore.QSize(bbSize, bbSize)
    bb = QtCore.QRect(mbbPos, size)
    bb = QtCore.QRectF(bb)

    return bb

def swapListIndices(inputList, oldIndex, newIndex):
    """
    Simply swap 2 indices in a the specified list.

    :param inputList: List that contains the elements to swap.
    :type  inputList: List.

    :param oldIndex: Index of the element to move.
    :type  oldIndex: Int.

    :param newIndex: Destination index of the element.
    :type  newIndex: Int.

    """
    if oldIndex == -1:
        oldIndex = len(inputList)-1


    if newIndex == -1:
        newIndex = len(inputList)

    value = inputList[oldIndex]
    inputList.pop(oldIndex)
    inputList.insert(newIndex, value)

############################################################################################################################
############################################################################################################################
############################################################################################################################


############################################################################################################################
##############################################   nodz_main.py   ############################################################
############################################################################################################################

class Nodz(QtWidgets.QGraphicsView):

    """
    The main view for the node graph representation.

    The node view implements a state pattern to control all the
    different user interactions.

    """

    signal_NodeCreated = QtCore.pyqtSignal(object)
    signal_NodeDeleted = QtCore.pyqtSignal(object)
    signal_NodeEdited = QtCore.pyqtSignal(object, object)
    signal_NodeSelected = QtCore.pyqtSignal(object)
    signal_NodeMoved = QtCore.pyqtSignal(str, object)

    signal_AttrCreated = QtCore.pyqtSignal(object, object)
    signal_AttrDeleted = QtCore.pyqtSignal(object, object)
    signal_AttrEdited = QtCore.pyqtSignal(object, object, object)

    signal_PlugConnected = QtCore.pyqtSignal(object, object, object, object)
    signal_PlugDisconnected = QtCore.pyqtSignal(object, object, object, object)
    signal_SocketConnected = QtCore.pyqtSignal(object, object, object, object)
    signal_SocketDisconnected = QtCore.pyqtSignal(object, object, object, object)

    signal_GraphSaved = QtCore.pyqtSignal()
    signal_GraphLoaded = QtCore.pyqtSignal()
    signal_GraphCleared = QtCore.pyqtSignal()
    signal_GraphEvaluated = QtCore.pyqtSignal()

    signal_KeyPressed = QtCore.pyqtSignal(object)
    signal_Dropped = QtCore.pyqtSignal()

    def __init__(self, parent, config_s):
        """
        Initialize the graphics view.

        """
        super(Nodz, self).__init__(parent)

        # Load nodz configuration.
        self.loadConfig(config_s)

        # General data.
        self.gridVisToggle = True
        self.gridSnapToggle = False
        self._nodeSnap = False
        self.selectedNodes = None

        # Connections data.
        self.drawingConnection = False
        self.currentHoveredNode = None
        self.sourceSlot = None

        # Display options.
        self.currentState = 'DEFAULT'
        self.pressedKeys = list()
        self.setDragMode(QtWidgets.QGraphicsView.ScrollHandDrag)

    def wheelEvent(self, event):
        """
        Zoom in the view with the mouse wheel.

        """
        self.currentState = 'ZOOM_VIEW'
        self.setTransformationAnchor(QtWidgets.QGraphicsView.AnchorUnderMouse)

        inFactor = 1.15
        outFactor = 1 / inFactor

        if event.delta() > 0:
            zoomFactor = inFactor
        else:
            zoomFactor = outFactor

        self.scale(zoomFactor, zoomFactor)
        self.currentState = 'DEFAULT'

    def mousePressEvent(self, event):
        """
        Initialize tablet zoom, drag canvas and the selection.

        """
        # Tablet zoom
        if (event.button() == QtCore.Qt.RightButton and
            event.modifiers() == QtCore.Qt.AltModifier):
            self.currentState = 'ZOOM_VIEW'
            self.initMousePos = event.pos()
            self.zoomInitialPos = event.pos()
            self.initMouse = QtGui.QCursor.pos()
            self.setInteractive(False)


        # Drag view
        elif (event.button() == QtCore.Qt.MiddleButton and
              event.modifiers() == QtCore.Qt.AltModifier):
            self.currentState = 'DRAG_VIEW'
            self.prevPos = event.pos()
            self.setCursor(QtCore.Qt.ClosedHandCursor)
            self.setInteractive(False)


        # Rubber band selection
        elif (event.button() == QtCore.Qt.LeftButton and
              event.modifiers() == QtCore.Qt.NoModifier and
              self.scene().itemAt(self.mapToScene(event.pos()), QtGui.QTransform()) is None):
            self.currentState = 'DRAG_WINDOW'
            #self._initRubberband(event.pos())
            self.setInteractive(False)


        # Drag Item
        elif (event.button() == QtCore.Qt.LeftButton and
              event.modifiers() == QtCore.Qt.NoModifier and
              self.scene().itemAt(self.mapToScene(event.pos()), QtGui.QTransform()) is not None):
            self.currentState = 'DRAG_ITEM'
            self.setInteractive(True)


        # Add selection
        elif (event.button() == QtCore.Qt.LeftButton and
              QtCore.Qt.Key_Shift in self.pressedKeys and
              QtCore.Qt.Key_Control in self.pressedKeys):
            self.currentState = 'ADD_SELECTION'
            self._initRubberband(event.pos())
            self.setInteractive(False)


        # Subtract selection
        elif (event.button() == QtCore.Qt.LeftButton and
              event.modifiers() == QtCore.Qt.ControlModifier):
            self.currentState = 'SUBTRACT_SELECTION'
            self._initRubberband(event.pos())
            self.setInteractive(False)


        # Toggle selection
        elif (event.button() == QtCore.Qt.LeftButton and
              event.modifiers() == QtCore.Qt.ShiftModifier):
            self.currentState = 'TOGGLE_SELECTION'
            self._initRubberband(event.pos())
            self.setInteractive(False)


        else:
            self.currentState = 'DEFAULT'

        super(Nodz, self).mousePressEvent(event)

    def mouseMoveEvent(self, event):
        """
        Update tablet zoom, canvas dragging and selection.

        """
        # Zoom.
        if self.currentState == 'ZOOM_VIEW':
            offset = self.zoomInitialPos.x() - event.pos().x()

            if offset > self.previousMouseOffset:
                self.previousMouseOffset = offset
                self.zoomDirection = -1
                self.zoomIncr -= 1

            elif offset == self.previousMouseOffset:
                self.previousMouseOffset = offset
                if self.zoomDirection == -1:
                    self.zoomDirection = -1
                else:
                    self.zoomDirection = 1

            else:
                self.previousMouseOffset = offset
                self.zoomDirection = 1
                self.zoomIncr += 1

            if self.zoomDirection == 1:
                zoomFactor = 1.03
            else:
                zoomFactor = 1 / 1.03

            # Perform zoom and re-center on initial click position.
            pBefore = self.mapToScene(self.initMousePos)
            self.setTransformationAnchor(QtWidgets.QGraphicsView.AnchorViewCenter)
            self.scale(zoomFactor, zoomFactor)
            pAfter = self.mapToScene(self.initMousePos)
            diff = pAfter - pBefore

            self.setTransformationAnchor(QtWidgets.QGraphicsView.NoAnchor)
            self.translate(diff.x(), diff.y())

        # Drag canvas.
        elif self.currentState == 'DRAG_VIEW':
            offset = self.prevPos - event.pos()
            self.prevPos = event.pos()
            self.verticalScrollBar().setValue(self.verticalScrollBar().value() + offset.y())
            self.horizontalScrollBar().setValue(self.horizontalScrollBar().value() + offset.x())

        # RuberBand selection.
        elif (self.currentState == 'SELECTION' or
              self.currentState == 'ADD_SELECTION' or
              self.currentState == 'SUBTRACT_SELECTION' or
              self.currentState == 'TOGGLE_SELECTION'):
            self.rubberband.setGeometry(QtCore.QRect(self.origin, event.pos()).normalized())

        super(Nodz, self).mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        """
        Apply tablet zoom, dragging and selection.

        """
        # Zoom the View.
        if self.currentState == '.ZOOM_VIEW':
            self.offset = 0
            self.zoomDirection = 0
            self.zoomIncr = 0
            self.setInteractive(True)


        # Drag View.
        elif self.currentState == 'DRAG_VIEW':
            self.setCursor(QtCore.Qt.ArrowCursor)
            self.setInteractive(True)

        elif self.currentState == 'DRAG_WINDOW':
            self.setInteractive(True)

        # Selection.
        elif self.currentState == 'SELECTION':
            self.rubberband.setGeometry(QtCore.QRect(self.origin,
                                                     event.pos()).normalized())
            painterPath = self._releaseRubberband()
            self.setInteractive(True)
            self.scene().setSelectionArea(painterPath)


        # Add Selection.
        elif self.currentState == 'ADD_SELECTION':
            self.rubberband.setGeometry(QtCore.QRect(self.origin,
                                                     event.pos()).normalized())
            painterPath = self._releaseRubberband()
            self.setInteractive(True)
            for item in self.scene().items(painterPath):
                item.setSelected(True)


        # Subtract Selection.
        elif self.currentState == 'SUBTRACT_SELECTION':
            self.rubberband.setGeometry(QtCore.QRect(self.origin,
                                                     event.pos()).normalized())
            painterPath = self._releaseRubberband()
            self.setInteractive(True)
            for item in self.scene().items(painterPath):
                item.setSelected(False)


        # Toggle Selection
        elif self.currentState == 'TOGGLE_SELECTION':
            self.rubberband.setGeometry(QtCore.QRect(self.origin,
                                                     event.pos()).normalized())
            painterPath = self._releaseRubberband()
            self.setInteractive(True)
            for item in self.scene().items(painterPath):
                if item.isSelected():
                    item.setSelected(False)
                else:
                    item.setSelected(True)

        self.currentState = 'DEFAULT'

        super(Nodz, self).mouseReleaseEvent(event)

    def keyPressEvent(self, event):
        """
        Save pressed key and apply shortcuts.

        Shortcuts are:
        DEL - Delete the selected nodes
        F - Focus view on the selection

        """
        if event.key() not in self.pressedKeys:
            self.pressedKeys.append(event.key())

        if event.key() == QtCore.Qt.Key_Delete:
            self._deleteSelectedNodes()

        if event.key() == QtCore.Qt.Key_F:
            self._focus()

        if event.key() == QtCore.Qt.Key_S:
            self._nodeSnap = True

        # Emit signal.
        self.signal_KeyPressed.emit(event.key())

    def keyReleaseEvent(self, event):
        """
        Clear the key from the pressed key list.

        """
        if event.key() == QtCore.Qt.Key_S:
            self._nodeSnap = False

        if event.key() in self.pressedKeys:
            self.pressedKeys.remove(event.key())

    def _initRubberband(self, position):
        """
        Initialize the rubber band at the given position.

        """
        self.rubberBandStart = position
        self.origin = position
        self.rubberband.setGeometry(QtCore.QRect(self.origin, QtCore.QSize()))
        self.rubberband.show()

    def _releaseRubberband(self):
        """
        Hide the rubber band and return the path.

        """
        painterPath = QtGui.QPainterPath()
        rect = self.mapToScene(self.rubberband.geometry())
        painterPath.addPolygon(rect)
        self.rubberband.hide()
        return painterPath

    def _focus(self):
        """
        Center on selected nodes or all of them if no active selection.

        """
        if self.scene().selectedItems():
            itemsArea = self._getSelectionBoundingbox()
            self.fitInView(itemsArea, QtCore.Qt.KeepAspectRatio)
        else:
            itemsArea = self.scene().itemsBoundingRect()
            self.fitInView(itemsArea, QtCore.Qt.KeepAspectRatio)

    def _getSelectionBoundingbox(self):
        """
        Return the bounding box of the selection.

        """
        bbx_min = None
        bbx_max = None
        bby_min = None
        bby_max = None
        bbw = 0
        bbh = 0
        for item in self.scene().selectedItems():
            pos = item.scenePos()
            x = pos.x()
            y = pos.y()
            w = x + item.boundingRect().width()
            h = y + item.boundingRect().height()

            # bbx min
            if bbx_min is None:
                bbx_min = x
            elif x < bbx_min:
                bbx_min = x
            # end if

            # bbx max
            if bbx_max is None:
                bbx_max = w
            elif w > bbx_max:
                bbx_max = w
            # end if

            # bby min
            if bby_min is None:
                bby_min = y
            elif y < bby_min:
                bby_min = y
            # end if

            # bby max
            if bby_max is None:
                bby_max = h
            elif h > bby_max:
                bby_max = h
            # end if
        # end if
        bbw = bbx_max - bbx_min
        bbh = bby_max - bby_min
        return QtCore.QRectF(QtCore.QRect(bbx_min, bby_min, bbw, bbh))

    def _deleteSelectedNodes(self):
        """
        Delete selected nodes.

        """
        selected_nodes = list()
        for node in self.scene().selectedItems():
            selected_nodes.append(node.name)
            node._remove()

        # Emit signal.
        self.signal_NodeDeleted.emit(selected_nodes)

    def _returnSelection(self):
        """
        Wrapper to return selected items.

        """
        selected_nodes = list()
        if self.scene().selectedItems():
            for node in self.scene().selectedItems():
                selected_nodes.append(node.name)

        # Emit signal.
        self.signal_NodeSelected.emit(selected_nodes)


    ##################################################################
    # API
    ##################################################################

    def loadConfig(self, d):
        """
        Set a specific configuration for this instance of Nodz.

        :type  filePath: str.
        :param filePath: The path to the config file that you want to
                         use.

        """
        self.config = d

    def initialize(self):
        """
        Setup the view's behavior.

        """
        # Setup view.
        config = self.config
        self.setRenderHint(QtGui.QPainter.Antialiasing, config['antialiasing'])
        self.setRenderHint(QtGui.QPainter.TextAntialiasing, config['antialiasing'])
        self.setRenderHint(QtGui.QPainter.HighQualityAntialiasing, config['antialiasing_boost'])
        self.setRenderHint(QtGui.QPainter.SmoothPixmapTransform, config['smooth_pixmap'])
        self.setRenderHint(QtGui.QPainter.NonCosmeticDefaultPen, True)
        self.setViewportUpdateMode(QtWidgets.QGraphicsView.FullViewportUpdate)
        self.setTransformationAnchor(QtWidgets.QGraphicsView.AnchorUnderMouse)
        self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAlwaysOff)
        self.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarAlwaysOff)
        self.rubberband = QtWidgets.QRubberBand(QtWidgets.QRubberBand.Rectangle, self)

        # Setup scene.
        scene = NodeScene(self)

        sceneWidth = config['scene_width']
        sceneHeight = config['scene_height']
        scene.setSceneRect(0, 0, sceneWidth, sceneHeight)
        self.setScene(scene)
        # Connect scene node moved signal
        scene.signal_NodeMoved.connect(self.signal_NodeMoved)

        # Tablet zoom.
        self.previousMouseOffset = 0
        self.zoomDirection = 0
        self.zoomIncr = 0

        # Connect signals.
        self.scene().selectionChanged.connect(self._returnSelection)


    # NODES
    def createNode(self, name='default', preset='node_default', position=None, alternate=True):
        """
        Create a new node with a given name, position and color.

        :type  name: str.
        :param name: The name of the node. The name has to be unique
                     as it is used as a key to store the node object.

        :type  preset: str.
        :param preset: The name of graphical preset in the config file.

        :type  position: QtCore.QPoint.
        :param position: The position of the node once created. If None,
                         it will be created at the center of the scene.

        :type  alternate: bool.
        :param alternate: The attribute color alternate state, if True,
                          every 2 attribute the color will be slightly
                          darker.

        :return : The created node

        """
        # Check for name clashes
        if name in self.scene().nodes.keys():
            print 'A node with the same name already exists : {0}'.format(name)
            print 'Node creation aborted !'
            return
        else:
            nodeItem = NodeItem(name=name, alternate=alternate, preset=preset,
                                config=self.config)

            # Store node in scene.
            self.scene().nodes[name] = nodeItem

            if not position:
                # Get the center of the view.
                p = self.viewport().rect().topLeft()
                p.setX(p.x() + 50)
                p.setY(p.y() + 50)
                position = self.mapToScene(p)

            # Set node position.
            self.scene().addItem(nodeItem)
            #nodeItem.setPos(position - nodeItem.nodeCenter)
            nodeItem.setPos(position)

            # Emit signal.
            self.signal_NodeCreated.emit(name)

            return nodeItem

    def deleteNode(self, node):
        """
        Delete the specified node from the view.

        :type  node: class.
        :param node: The node instance that you want to delete.

        """
        if not node in self.scene().nodes.values():
            print 'Node object does not exist !'
            print 'Node deletion aborted !'
            return

        if node in self.scene().nodes.values():
            nodeName = node.name
            node._remove()

            # Emit signal.
            self.signal_NodeDeleted.emit([nodeName])

    def editNode(self, node, newName=None):
        """
        Rename an existing node.

        :type  node: class.
        :param node: The node instance that you want to delete.

        :type  newName: str.
        :param newName: The new name for the given node.

        """
        if not node in self.scene().nodes.values():
            print 'Node object does not exist !'
            print 'Node edition aborted !'
            return

        oldName = node.name

        if newName != None:
            # Check for name clashes
            if newName in self.scene().nodes.keys():
                print 'A node with the same name already exists : {0}'.format(newName)
                print 'Node edition aborted !'
                return
            else:
                node.name = newName

        # Replace node data.
        self.scene().nodes[newName] = self.scene().nodes[oldName]
        self.scene().nodes.pop(oldName)

        # Store new node name in the connections
        if node.sockets:
            for socket in node.sockets.values():
                for connection in socket.connections:
                    connection.socketNode = newName

        if node.plugs:
            for plug in node.plugs.values():
                for connection in plug.connections:
                    connection.plugNode = newName

        node.update()

        # Emit signal.
        self.signal_NodeEdited.emit(oldName, newName)


    # ATTRS
    def createAttribute(self, node, name='default', index=-1, preset='attr_default', plug=True, socket=True, dataType=None):
        """
        Create a new attribute with a given name.

        :type  node: class.
        :param node: The node instance that you want to delete.

        :type  name: str.
        :param name: The name of the attribute. The name has to be
                     unique as it is used as a key to store the node
                     object.

        :type  index: int.
        :param index: The index of the attribute in the node.

        :type  preset: str.
        :param preset: The name of graphical preset in the config file.

        :type  plug: bool.
        :param plug: Whether or not this attribute can emit connections.

        :type  socket: bool.
        :param socket: Whether or not this attribute can receive
                       connections.

        :type  dataType: type.
        :param dataType: Type of the data represented by this attribute
                         in order to highlight attributes of the same
                         type while performing a connection.

        """
        if not node in self.scene().nodes.values():
            print 'Node object does not exist !'
            print 'Attribute creation aborted !'
            return

        if name in node.attrs:
            print 'An attribute with the same name already exists : {0}'.format(name)
            print 'Attribute creation aborted !'
            return

        node._createAttribute(name=name, index=index, preset=preset, plug=plug, socket=socket, dataType=dataType)

        # Emit signal.
        self.signal_AttrCreated.emit(node.name, index)

    def deleteAttribute(self, node, index):
        """
        Delete the specified attribute.

        :type  node: class.
        :param node: The node instance that you want to delete.

        :type  index: int.
        :param index: The index of the attribute in the node.

        """
        if not node in self.scene().nodes.values():
            print 'Node object does not exist !'
            print 'Attribute deletion aborted !'
            return

        node._deleteAttribute(index)

        # Emit signal.
        self.signal_AttrDeleted.emit(node.name, index)

    def editAttribute(self, node, index, newName=None, newIndex=None):
        """
        Edit the specified attribute.

        :type  node: class.
        :param node: The node instance that you want to delete.

        :type  index: int.
        :param index: The index of the attribute in the node.

        :type  newName: str.
        :param newName: The new name for the given attribute.

        :type  newIndex: int.
        :param newIndex: The index for the given attribute.

        """
        if not node in self.scene().nodes.values():
            print 'Node object does not exist !'
            print 'Attribute creation aborted !'
            return

        if newName != None:
            if newName in node.attrs:
                print 'An attribute with the same name already exists : {0}'.format(newName)
                print 'Attribute edition aborted !'
                return
            else:
                oldName = node.attrs[index]

            # Rename in the slot item(s).
            if node.attrsData[oldName]['plug']:
                node.plugs[oldName].attribute = newName
                node.plugs[newName] = node.plugs[oldName]
                node.plugs.pop(oldName)
                for connection in node.plugs[newName].connections:
                    connection.plugAttr = newName

            if node.attrsData[oldName]['socket']:
                node.sockets[oldName].attribute = newName
                node.sockets[newName] = node.sockets[oldName]
                node.sockets.pop(oldName)
                for connection in node.sockets[newName].connections:
                    connection.socketAttr = newName

            # Replace attribute data.
            node.attrsData[oldName]['name'] = newName
            node.attrsData[newName] = node.attrsData[oldName]
            node.attrsData.pop(oldName)
            node.attrs[index] = newName

        if isinstance(newIndex, int):
            attrName = node.attrs[index]

            swapListIndices(node.attrs, index, newIndex)

            # Refresh connections.
            for plug in node.plugs.values():
                plug.update()
                if plug.connections:
                    for connection in plug.connections:
                        if isinstance(connection.source, PlugItem):
                            connection.source = plug
                            connection.source_point = plug.center()
                        else:
                            connection.target = plug
                            connection.target_point = plug.center()
                        if newName:
                            connection.plugAttr = newName
                        connection.updatePath()

            for socket in node.sockets.values():
                socket.update()
                if socket.connections:
                    for connection in socket.connections:
                        if isinstance(connection.source, SocketItem):
                            connection.source = socket
                            connection.source_point = socket.center()
                        else:
                            connection.target = socket
                            connection.target_point = socket.center()
                        if newName:
                            connection.socketAttr = newName
                        connection.updatePath()

            self.scene().update()

        node.update()

        # Emit signal.
        if newIndex:
            self.signal_AttrEdited.emit(node.name, index, newIndex)
        else:
            self.signal_AttrEdited.emit(node.name, index, index)

    def createConnection(self, sourceNode, sourceAttr, targetNode, targetAttr):
        """
        Create a manual connection.

        :type  sourceNode: str.
        :param sourceNode: Node that emits the connection.

        :type  sourceAttr: str.
        :param sourceAttr: Attribute that emits the connection.

        :type  targetNode: str.
        :param targetNode: Node that receives the connection.

        :type  targetAttr: str.
        :param targetAttr: Attribute that receives the connection.

        """
        plug = self.scene().nodes[sourceNode].plugs[sourceAttr]
        socket = self.scene().nodes[targetNode].sockets[targetAttr]

        connection = ConnectionItem(plug.center(), socket.center(), plug, socket)

        connection.plugNode = plug.parentItem().name
        connection.plugAttr = plug.attribute
        connection.socketNode = socket.parentItem().name
        connection.socketAttr = socket.attribute

        plug.connect(socket, connection)
        socket.connect(plug, connection)

        connection.updatePath()

        self.scene().addItem(connection)

        return connection

    def evaluateGraph(self):
        """
        Create a list of connection tuples.
        [("sourceNode.attribute", "TargetNode.attribute"), ...]

        """
        scene = self.scene()

        data = list()

        for item in scene.items():
            if isinstance(item, ConnectionItem):
                connection = item

                data.append(connection._outputConnectionData())

        # Emit Signal
        self.signal_GraphEvaluated.emit()

        return data

    def clearGraph(self):
        """
        Clear the graph.

        """
        self.scene().clear()
        self.scene().nodes = dict()

        # Emit signal.
        self.signal_GraphCleared.emit()

    ##################################################################
    # END API
    ##################################################################


class NodeScene(QtWidgets.QGraphicsScene):

    """
    The scene displaying all the nodes.

    """
    signal_NodeMoved = QtCore.pyqtSignal(str, object)

    def __init__(self, parent):
        """
        Initialize the class.

        """
        super(NodeScene, self).__init__(parent)

        # General.
        self.gridSize = parent.config['grid_size']

        # Nodes storage.
        self.nodes = dict()

    def dragEnterEvent(self, event):
        """
        Make the dragging of nodes into the scene possible.

        """
        event.setDropAction(QtCore.Qt.MoveAction)
        event.accept()

    def dragMoveEvent(self, event):
        """
        Make the dragging of nodes into the scene possible.

        """
        event.setDropAction(QtCore.Qt.MoveAction)
        event.accept()

    def dropEvent(self, event):
        """
        Create a node from the dropped item.

        """
        # Emit signal.
        self.signal_Dropped.emit(event.scenePos())

        event.accept()

    def drawBackground(self, painter, rect):
        """
        Draw a grid in the background.

        """
        if self.views()[0].gridVisToggle:
            leftLine = rect.left() - rect.left() % self.gridSize
            topLine = rect.top() - rect.top() % self.gridSize
            lines = list()

            i = int(leftLine)
            while i < int(rect.right()):
                lines.append(QtCore.QLineF(i, rect.top(), i, rect.bottom()))
                i += self.gridSize

            u = int(topLine)
            while u < int(rect.bottom()):
                lines.append(QtCore.QLineF(rect.left(), u, rect.right(), u))
                u += self.gridSize

            self.pen = QtGui.QPen()
            config = self.parent().config
            self.pen.setColor(convertDataToColor(config['grid_color']))
            self.pen.setWidth(0)
            painter.setPen(self.pen)
            painter.drawLines(lines)

    def updateScene(self):
        """
        Update the connections position.

        """
        for connection in [i for i in self.items() if isinstance(i, ConnectionItem)]:
            connection.target_point = connection.target.center()
            connection.source_point = connection.source.center()
            connection.updatePath()


class NodeItem(QtWidgets.QGraphicsItem):

    """
    A graphic representation of a node containing attributes.

    """

    def __init__(self, name, alternate, preset, config):
        """
        Initialize the class.

        :type  name: str.
        :param name: The name of the node. The name has to be unique
                     as it is used as a key to store the node object.

        :type  alternate: bool.
        :param alternate: The attribute color alternate state, if True,
                          every 2 attribute the color will be slightly
                          darker.

        :type  preset: str.
        :param preset: The name of graphical preset in the config file.

        """
        super(NodeItem, self).__init__()

        self.setZValue(1)

        # Storage
        self.name = name
        self.alternate = alternate
        self.nodePreset = preset
        self.attrPreset = None

        # Attributes storage.
        self.attrs = list()
        self.attrsData = dict()
        self.attrCount = 0
        self.currentDataType = None

        self.plugs = dict()
        self.sockets = dict()

        # Methods.
        self._createStyle(config)

    @property
    def height(self):
        """
        Increment the final height of the node every time an attribute
        is created.

        """
        if self.attrCount > 0:
            return (self.baseHeight +
                    self.attrHeight * self.attrCount +
                    self.border +
                    0.5 * self.radius)
        else:
            return self.baseHeight

    @property
    def width(self):
        """
        Calculate width dynamically.
        :return:
        """

        w = self.baseWidth #min width
        for attr in self.attrs:
            w = max(w,self._attrFontMetrics.width(attr)+self.radius*2+self.border)
        return w

    @property
    def pen(self):
        """
        Return the pen based on the selection state of the node.

        """
        if self.isSelected():
            return self._penSel
        else:
            return self._pen

    def _createStyle(self, config):
        """
        Read the node style from the configuration file.

        """
        self.setAcceptHoverEvents(True)
        self.setFlag(QtWidgets.QGraphicsItem.ItemIsMovable)
        self.setFlag(QtWidgets.QGraphicsItem.ItemIsSelectable)

        # Dimensions.
        self.baseWidth  = config['node_width']
        self.baseHeight = config['node_height']
        self.attrHeight = config['node_attr_height']
        self.border = config['node_border']
        self.radius = config['node_radius']

        self.nodeCenter = QtCore.QPointF()
        self.nodeCenter.setX(self.width / 2.0)
        self.nodeCenter.setY(self.height / 2.0)

        self._brush = QtGui.QBrush()
        self._brush.setStyle(QtCore.Qt.SolidPattern)
        self._brush.setColor(convertDataToColor(config[self.nodePreset]['bg']))

        self._pen = QtGui.QPen()
        self._pen.setStyle(QtCore.Qt.SolidLine)
        self._pen.setWidth(self.border)
        self._pen.setColor(convertDataToColor(config[self.nodePreset]['border']))

        self._penSel = QtGui.QPen()
        self._penSel.setStyle(QtCore.Qt.SolidLine)
        self._penSel.setWidth(self.border)
        self._penSel.setColor(convertDataToColor(config[self.nodePreset]['border_sel']))

        self._textPen = QtGui.QPen()
        self._textPen.setStyle(QtCore.Qt.SolidLine)
        self._textPen.setColor(convertDataToColor(config[self.nodePreset]['text']))

        self._nodeTextFont = QtGui.QFont(config['node_font'], config['node_font_size'], QtGui.QFont.Bold)
        self._attrTextFont = QtGui.QFont(config['attr_font'], config['attr_font_size'], QtGui.QFont.Normal)
        self._attrFontMetrics = QtGui.QFontMetrics(self._attrTextFont)

        self._attrBrush = QtGui.QBrush()
        self._attrBrush.setStyle(QtCore.Qt.SolidPattern)

        self._attrBrushAlt = QtGui.QBrush()
        self._attrBrushAlt.setStyle(QtCore.Qt.SolidPattern)

        self._attrPen = QtGui.QPen()
        self._attrPen.setStyle(QtCore.Qt.SolidLine)

    def _createAttribute(self, name, index, preset, plug, socket, dataType):
        """
        Create an attribute by expanding the node, adding a label and
        connection items.

        :type  name: str.
        :param name: The name of the attribute. The name has to be
                     unique as it is used as a key to store the node
                     object.

        :type  index: int.
        :param index: The index of the attribute in the node.

        :type  preset: str.
        :param preset: The name of graphical preset in the config file.

        :type  plug: bool.
        :param plug: Whether or not this attribute can emit connections.

        :type  socket: bool.
        :param socket: Whether or not this attribute can receive
                       connections.

        :type  dataType: type.
        :param dataType: Type of the data represented by this attribute
                         in order to highlight attributes of the same
                         type while performing a connection.

        """
        if name in self.attrs:
            print 'An attribute with the same name already exists on this node : {0}'.format(name)
            print 'Attribute creation aborted !'
            return

        self.attrPreset = preset

        # Create a plug connection item.
        if plug:
            plugInst = PlugItem(parent=self,
                                attribute=name,
                                index=self.attrCount,
                                preset=preset,
                                dataType=dataType)
            self.plugs[name] = plugInst

        # Create a socket connection item.
        if socket:
            socketInst = SocketItem(parent=self,
                                    attribute=name,
                                    index=self.attrCount,
                                    preset=preset,
                                    dataType=dataType)

            self.sockets[name] = socketInst

        self.attrCount += 1

        # Add the attribute based on its index.
        if index == -1 or index > self.attrCount:
            self.attrs.append(name)
        else:
            self.attrs.insert(index, name)

        # Store attr data.
        self.attrsData[name] = {'name': name,
                                'socket': socket,
                                'plug': plug,
                                'preset': preset,
                                'dataType': dataType}

        # Update node height.
        self.update()

    def _deleteAttribute(self, index):
        """
        Remove an attribute by reducing the node, removing the label
        and the connection items.

        :type  index: int.
        :param index: The index of the attribute in the node.

        """
        name = self.attrs[index]

        # Remove socket and its connections.
        if name in self.sockets.keys():
            for connection in self.sockets[name].connections:
                connection._remove()

            self.scene().removeItem(self.sockets[name])
            self.sockets.pop(name)

        # Remove plug and its connections.
        if name in self.plugs.keys():
            for connection in self.plugs[name].connections:
                connection._remove()

            self.scene().removeItem(self.plugs[name])
            self.plugs.pop(name)

        # Reduce node height.
        if self.attrCount > 0:
            self.attrCount -= 1

        # Remove attribute from node.
        if name in self.attrs:
            self.attrs.remove(name)

        self.update()

    def _remove(self):
        """
        Remove this node instance from the scene.

        Make sure that all the connections to this node are also removed
        in the process

        """
        self.scene().nodes.pop(self.name)

        # Remove all sockets connections.
        for socket in self.sockets.values():
            while len(socket.connections)>0:
                socket.connections[0]._remove()

        # Remove all plugs connections.
        for plug in self.plugs.values():
            while len(plug.connections)>0:
                plug.connections[0]._remove()

        # Remove node.
        scene = self.scene()
        scene.removeItem(self)
        scene.update()

    def boundingRect(self):
        """
        The bounding rect based on the width and height variables.

        """
        rect = QtCore.QRect(0, 0, self.width, self.height)
        rect = QtCore.QRectF(rect)
        return rect

    def shape(self):
        """
        The shape of the item.

        """
        path = QtGui.QPainterPath()
        path.addRect(self.boundingRect())
        return path

    def paint(self, painter, option, widget):
        """
        Paint the node and attributes.

        """
        # Node base.
        painter.setBrush(self._brush)
        painter.setPen(self.pen)

        painter.drawRoundedRect(0, 0,
                                self.width,
                                self.height,
                                self.radius,
                                self.radius)

        # Node label.
        painter.setPen(self._textPen)
        painter.setFont(self._nodeTextFont)

        metrics = QtGui.QFontMetrics(painter.font())
        text_width = metrics.boundingRect(self.name).width() + 14
        text_height = metrics.boundingRect(self.name).height() + 14
        margin = (text_width - self.baseWidth) * 0.5
        textRect = QtCore.QRect(-margin,
                                -text_height,
                                text_width,
                                text_height)

        painter.drawText(textRect,
                         QtCore.Qt.AlignCenter,
                         self.name)


        # Attributes.
        offset = 0
        for attr in self.attrs:
            nodzInst = self.scene().views()[0]
            config = nodzInst.config

            # Attribute rect.
            rect = QtCore.QRect(self.border / 2,
                                self.baseHeight - self.radius + offset,
                                self.width - self.border,
                                self.attrHeight)



            attrData = self.attrsData[attr]
            name = attr

            preset = attrData['preset']


            # Attribute base.
            self._attrBrush.setColor(convertDataToColor(config[preset]['bg']))
            if self.alternate:
                self._attrBrushAlt.setColor(convertDataToColor(config[preset]['bg'], True, config['alternate_value']))

            self._attrPen.setColor(convertDataToColor([0, 0, 0, 0]))
            painter.setPen(self._attrPen)
            painter.setBrush(self._attrBrush)
            if (offset / self.attrHeight) % 2:
                painter.setBrush(self._attrBrushAlt)

            painter.drawRect(rect)

            # Attribute label.
            painter.setPen(convertDataToColor(config[preset]['text']))
            painter.setFont(self._attrTextFont)

            # Search non-connectable attributes.
            if nodzInst.drawingConnection:
                if self == nodzInst.currentHoveredNode:
                    if (attrData['dataType'] != nodzInst.sourceSlot.dataType or
                        (nodzInst.sourceSlot.slotType == 'plug' and attrData['socket'] == False or
                         nodzInst.sourceSlot.slotType == 'socket' and attrData['plug'] == False)):
                        # Set non-connectable attributes color.
                        painter.setPen(convertDataToColor(config['non_connectable_color']))

            textRect = QtCore.QRect(rect.left() + self.radius,
                                     rect.top(),
                                     rect.width() - 2*self.radius,
                                     rect.height())
            painter.drawText(textRect, QtCore.Qt.AlignVCenter, name)

            offset += self.attrHeight

    def mousePressEvent(self, event):
        """
        Keep the selected node on top of the others.

        """
        nodes = self.scene().nodes
        for node in nodes.values():
            node.setZValue(1)

        for item in self.scene().items():
            if isinstance(item, ConnectionItem):
                item.setZValue(1)

        self.setZValue(2)

        super(NodeItem, self).mousePressEvent(event)

    def mouseMoveEvent(self, event):
        """
        .

        """
        if self.scene().views()[0].gridVisToggle:
            if self.scene().views()[0].gridSnapToggle or self.scene().views()[0]._nodeSnap:
                gridSize = self.scene().gridSize

                currentPos = self.mapToScene(event.pos().x() - self.width / 2,
                                             event.pos().y() - self.height / 2)

                snap_x = (round(currentPos.x() / gridSize) * gridSize) - gridSize/4
                snap_y = (round(currentPos.y() / gridSize) * gridSize) - gridSize/4
                snap_pos = QtCore.QPointF(snap_x, snap_y)
                self.setPos(snap_pos)

                self.scene().updateScene()
            else:
                self.scene().updateScene()
                super(NodeItem, self).mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        """
        .

        """
        # Emit node moved signal.
        self.scene().signal_NodeMoved.emit(self.name, self.pos())
        super(NodeItem, self).mouseReleaseEvent(event)

    def hoverLeaveEvent(self, event):
        """
        .

        """
        nodzInst = self.scene().views()[0]

        for item in nodzInst.scene().items():
            if isinstance(item, ConnectionItem):
                item.setZValue(0)

        super(NodeItem, self).hoverLeaveEvent(event)


class SlotItem(QtWidgets.QGraphicsItem):

    """
    The base class for graphics item representing attributes hook.

    """

    def __init__(self, parent, attribute, preset, index, dataType):
        """
        Initialize the class.

        :param parent: The parent item of the slot.
        :type  parent: QtWidgets.QGraphicsItem instance.

        :param attribute: The attribute associated to the slot.
        :type  attribute: String.

        :param index: int.
        :type  index: The index of the attribute in the node.

        :type  preset: str.
        :param preset: The name of graphical preset in the config file.

        :param dataType: The data type associated to the attribute.
        :type  dataType: Type.

        """
        super(SlotItem, self).__init__(parent)

        # Status.
        self.setAcceptHoverEvents(True)

        # Storage.
        self.slotType = None
        self.attribute = attribute
        self.preset = preset
        self.index = index
        self.dataType = dataType

        # Style.
        self.brush = QtGui.QBrush()
        self.brush.setStyle(QtCore.Qt.SolidPattern)

        self.pen = QtGui.QPen()
        self.pen.setStyle(QtCore.Qt.SolidLine)

        # Connections storage.
        self.connected_slots = list()
        self.newConnection = None
        self.connections = list()

    def mousePressEvent(self, event):
        """
        Start the connection process.

        """
        if event.button() == QtCore.Qt.LeftButton:
            self.newConnection = ConnectionItem(self.center(),
                                                self.mapToScene(event.pos()),
                                                self,
                                                None)

            self.connections.append(self.newConnection)
            self.scene().addItem(self.newConnection)

            nodzInst = self.scene().views()[0]
            nodzInst.drawingConnection = True
            nodzInst.sourceSlot = self
            nodzInst.currentDataType = self.dataType
        else:
            super(SlotItem, self).mousePressEvent(event)

    def mouseMoveEvent(self, event):
        """
        Update the new connection's end point position.

        """
        nodzInst = self.scene().views()[0]
        config = nodzInst.config
        if nodzInst.drawingConnection:
            mbb = createPointerBoundingBox(pointerPos=event.scenePos().toPoint(),
                                                  bbSize=config['mouse_bounding_box'])

            # Get nodes in pointer's bounding box.
            targets = self.scene().items(mbb)

            if any(isinstance(target, NodeItem) for target in targets):
                if self.parentItem() not in targets:
                    for target in targets:
                        if isinstance(target, NodeItem):
                            nodzInst.currentHoveredNode = target
            else:
                nodzInst.currentHoveredNode = None

            # Set connection's end point.
            self.newConnection.target_point = self.mapToScene(event.pos())
            self.newConnection.updatePath()
        else:
            super(SlotItem, self).mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        """
        Apply the connection if target_slot is valid.

        """
        nodzInst = self.scene().views()[0]
        if event.button() == QtCore.Qt.LeftButton:
            nodzInst.drawingConnection = False
            nodzInst.currentDataType = None

            target = self.scene().itemAt(event.scenePos().toPoint(), QtGui.QTransform())

            if not isinstance(target, SlotItem):
                self.newConnection._remove()
                super(SlotItem, self).mouseReleaseEvent(event)
                return

            if target.accepts(self):
                self.newConnection.target = target
                self.newConnection.source = self
                self.newConnection.target_point = target.center()
                self.newConnection.source_point = self.center()

                # Perform the ConnectionItem.
                self.connect(target, self.newConnection)
                target.connect(self, self.newConnection)

                self.newConnection.updatePath()
            else:
                self.newConnection._remove()
        else:
            super(SlotItem, self).mouseReleaseEvent(event)

        nodzInst.currentHoveredNode = None

    def shape(self):
        """
        The shape of the Slot is a circle.

        """
        path = QtGui.QPainterPath()
        path.addRect(self.boundingRect())
        return path

    def paint(self, painter, option, widget):
        """
        Paint the Slot.

        """
        painter.setBrush(self.brush)
        painter.setPen(self.pen)

        nodzInst = self.scene().views()[0]
        config = nodzInst.config
        if nodzInst.drawingConnection:
            if self.parentItem() == nodzInst.currentHoveredNode:
                painter.setBrush(convertDataToColor(config['non_connectable_color']))
                if (self.slotType == nodzInst.sourceSlot.slotType or (self.slotType != nodzInst.sourceSlot.slotType and self.dataType != nodzInst.sourceSlot.dataType)):
                    painter.setBrush(convertDataToColor(config['non_connectable_color']))
                else:
                    _penValid = QtGui.QPen()
                    _penValid.setStyle(QtCore.Qt.SolidLine)
                    _penValid.setWidth(2)
                    _penValid.setColor(QtGui.QColor(255, 255, 255, 255))
                    painter.setPen(_penValid)
                    painter.setBrush(self.brush)

        painter.drawEllipse(self.boundingRect())

    def center(self):
        """
        Return The center of the Slot.

        """
        rect = self.boundingRect()
        center = QtCore.QPointF(rect.x() + rect.width() * 0.5,
                                rect.y() + rect.height() * 0.5)

        return self.mapToScene(center)


class PlugItem(SlotItem):

    """
    A graphics item representing an attribute out hook.

    """

    def __init__(self, parent, attribute, index, preset, dataType):
        """
        Initialize the class.

        :param parent: The parent item of the slot.
        :type  parent: QtWidgets.QGraphicsItem instance.

        :param attribute: The attribute associated to the slot.
        :type  attribute: String.

        :param index: int.
        :type  index: The index of the attribute in the node.

        :type  preset: str.
        :param preset: The name of graphical preset in the config file.

        :param dataType: The data type associated to the attribute.
        :type  dataType: Type.

        """
        super(PlugItem, self).__init__(parent, attribute, preset, index, dataType)
        # Storage.
        self.attributte = attribute
        self.preset = preset
        self.slotType = 'plug'

        # Methods.
        self._createStyle(parent)

    def _createStyle(self, parent):
        """
        Read the attribute style from the configuration file.

        """
        config = parent.scene().views()[0].config
        self.brush = QtGui.QBrush()
        self.brush.setStyle(QtCore.Qt.SolidPattern)
        self.brush.setColor(convertDataToColor(config[self.preset]['plug']))

    def boundingRect(self):
        """
        The bounding rect based on the width and height variables.

        """
        width = height = self.parentItem().attrHeight / 4.0

        nodzInst = self.scene().views()[0]
        config = nodzInst.config

        x = self.parentItem().width - (width / 2.0)
        y = (self.parentItem().baseHeight - config['node_radius'] +
             self.parentItem().attrHeight * (3.0/8.0) +
             self.parentItem().attrs.index(self.attribute) * self.parentItem().attrHeight)

        rect = QtCore.QRectF(QtCore.QRect(x, y, width, height))
        return rect

    def accepts(self, socket_item):
        """
        Only accepts socket items that belong to other nodes.

        """
        if isinstance(socket_item, SocketItem):
            if self.parentItem() != socket_item.parentItem():
                if socket_item.dataType == self.dataType:
                    if socket_item in self.connected_slots:
                        return False
                    else:
                        return True
            else:
                return False
        else:
            return False

    def connect(self, socket_item, connection):
        """
        Connect to the given socket_item.

        """
        # Populate connection.
        connection.socketItem = socket_item
        connection.plugNode = self.parentItem().name
        connection.plugAttr = self.attribute

        # Add socket to connected slots.
        if socket_item in self.connected_slots:
            self.connected_slots.remove(socket_item)
        self.connected_slots.append(socket_item)

        # Add connection.
        if connection not in self.connections:
            self.connections.append(connection)

        # Emit signal.
        nodzInst = self.scene().views()[0]
        nodzInst.signal_PlugConnected.emit(connection.plugNode, connection.plugAttr, connection.socketNode, connection.socketAttr)

    def disconnect(self, connection):
        """
        Disconnect the given connection from this plug item.

        """
        # Emit signal.
        nodzInst = self.scene().views()[0]
        nodzInst.signal_PlugDisconnected.emit(connection.plugNode, connection.plugAttr, connection.socketNode, connection.socketAttr)

        # Remove connected socket from plug
        if connection.socketItem in self.connected_slots:
            self.connected_slots.remove(connection.socketItem)
        # Remove connection
        self.connections.remove(connection)


class SocketItem(SlotItem):

    """
    A graphics item representing an attribute in hook.

    """

    def __init__(self, parent, attribute, index, preset, dataType):
        """
        Initialize the socket.

        :param parent: The parent item of the slot.
        :type  parent: QtWidgets.QGraphicsItem instance.

        :param attribute: The attribute associated to the slot.
        :type  attribute: String.

        :param index: int.
        :type  index: The index of the attribute in the node.

        :type  preset: str.
        :param preset: The name of graphical preset in the config file.

        :param dataType: The data type associated to the attribute.
        :type  dataType: Type.

        """
        super(SocketItem, self).__init__(parent, attribute, preset, index, dataType)
        # Storage.

        self.attributte = attribute
        self.preset = preset
        self.slotType = 'socket'

        # Methods.
        self._createStyle(parent)

    def _createStyle(self, parent):
        """
        Read the attribute style from the configuration file.

        """
        config = parent.scene().views()[0].config
        self.brush = QtGui.QBrush()
        self.brush.setStyle(QtCore.Qt.SolidPattern)
        self.brush.setColor(convertDataToColor(config[self.preset]['socket']))

    def boundingRect(self):
        """
        The bounding rect based on the width and height variables.

        """
        width = height = self.parentItem().attrHeight / 4.0

        nodzInst = self.scene().views()[0]
        config = nodzInst.config

        x = - width / 2.0
        y = (self.parentItem().baseHeight - config['node_radius'] +
             self.parentItem().attrHeight * (3.0 / 8.0) +
             self.parentItem().attrs.index(self.attribute) * self.parentItem().attrHeight )

        rect = QtCore.QRectF(QtCore.QRect(x, y, width, height))
        return rect

    def accepts(self, plug_item):
        """
        Only accepts plug items that belong to other nodes.

        """
        if isinstance(plug_item, PlugItem):
            if (self.parentItem() != plug_item.parentItem() and
                len(self.connected_slots) <= 1):
                if plug_item.dataType == self.dataType:
                    if plug_item in self.connected_slots:
                        return False
                    else:
                        return True
            else:
                return False
        else:
            return False

    def connect(self, plug_item, connection):
        """
        Connect to the given plug item.

        """
        if len(self.connected_slots) > 0:
            # Already connected.
            self.connections[0]._remove()
            self.connected_slots = list()

        # Populate connection.
        connection.plugItem = plug_item
        connection.socketNode = self.parentItem().name
        connection.socketAttr = self.attribute

        # Add plug to connected slots.
        self.connected_slots.append(plug_item)

        # Add connection.
        if connection not in self.connections:
            self.connections.append(connection)

        # Emit signal.
        nodzInst = self.scene().views()[0]
        nodzInst.signal_SocketConnected.emit(connection.plugNode, connection.plugAttr, connection.socketNode, connection.socketAttr)

    def disconnect(self, connection):
        """
        Disconnect the given connection from this socket item.

        """
        # Emit signal.
        nodzInst = self.scene().views()[0]
        nodzInst.signal_SocketDisconnected.emit(connection.plugNode, connection.plugAttr, connection.socketNode, connection.socketAttr)

        # Remove connected plugs
        if connection.plugItem in self.connected_slots:
            self.connected_slots.remove(connection.plugItem)
        # Remove connections
        self.connections.remove(connection)


class ConnectionItem(QtWidgets.QGraphicsPathItem):

    """
    A graphics path representing a connection between two attributes.

    """

    def __init__(self, source_point, target_point, source, target):
        """
        Initialize the class.

        :param sourcePoint: Source position of the connection.
        :type  sourcePoint: QPoint.

        :param targetPoint: Target position of the connection
        :type  targetPoint: QPoint.

        :param source: Source item (plug or socket).
        :type  source: class.

        :param target: Target item (plug or socket).
        :type  target: class.

        """
        super(ConnectionItem, self).__init__()

        self.setZValue(1)

        # Storage.
        self.socketNode = None
        self.socketAttr = None
        self.plugNode = None
        self.plugAttr = None

        self.source_point = source_point
        self.target_point = target_point
        self.source = source
        self.target = target

        self.plugItem = None
        self.socketItem = None

        self.movable_point = None

        self.data = tuple()

        # Methods.
        self._createStyle()

    def _createStyle(self):
        """
        Read the connection style from the configuration file.

        """
        config = self.source.scene().views()[0].config
        self.setAcceptHoverEvents(True)
        self.setZValue(-1)

        self._pen = QtGui.QPen(convertDataToColor(config['connection_color']))
        self._pen.setWidth(config['connection_width'])

    def _outputConnectionData(self):
        """
        .

        """
        return ("{0}.{1}".format(self.plugNode, self.plugAttr),
                "{0}.{1}".format(self.socketNode, self.socketAttr))

    def mousePressEvent(self, event):
        """
        Snap the Connection to the mouse.

        """
        nodzInst = self.scene().views()[0]

        for item in nodzInst.scene().items():
            if isinstance(item, ConnectionItem):
                item.setZValue(0)

        nodzInst.drawingConnection = True

        d_to_target = (event.pos() - self.target_point).manhattanLength()
        d_to_source = (event.pos() - self.source_point).manhattanLength()
        if d_to_target < d_to_source:
            self.target_point = event.pos()
            self.movable_point = 'target_point'
            self.target.disconnect(self)
            self.target = None
            nodzInst.sourceSlot = self.source
        else:
            self.source_point = event.pos()
            self.movable_point = 'source_point'
            self.source.disconnect(self)
            self.source = None
            nodzInst.sourceSlot = self.target

        self.updatePath()

    def mouseMoveEvent(self, event):
        """
        Move the Connection with the mouse.

        """
        nodzInst = self.scene().views()[0]
        config = nodzInst.config

        mbb = createPointerBoundingBox(pointerPos=event.scenePos().toPoint(),
                                              bbSize=config['mouse_bounding_box'])

        # Get nodes in pointer's bounding box.
        targets = self.scene().items(mbb)

        if any(isinstance(target, NodeItem) for target in targets):

            if nodzInst.sourceSlot.parentItem() not in targets:
                for target in targets:
                    if isinstance(target, NodeItem):
                        nodzInst.currentHoveredNode = target
        else:
            nodzInst.currentHoveredNode = None

        if self.movable_point == 'target_point':
            self.target_point = event.pos()
        else:
            self.source_point = event.pos()

        self.updatePath()

    def mouseReleaseEvent(self, event):
        """
        Create a Connection if possible, otherwise delete it.

        """
        nodzInst = self.scene().views()[0]
        nodzInst.drawingConnection = False

        slot = self.scene().itemAt(event.scenePos().toPoint(), QtGui.QTransform())

        if not isinstance(slot, SlotItem):
            self._remove()
            self.updatePath()
            super(ConnectionItem, self).mouseReleaseEvent(event)
            return

        if self.movable_point == 'target_point':
            if slot.accepts(self.source):
                # Plug reconnection.
                self.target = slot
                self.target_point = slot.center()
                plug = self.source
                socket = self.target

                # Reconnect.
                socket.connect(plug, self)

                self.updatePath()
            else:
                self._remove()

        else:
            if slot.accepts(self.target):
                # Socket Reconnection
                self.source = slot
                self.source_point = slot.center()
                socket = self.target
                plug = self.source

                # Reconnect.
                plug.connect(socket, self)

                self.updatePath()
            else:
                self._remove()

    def _remove(self):
        """
        Remove this Connection from the scene.

        """
        if self.source is not None:
            self.source.disconnect(self)
        if self.target is not None:
            self.target.disconnect(self)

        scene = self.scene()
        scene.removeItem(self)
        scene.update()

    def updatePath(self):
        """
        Update the path.

        """
        self.setPen(self._pen)

        path = QtGui.QPainterPath()
        path.moveTo(self.source_point)
        dx = (self.target_point.x() - self.source_point.x()) * 0.5
        dy = self.target_point.y() - self.source_point.y()
        ctrl1 = QtCore.QPointF(self.source_point.x() + dx, self.source_point.y() + dy * 0)
        ctrl2 = QtCore.QPointF(self.source_point.x() + dx, self.source_point.y() + dy * 1)
        path.cubicTo(ctrl1, ctrl2, self.target_point)

        self.setPath(path)

############################################################################################################################
############################################################################################################################
############################################################################################################################


######################################################################
# Test signals
######################################################################

# Nodes
@QtCore.pyqtSlot(str)
def on_nodeCreated(nodeName):
    print 'node created : ', nodeName

@QtCore.pyqtSlot(str)
def on_nodeDeleted(nodeName):
    print 'node deleted : ', nodeName

@QtCore.pyqtSlot(str, str)
def on_nodeEdited(nodeName, newName):
    print 'node edited : {0}, new name : {1}'.format(nodeName, newName)

@QtCore.pyqtSlot(str)
def on_nodeSelected(nodesName):
    print 'node selected : ', nodesName

@QtCore.pyqtSlot(str, object)
def on_nodeMoved(nodeName, nodePos):
    print 'node {0} moved to {1}'.format(nodeName, nodePos)

# Attrs
@QtCore.pyqtSlot(str, int)
def on_attrCreated(nodeName, attrId):
    print 'attr created : {0} at index : {1}'.format(nodeName, attrId)

@QtCore.pyqtSlot(str, int)
def on_attrDeleted(nodeName, attrId):
    print 'attr Deleted : {0} at old index : {1}'.format(nodeName, attrId)

@QtCore.pyqtSlot(str, int, int)
def on_attrEdited(nodeName, oldId, newId):
    print 'attr Edited : {0} at old index : {1}, new index : {2}'.format(nodeName, oldId, newId)

# Connections
@QtCore.pyqtSlot(str, str, str, str)
def on_connected(srcNodeName, srcPlugName, destNodeName, dstSocketName):
    print 'connected src: "{0}" at "{1}" to dst: "{2}" at "{3}"'.format(srcNodeName, srcPlugName, destNodeName, dstSocketName)

@QtCore.pyqtSlot(str, str, str, str)
def on_disconnected(srcNodeName, srcPlugName, destNodeName, dstSocketName):
    print 'disconnected src: "{0}" at "{1}" from dst: "{2}" at "{3}"'.format(srcNodeName, srcPlugName, destNodeName, dstSocketName)

# Graph
@QtCore.pyqtSlot()
def on_graphSaved():
    print 'graph saved !'

@QtCore.pyqtSlot()
def on_graphLoaded():
    print 'graph loaded !'

@QtCore.pyqtSlot()
def on_graphCleared():
    print 'graph cleared !'

@QtCore.pyqtSlot()
def on_graphEvaluated():
    print 'graph evaluated !'

# Other
@QtCore.pyqtSlot(object)
def on_keyPressed(key):
    print 'key pressed : ', key


def u(value, size):
    global endian
    if endian == 'big':
        e = ">"
    else:
        e = "<"
    if size == 2:
        return struct.unpack(e+"H", value)[0]
    elif size == 4:
        return struct.unpack(e+"I", value)[0]
    elif size == 8:
        return struct.unpack(e+"Q", value)[0]


class NotDefinedObjectException(Exception):
    def __init__(self, msg):
        super(NotDefinedObjectException, self).__init__(msg)

class NoMemberFoundException(Exception):
    def __init__(self, msg):
        super(NoMemberFoundException, self).__init__(msg)

class CMember(object):
    global bits, pat, nodz

    def __init__(self, address, offset, name, size, flag, member_id, idx, cobject=None):
        self.address = address
        self.offset = offset
        self.name = name
        self.size = size
        self.flag = flag
        self.member_id = member_id
        self.type = idc.get_type(self.member_id)
        self.is_array = False
        self.idx = idx
        self.cobject = cobject
        self.connected_cobject = None
        # even if use_dbg=False, get_bytes read memory from debugger. TODO check if this is true.
        if self.type is None:
            # Didn't defined type explicitly.

            # Default value is integer(size=1,2,4,8).
            if (idc.is_byte(self.flag) and self.size == 1) or (idc.is_word(self.flag) and self.size == 2) or (idc.is_dword(self.flag) and self.size == 4) or (idc.is_qword(self.flag) and self.size == 8):
                self.value = u(idc.get_bytes(address, size, False), size)
            else:
                # maybe list
                # TODO handle correctly if type isn't array (enum, bitfield, ...)
                if idc.is_enum0(self.flag) or idc.is_bf(self.flag):
                    raise Exception("Not implemented")
                self.is_array = True
                self.value = idc.get_bytes(address, size, False)
        else:
            self.type = re.sub(pat, '', self.type)
            # Struct or type defined explicitly.
            if self.is_ptr:
                self.value = u(idc.get_bytes(address, size, False), size)
            else:
                # handle type.
                self.value = idc.get_bytes(address, size, False)
        # TODO
        # if self.is_ptr() -> Find CMember to which it points.(via CObjectManager?)

    def connect(self, target_cmember):
        self.connected_cobject = target_cmember
        nodz.createConnection(str(self.cobject), str(self), str(target_cmember.cobject), str(target_cmember))

    @property
    def is_ptr(self):
        return idc.is_off0(self.flag)

    @property
    def is_valid_ptr(self):
        """

        :return:
        """
        global bits
        return self.is_ptr and idc.is_mapped(self.value)

    @property
    def ptr_struct_name(self):
        """
        Handy dereference structure name.
        struc_x *** -> struc_x **
        :return: structure name to which CMember point.
        """
        assert self.is_ptr
        if self.type is None:
            return None
        assert self.type[-1] == '*'
        return self.type[:-1].rstrip(' ')

    @property
    def is_struct(self):
        return idc.is_struct(self.flag)

    @property
    def bottom_y(self):
        """

        :return: max(self.bottom, max(self.childlen's bottom))
        """
        if self.idx == 0:
            return self.cobject.node.pos().y()

    def __repr__(self):
        # Do you want type name?
        # TODO if string, preview string?
        if isinstance(self.value, (int, long)):
            return self.name + '  ' + ("0x{0:0" + str(self.size * 2) + "x}").format(self.value)
        else:
            return self.name + '  ' + "PREVIEW"


class CObject(object):
    global nodz

    def __init__(self, address, struct_name, pos, cmanager=None, parent_cmember=None):
        self.address = address
        self.members = []
        self.struct_name = re.sub(pat, '', struct_name)
        self.struct_id = idaapi.get_struc_id(self.struct_name)
        self.size = idc.get_struc_size(self.struct_id)
        self.cmanager = cmanager
        # set parent object
        self.parent_cmember = parent_cmember
        self.node = nodz.createNode(name=str(self), preset='node_preset_1', position=pos)
        if self.struct_id == idc.BADADDR:
            raise NotDefinedObjectException(self.struct_name + ' isn\'t defined. Please insert into structure window.')
        if idc.is_union(self.struct_id):
            raise NotDefinedObjectException("Union Not supported now.")

        idx = 0 # TODO remove this and bottom_y of CMember
        for member in idautils.StructMembers(self.struct_id):
            offset, name, size = member
            # TODO if member is struct, expand struct members.
            # TODO if member is array, expand array members.( But db array is maybe string, and user don't want it to expand... ) Only expand dd|dw|dq array?
            cmember = CMember(address + offset, offset, name, size, idc.get_member_flag(self.struct_id, offset), idc.get_member_id(self.struct_id, offset), idx, cobject=self)
            self.members.append(cmember)
            idx += 1
        if self.members is []:
            raise NoMemberFoundException("No member found at " + struct_name)

        for member in self.members:
            nodz.createAttribute(node=self.node, name=str(member), index=-1, preset='attr_preset_1', plug=True, socket=True, dataType=str)

    def connect(self):
        pos = self.node.pos()
        pos.setX(self.right_end + 40) # shift right a little bit
        for member in self.members:
            if member.is_valid_ptr:
                if self.cmanager.is_contain(member.value):
                    # address already exists.
                    member.connect(self.cmanager.search_cmember(member.value))
                else:
                    # create cobject
                    cobj = self.cmanager.add_cobject(member.value, member.ptr_struct_name, pos, parent_cmember=member)
                    member.connect(cobj.members[0]) # connect to top member
                    pos.setY(cobj.bottom_y + 40) # shift down a little bit

    def is_contain(self, address):
        if self.address <= address < self.address + self.size:
            return True
        return False

    def search_cmember(self, address):
        for member in self.members:
            if member.address <= address < member.address + member.size:
                return member
        return None

    @property
    def right_end(self):
        """

        :return:
        """
        return self.node.pos().x() + self.node.width

    @property
    def bottom_y(self):
        """

        :return: max(self.bottom, max(self.childlen's bottom))
        """
        bottom = self.node.pos().y() + self.node.height
        for member in self.members:
            if member.is_valid_ptr:
                bottom = max(bottom, member.connected_cobject.bottom_y)
        return bottom

    def __repr__(self):
        return self.struct_name + '@' + hex(self.address)


class CObjectManager(object):
    global pat

    def __init__(self, nodz, max_depth, main_address, main_struct_name):
        self.nodz = nodz
        self.max_depth = max_depth
        self.main_address = main_address
        self.main_struct_name = main_struct_name
        self.cobjects = []
        self.add_cobject(main_address, main_struct_name, None)

    def debug_dump(self):
        for cobject in self.cobjects:
            print "CObject : " + hex(cobject.address)
            print "=============================="
            for member in cobject.members:
                print member
            print "=============================="

    def add_cobject(self, address, struct_name, pos, parent_cmember=None):
        if self.is_contain(address):
            return
        cobj = CObject(address, struct_name, pos, cmanager=self, parent_cmember=parent_cmember)
        self.cobjects.append(cobj)
        cobj.connect()
        return cobj

    def search_cmember(self, address):
        obj = self.search_cobject(address)
        return obj.search_cmember(address)

    def search_cobject(self, address):
        for obj in self.cobjects:
            if obj.is_contain(address):
                return obj
        return None

    def is_contain(self, address):
        """

        :param address: Memory address.
        :return: True when object with specified address already added.
        """
        for obj in self.cobjects:
            if obj.is_contain(address):
                return True
        return False

    def auto_layout(self):
        for cobject in self.cobjects:
            print cobject.bottom_y

def object_view_main():
    global nodz #VERY IMPORTANT!!!!
    global max_depth
    global dbg_active
    #check if debugger active

    highlighted = ida_kernwin.get_highlight(ida_kernwin.get_current_widget())
    if highlighted is None:
        print "No highlighted item."
        return
    name,flag = highlighted
    #check if debugger is active
    if flag == 3:
        #name is register name (e.g. RAX)
        try:
            address = idc.get_reg_value(name)
        except:
            address = idc.BADADDR
    else:
        #maybe flag == 1 (what's 2?)
        address = ida_kernwin.str2ea(name)
    s = idc.get_type(address)
    if not s:
        s = ''
    struct_name = ida_kernwin.ask_str(s, 0, "Type struct")
    if not struct_name:
        return


    config_s = {
        "scene_width": 2000,
        "scene_height": 2000,
        "grid_size": 36,
        "antialiasing": True,
        "antialiasing_boost": True,
        "smooth_pixmap": True,

        "node_font": "Arial",
        "node_font_size": 12,
        "attr_font": "Arial",
        "attr_font_size": 10,
        "mouse_bounding_box": 80,

        "node_width": 200,
        "node_height": 25,
        "node_radius": 10,
        "node_border": 2,
        "node_attr_height": 30,
        "connection_width": 2,

        "alternate_value": 20,
        "grid_color": [50, 50, 50, 255],
        "slot_border": [50, 50, 50, 255],
        "non_connectable_color": [100, 100, 100, 255],
        "connection_color": [255, 155, 0, 255],

        "node_default": {
            "bg": [130, 130, 130, 255],
            "border": [50, 50, 50, 255],
            "border_sel": [250, 250, 250, 255],
            "text": [255, 255, 255, 255]
        },

        "attr_default": {
            "bg": [160, 160, 160, 255],
            "text": [220, 220, 220, 255],
            "plug": [255, 155, 0, 255],
            "socket": [255, 155, 0, 255]
        },

        "node_preset_1": {
            "bg": [80, 80, 80, 255],
            "border": [50, 50, 50, 255],
            "border_sel": [170, 80, 80, 255],
            "text": [230, 230, 230, 255]
        },

        "attr_preset_1": {
            "bg": [60, 60, 60, 255],
            "text": [220, 220, 220, 255],
            "plug": [255, 155, 0, 255],
            "socket": [255, 155, 0, 255]
        },

        "attr_preset_2": {
            "bg": [250, 120, 120, 255],
            "text": [220, 220, 220, 255],
            "plug": [255, 155, 0, 255],
            "socket": [255, 155, 0, 255]
        },

        "attr_preset_3": {
            "bg": [160, 160, 160, 255],
            "text": [220, 220, 220, 255],
            "plug": [255, 155, 0, 255],
            "socket": [255, 155, 0, 255]
        }
    }
    #mainwindow = [x for x in QtWidgets.QApplication.topLevelWidgets() if type(x) == QtWidgets.QMainWindow][0]
    #nodz = Nodz(mainwindow, config_s) #will close immediately. fix it.(set parent?)
    nodz = Nodz(None, config_s) #will close immediately. fix it.(set parent?)

    nodz.initialize()
    nodz.show()
    nodz.signal_NodeCreated.connect(on_nodeCreated)
    nodz.signal_NodeDeleted.connect(on_nodeDeleted)
    nodz.signal_NodeEdited.connect(on_nodeEdited)
    nodz.signal_NodeSelected.connect(on_nodeSelected)
    nodz.signal_NodeMoved.connect(on_nodeMoved)

    nodz.signal_AttrCreated.connect(on_attrCreated)
    nodz.signal_AttrDeleted.connect(on_attrDeleted)
    nodz.signal_AttrEdited.connect(on_attrEdited)

    nodz.signal_PlugConnected.connect(on_connected)
    nodz.signal_SocketConnected.connect(on_connected)
    nodz.signal_PlugDisconnected.connect(on_disconnected)
    nodz.signal_SocketDisconnected.connect(on_disconnected)

    nodz.signal_GraphSaved.connect(on_graphSaved)
    nodz.signal_GraphLoaded.connect(on_graphLoaded)
    nodz.signal_GraphCleared.connect(on_graphCleared)
    nodz.signal_GraphEvaluated.connect(on_graphEvaluated)

    nodz.signal_KeyPressed.connect(on_keyPressed)

    try:
        com = CObjectManager(nodz, max_depth, address, struct_name)
        com.auto_layout()
    except NotDefinedObjectException as e:
        print e
        print "Please report to me. X("
        return
    com.debug_dump()

    """
    # Node A
    nodeA = nodz.createNode(name='nodeA', preset='node_preset_1', position=None)

    nodz.createAttribute(node=nodeA, name='Aattr1AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABCD', index=-1, preset='attr_preset_1',
                         plug=True, socket=False, dataType=str)

    nodz.createAttribute(node=nodeA, name='Aattr2', index=-1, preset='attr_preset_1',
                         plug=False, socket=False, dataType=int)

    nodz.createAttribute(node=nodeA, name='Aattr3', index=-1, preset='attr_preset_2',
                         plug=True, socket=True, dataType=int)

    nodz.createAttribute(node=nodeA, name='Aattr4', index=-1, preset='attr_preset_2',
                         plug=True, socket=True, dataType=str)



    # Node B
    nodeB = nodz.createNode(name='nodeB', preset='node_preset_1')

    nodz.createAttribute(node=nodeB, name='Battr1', index=-1, preset='attr_preset_1',
                         plug=True, socket=False, dataType=str)

    nodz.createAttribute(node=nodeB, name='Battr2', index=-1, preset='attr_preset_1',
                         plug=True, socket=False, dataType=int)



    # Node C
    nodeC = nodz.createNode(name='nodeC', preset='node_preset_1')

    nodz.createAttribute(node=nodeC, name='Cattr1', index=-1, preset='attr_preset_1',
                         plug=False, socket=True, dataType=str)

    nodz.createAttribute(node=nodeC, name='Cattr2', index=-1, preset='attr_preset_1',
                         plug=True, socket=False, dataType=int)

    nodz.createAttribute(node=nodeC, name='Cattr3', index=-1, preset='attr_preset_1',
                         plug=True, socket=False, dataType=str)

    nodz.createAttribute(node=nodeC, name='Cattr4', index=-1, preset='attr_preset_2',
                         plug=False, socket=True, dataType=str)

    nodz.createAttribute(node=nodeC, name='Cattr5', index=-1, preset='attr_preset_2',
                         plug=False, socket=True, dataType=int)

    nodz.createAttribute(node=nodeC, name='Cattr6', index=-1, preset='attr_preset_3',
                         plug=True, socket=False, dataType=str)

    nodz.createAttribute(node=nodeC, name='Cattr7', index=-1, preset='attr_preset_3',
                         plug=True, socket=False, dataType=str)

    nodz.createAttribute(node=nodeC, name='Cattr8', index=-1, preset='attr_preset_3',
                         plug=True, socket=False, dataType=int)


    # Please note that this is a local test so once the graph is cleared
    # and reloaded, all the local variables are not valid anymore, which
    # means the following code to alter nodes won't work but saving/loading/
    # clearing/evaluating will.

    # Connection creation
    nodz.createConnection('nodeB', 'Battr2', 'nodeA', 'Aattr3')
    nodz.createConnection('nodeB', 'Battr1', 'nodeA', 'Aattr4')

    # Attributes Edition
    nodz.editAttribute(node=nodeC, index=0, newName=None, newIndex=-1)
    nodz.editAttribute(node=nodeC, index=-1, newName='NewAttrName', newIndex=None)

    # Attributes Deletion
    nodz.deleteAttribute(node=nodeC, index=-1)
    # Nodes Edition
    nodz.editNode(node=nodeC, newName='newNodeName')
    # Nodes Deletion
    nodz.deleteNode(node=nodeC)
    # Graph
    print nodz.evaluateGraph()
    print
    """


    '''
    if app:
        # command line stand alone test... run our own event loop
        app.exec_()
    '''


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

    def activate(self, ctx):
        try:
            action = object_view_main()
        except Exception as e:
            print(e)
        return 1

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

class UIHook(idaapi.UI_Hooks):
    form_type_list = [idaapi.BWN_DISASM,idaapi.BWN_DUMP,idaapi.BWN_NAMES,idaapi.BWN_PSEUDOCODE,idaapi.BWN_STACK,idaapi.BWN_STKVIEW,idaapi.BWN_STRINGS]
    def __init__(self):
        idaapi.UI_Hooks.__init__(self)

    def finish_populating_tform_popup(self, form, popup):
        form_type = idaapi.get_tform_type(form)

        if form_type in self.form_type_list:
            idaapi.attach_action_to_popup(form, popup, "Object View", None)

class ObjectViewerPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX | idaapi.PLUGIN_HIDE
    comment = "Graphical object viewer for IDA"
    help = ""
    wanted_name = "IDAObjectViewer"
    wanted_hotkey = ""

    def init(self):

        global bits
        info = idaapi.get_inf_structure()
        if info.is_64bit():
            bits = 64
        elif info.is_32bit():
            bits = 32
        else:
            bits = 16 #not tested

        print "Object Viewer Plugin loaded."
        self.ui_hook = UIHook()
        self.ui_hook.hook()
        self.action = idaapi.action_desc_t("Object View", "Object View", object_viewer_handler(), "")
        idaapi.register_action(self.action)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        self.ui_hook.unhook()
        idaapi.unregister_action("Object View")
        pass

def PLUGIN_ENTRY():
    return ObjectViewerPlugin()


