# plugins

repository: awgil/ffxiv_reverse
file: idaplugins/ffnetwork.py
description: Various utilities for reversing FF14
wanted_name: ffnetwork
comment: Build opcode map
created_at: 2023-02-06 18:28:16+00:00
pushed_at: 2023-10-03 10:25:46+00:00
url:
https://github.com/awgil/ffxiv_reverse/blob/f35b6226c1478234ca2b7149f82d251cffca2f56/idaplugins/ffnetwork.py

content:
import idaapi
import ida_ida
import ida_bytes
import ida_search

packet_names = {
	2: 'Ping', # note: could be +1/-1
	3: 'Init', # note: could be +1/-1
	8: 'Logout', # note: could be +1/-1
	11: 'CFCancel', # note: could be +1/-1
	13: 'CFDutyInfo', # note: could be +1/-1
	14: 'CFNotify', # note: could be +1/-1
	18: 'CFPreferredRole', # note: could be +1/-1
	81: 'CrossWorldLinkshellList', # note: could be +1/-1
	89: 'FellowshipList', # note: could be +1/-1
	111: 'Playtime', # note: could be +1/-1
	112: 'CFRegistered', # note: could be +1/-1
	115: 'Chat', # note: could be +1/-1
	121: 'RSVData',
	122: 'RSFData',
	123: 'SocialMessage',
	124: 'SocialMessage2',
	126: 'SocialList',
	127: 'SocialRequestResponse',
	128: 'ExamineSearchInfo',
	129: 'UpdateSearchInfo',
	130: 'InitSearchInfo',
	131: 'ExamineSearchComment',
	134: 'ServerNoticeShort',
	135: 'ServerNotice',
	136: 'SetOnlineStatus',
	137: 'LogMessage',
	141: 'Countdown',
	142: 'CountdownCancel',
	147: 'PartyMessage',
	149: 'PlayerAddedToBlacklist',
	150: 'PlayerRemovedFromBlacklist',
	151: 'BlackList',
	157: 'LinkshellList',
	158: 'MailDeleteRequest',
	162: 'MarketBoardItemListingCount',
	164: 'MarketBoardItemListing',
	166: 'MarketBoardPurchase',
	168: 'MarketBoardItemListingHistory',
	169: 'RetainerSaleHistory',
	171: 'MarketBoardSearchResult',
	173: 'FreeCompanyInfo',
	175: 'ExamineFreeCompanyInfo',
	176: 'FreeCompanyDialog',
	201: 'StatusEffectList',
	202: 'StatusEffectListEureka',
	203: 'StatusEffectListBozja',
	204: 'StatusEffectListDouble',
	206: 'EffectResult1',
	207: 'EffectResult4',
	208: 'EffectResult8',
	209: 'EffectResult16',
	211: 'EffectResultBasic1',
	212: 'EffectResultBasic4',
	213: 'EffectResultBasic8',
	214: 'EffectResultBasic16',
	215: 'EffectResultBasic32',
	216: 'EffectResultBasic64',
	217: 'ActorControl',
	218: 'ActorControlSelf',
	219: 'ActorControlTarget',
	220: 'UpdateHpMpTp',
	221: 'ActionEffect1',
	224: 'ActionEffect8',
	225: 'ActionEffect16',
	226: 'ActionEffect24',
	227: 'ActionEffect32',
	230: 'StatusEffectListPlayer',
	231: 'StatusEffectListPlayerDouble',
	233: 'UpdateRecastTimes',
	235: 'UpdateAllianceNormal',
	236: 'UpdateAllianceSmall',
	237: 'UpdatePartyMemberPositions',
	238: 'UpdateAllianceNormalMemberPositions',
	239: 'UpdateAllianceSmallMemberPositions',
	241: 'GCAffiliation', # ?
	260: 'SpawnPlayer',
	261: 'SpawnNPC',
	262: 'SpawnBoss',
	263: 'DespawnCharacter',
	264: 'ActorMove',
	266: 'Transfer',
	267: 'ActorSetPos',
	269: 'ActorCast',
	270: 'PlayerUpdateLook',
	271: 'UpdateParty',
	272: 'InitZone',
	273: 'ApplyIDScramble',
	274: 'UpdateHate',
	275: 'UpdateHater',
	276: 'SpawnObject',
	277: 'DespawnObject',
	278: 'UpdateClassInfo',
	279: 'UpdateClassInfoEureka',
	280: 'UpdateClassInfoBozja',
	281: 'PlayerSetup',
	282: 'PlayerStats',
	283: 'FirstAttack',
	284: 'PlayerStateFlags',
	285: 'PlayerClassInfo',
	287: 'ModelEquip',
	288: 'Examine',
	291: 'CharaNameReq',
	295: 'RetainerInformation',
	297: 'ItemMarketBoardInfo',
	299: 'ItemInfo',
	300: 'ContainerInfo',
	301: 'InventoryTransactionFinish',
	302: 'InventoryTransaction',
	303: 'CurrencyCrystalInfo',
	305: 'InventoryActionAck',
	306: 'UpdateInventorySlot',
	308: 'OpenTreasure',
	311: 'LootMessage',
	315: 'CreateTreasure',
	316: 'TreasureFadeOut',
	317: 'HuntingLogEntry',
	319: 'EventPlay',
	320: 'EventPlay4',
	321: 'EventPlay8',
	322: 'EventPlay16',
	323: 'EventPlay32',
	324: 'EventPlay64',
	325: 'EventPlay128',
	326: 'EventPlay255',
	328: 'EventStart',
	329: 'EventFinish',
	340: 'EventContinue',
	342: 'ResultDialog',
	343: 'DesynthResult',
	348: 'QuestActiveList',
	349: 'QuestUpdate',
	350: 'QuestCompleteList',
	351: 'QuestFinish',
	354: 'MSQTrackerComplete',
	366: 'QuestTracker',
	387: 'Mount',
	389: 'DirectorVars',
	390: 'ContentDirectorSync',
	392: 'EnvControl',
	398: 'SystemLogMessage1',
	399: 'SystemLogMessage2',
	400: 'SystemLogMessage4',
	401: 'SystemLogMessage8',
	402: 'SystemLogMessage16',
	404: 'BattleTalk2',
	405: 'BattleTalk4',
	406: 'BattleTalk8',
	408: 'MapUpdate',
	409: 'MapUpdate4',
	410: 'MapUpdate8',
	411: 'MapUpdate16',
	412: 'MapUpdate32',
	413: 'MapUpdate64',
	414: 'MapUpdate128',
	416: 'BalloonTalk2',
	417: 'BalloonTalk4',
	418: 'BalloonTalk8',
	420: 'WeatherChange',
	421: 'PlayerTitleList',
	422: 'Discovery',
	424: 'EorzeaTimeOffset',
	437: 'EquipDisplayFlags',
	438: 'NpcYell',
	443: 'FateInfo',
	457: 'LandSetInitialize',
	458: 'LandUpdate',
	459: 'YardObjectSpawn',
	460: 'HousingIndoorInitialize',
	461: 'LandAvailability',
	463: 'LandPriceUpdate',
	464: 'LandInfoSign',
	465: 'LandRename',
	466: 'HousingEstateGreeting',
	467: 'HousingUpdateLandFlagsSlot',
	468: 'HousingLandFlags',
	469: 'HousingShowEstateGuestAccess',
	471: 'HousingObjectInitialize',
	472: 'HousingInternalObjectSpawn',
	474: 'HousingWardInfo',
	475: 'HousingObjectMove',
	476: 'HousingObjectDye',
	488: 'SharedEstateSettingsResponse',
	500: 'DailyQuests',
	502: 'DailyQuestRepeatFlags',
	504: 'LandUpdateHouseName',
	515: 'AirshipTimers',
	518: 'PlaceMarker',
	519: 'WaymarkPreset',
	520: 'Waymark',
	523: 'UnMount', # note: could be up to -2
	526: 'CeremonySetActorAppearance', # note: could be up to -2
	532: 'AirshipStatusList', # note: could be up to -2
	533: 'AirshipStatus', # note: could be up to -2
	534: 'AirshipExplorationResult', # note: could be up to -2
	535: 'SubmarineStatusList', # note: could be up to -2
	536: 'SubmarineProgressionStatus', # note: could be up to -2
	537: 'SubmarineExplorationResult', # note: could be up to -2
	539: 'SubmarineTimers', # note: could be up to -2
	569: 'PrepareZoning', # note: could be up to -1
	570: 'ActorGauge',
	571: 'CharaVisualEffect',
	572: 'LandSetMap',
	573: 'Fall',
	622: 'PlayMotionSync',
	631: 'CEDirector',
	653: 'IslandWorkshopSupplyDemand',
}

def find_next_func_by_sig(ea, pattern):
	return ida_search.find_binary(ea, ida_ida.inf_get_max_ea(), pattern, 16, ida_search.SEARCH_DOWN)

def find_single_func_by_sig(pattern):
	ea_first = find_next_func_by_sig(ida_ida.inf_get_min_ea(), pattern)
	if ea_first == idaapi.BADADDR:
		print(f'Could not find function by pattern {pattern}')
		return 0
	if find_next_func_by_sig(ea_first + 1, pattern) != idaapi.BADADDR:
		print(f'Multiple functions match pattern {pattern}')
		return 0
	return ea_first

def read_signed_byte(ea):
	v = ida_bytes.get_byte(ea)
	return v - 0x100 if v & 0x80 else v

def read_signed_dword(ea):
	v = ida_bytes.get_dword(ea)
	return v - 0x100000000 if v & 0x80000000 else v

def read_rva(ea):
	return ea + 4 + read_signed_dword(ea)

def get_vfoff_for_body(body):
	# assume each case has the following body:
	# mov rax, [rcx]
	# lea r9, [r10+10h]
	# jmp qword ptr [rax+<vfoff>]
	if ida_bytes.get_byte(body) != 0x48 or ida_bytes.get_byte(body + 1) != 0x8B or ida_bytes.get_byte(body + 2) != 0x01:
		return -1
	if ida_bytes.get_byte(body + 3) != 0x4D or ida_bytes.get_byte(body + 4) != 0x8D or ida_bytes.get_byte(body + 5) != 0x4A or ida_bytes.get_byte(body + 6) != 0x10:
		return -1
	if ida_bytes.get_byte(body + 7) != 0x48 or ida_bytes.get_byte(body + 8) != 0xFF:
		return -1
	sz = ida_bytes.get_byte(body + 9)
	if sz == 0x60:
		return read_signed_byte(body + 10)
	elif sz == 0xA0:
		return read_signed_dword(body + 10)
	else:
		return -1

def vfoff_to_index(vfoff):
	if vfoff < 0x10:
		return -1 # first two vfs are dtor and exec
	if (vfoff & 7) != 0:
		return -1 # vf contains qwords
	return (vfoff >> 3) - 2

class ffnetwork(idaapi.plugin_t):
	flags = idaapi.PLUGIN_UNL
	comment = 'Build opcode map'
	help = ''
	wanted_name = 'ffnetwork'
	wanted_hotkey = ''

	_unknown_in_output = False

	def init(self):
		return idaapi.PLUGIN_OK

	def run(self, arg=None):
		# assume func starts with:
		# mov rax, [r8+10h]
		# mov r10, [rax+38h]
		# movzx eax, word ptr [r10+2]
		# add eax, -<min_case>
		# cmp eax, <max_case-min_case>
		# ja <default_off>
		# lea r11, <__ImageBase_off>
		# cdqe
		# mov r9d, ds::<jumptable_rva>[r11+rax*4]
		func = find_single_func_by_sig('49 8B 40 10  4C 8B 50 38  41 0F B7 42 02  83 C0 ??  3D ?? ?? ?? ??  0F 87 ?? ?? ?? ??  4C 8D 1D ?? ?? ?? ??  48 98  45 8B 8C 83 ?? ?? ?? ??')
		if func == 0:
			return
		min_case = -read_signed_byte(func + 15) # this is a negative
		jumptable_size = read_signed_dword(func + 17) + 1
		def_addr = read_rva(func + 23)
		imagebase = read_rva(func + 30)
		jumptable = imagebase + read_signed_dword(func + 40)
		opcodemap = {}
		for i in range(jumptable_size):
			body = imagebase + read_signed_dword(jumptable + 4 * i)
			if body == def_addr:
				continue
			case = i + min_case
			voff = get_vfoff_for_body(body)
			index = vfoff_to_index(voff)
			if index < 0:
				print(f'Unexpected body for case {case}')
				continue
			if index in opcodemap:
				print(f'Multiple opcodes map to single operation {index}: {hex(opcodemap[index])} and {hex(case)}')
				continue
			opcodemap[index] = case
		for k, v in sorted(opcodemap.items()):
			if k in packet_names:
				print(f'{packet_names[k]} = 0x{v:0{4}X},')
			elif self._unknown_in_output:
				print(f'Packet{k} = 0x{v:0{4}X},')

	def term(self):
		pass

def PLUGIN_ENTRY():
	return ffnetwork()


repository: cellebrite-labs/LabSync
file: labsync.py
url: https://github.com/cellebrite-labs/LabSync/blob/07e66fb25e11c85738b32b64248b6620615036b1/labsync.py
wanted_name: LabSync
comment: helps to partially synchronize IDBs over git
created_at: 2024-12-31 09:54:10+00:00
pushed_at: 2025-01-08 07:38:31+00:00
description:
An IDA plugin that can be used to partially synchronize IDBs between different users reversing the same binaries

content:
from __future__ import annotations

import configparser
import contextlib
import dataclasses
import functools
import heapq
import io
import logging
import os
import pathlib
import re
import string
import tempfile
import time
import uuid
from collections.abc import Generator
from typing import IO, Any, Optional

import git
import parse
import yaml

try:
    import functioninliner
except ModuleNotFoundError:
    # define a mock that will allow us to work without functioninliner installed
    class functioninliner:  # noqa: N801
        class ClonesStorage(dict):
            def update_from_storage(self) -> None:
                pass

try:
    import ida_diskio
    import ida_idaapi
    import ida_idp
    import ida_kernwin
    import ida_loader
    import ida_nalt
    import ida_name
    import ida_segment
    import ida_typeinf
    import ida_xref
    import idautils
    import netnode
    import sark
except ModuleNotFoundError:
    # define mocks to support importing outside of IDA for testing
    class ida_kernwin:  # noqa: N801
        class action_handler_t:  # noqa: N801
            pass

        class UI_Hooks:  # noqa: N801
            pass

    class ida_idaapi:  # noqa: N801
        PLUGIN_MOD = 0
        PLUGIN_HIDE = 0
        BADADDR = 0

        class plugin_t:  # noqa: N801
            pass

    class ida_idp:  # noqa: N801
        IDP_INTERFACE_VERSION = 0

    class ida_typeinf:  # noqa: N801
        class text_sink_t:  # noqa: N801
            pass

        class tinfo_t:  # noqa: N801
            pass

    class netnode:  # noqa: N801
        class Netnode:
            pass


# CONFIGURATION


# we decided not to normalize prototypes because it makes it much harder to resolve conflicts since
# you don't know which function you're looking at
#
# the downside is that a conflict on a function name change will result in two conflicts (one on
# the name and one on the prototype)
NORMALIZE_PROTOTYPES = False

# we decided not to remove names and prototypes that are missing in the YAML since some times
# exporting/importing them remotely can be an issue and so they will be removed locally as well
REMOVE_MISSING_NAMES_AND_PROTOTYPES = False

LOCAL_TYPES_COMMENT_FMT = "/* >> LABSYNC DO NOT TOUCH: {} << */"

LOCKFILE = "labsync.lock"
DEFAULT_LOCK_TIMEOUT = 60  # sec


# LOGGING


class LoggerWithTrace(logging.getLoggerClass()):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        logging.TRACE = 5
        logging.addLevelName(logging.TRACE, "TRACE")

    def trace(self, msg: str, *args, **kwargs) -> None:
        self.log(logging.TRACE, msg, *args, **kwargs)


logger = LoggerWithTrace("LabSync")


# EXCEPTIONS


class LabSyncError(Exception):
    pass


class LabSyncLockError(LabSyncError):
    pass


class LabSyncBinaryMatchingError(LabSyncError):
    pass


# HELPERS


class LabSyncYAMLDumper(yaml.CDumper):
    @staticmethod
    def _hex_representer(dumper: yaml.Dumper, data: int) -> str:
        return dumper.represent_scalar("tag:yaml.org,2002:int", hex(data))

    @staticmethod
    def _str_representer(dumper: yaml.dumper, data: str) -> str:
        if "\n" in data:
            return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
        else:
            return dumper.represent_scalar("tag:yaml.org,2002:str", data)


LabSyncYAMLDumper.add_representer(int, LabSyncYAMLDumper._hex_representer)  # noqa: SLF001
LabSyncYAMLDumper.add_representer(str, LabSyncYAMLDumper._str_representer)  # noqa: SLF001


@contextlib.contextmanager
def wait_box(msg: str, *, hide_cancel: bool = False) -> None:
    prefix = "HIDECANCEL\n" if hide_cancel else ""
    ida_kernwin.show_wait_box(prefix + msg)
    try:
        yield None
    finally:
        ida_kernwin.hide_wait_box()


class StringIOTextSink(ida_typeinf.text_sink_t):
    def __init__(self):
        super().__init__()
        self.sio = io.StringIO()

    def _print(self, thing: str) -> int:
        self.sio.write(thing)
        return 0


def local_types() -> Generator[str]:
    name = ida_typeinf.first_named_type(None, ida_typeinf.NTF_TYPE)
    while name:
        yield name
        name = ida_typeinf.next_named_type(None, name, ida_typeinf.NTF_TYPE)


# EXPORT LOGIC


@dataclasses.dataclass(eq=True, order=True)
class SyncedBinary:
    idb_id: str = dataclasses.field(compare=False)
    start_ea: int = 0  # must be the first field, because we rely on it when sorting
    end_ea: int = ida_idaapi.BADADDR  # exclusive
    base_ea: int = 0
    seg_prefix: Optional[str] = None

    def contains(self, ea: int) -> bool:
        return self.start_ea <= ea < self.end_ea

    def ea2dump(self, ea: int) -> int:
        return ea - self.base_ea

    def dump2ea(self, dea: int) -> int:
        return dea + self.base_ea


def dump_names(binary: SyncedBinary, storage: functioninliner.ClonesStorage) -> dict[int, str]:
    d = {}
    for ea, name in idautils.Names():
        if not binary.contains(ea):
            continue

        seg = ida_segment.getseg(ea)
        seg_name = ida_segment.get_segm_name(seg)

        # skip names that are in inlined chunks
        if seg_name.startswith("inlined_"):
            continue

        # skip names for inlined functions
        if ea in storage:
            continue

        dea = binary.ea2dump(ea)
        d[dea] = name

    return d


def dump_inlined_funcs(binary: SyncedBinary, storage: functioninliner.ClonesStorage) -> list[int]:
    funcs = (binary.ea2dump(ea) for ea in storage if binary.contains(ea))
    return list(sorted(funcs))  # noqa: C413


def stable_topological_sort(graph: dict[Any, set[Any]]) -> Generator[Any]:
    heap = []
    next_v = None

    while heap or graph:
        for v, edges in list(graph.items()):
            if next_v is not None:
                edges.discard(next_v)
            if not edges:
                heapq.heappush(heap, v)
                del graph[v]

        if not heap:
            msg = "graph contains unsolvable dependencies"
            raise ValueError(msg)

        next_v = heapq.heappop(heap)
        yield next_v


def dump_local_types(types: netnode.Netnode) -> str:
    # we emulate print_decls() ourselves because it internally uses PRTYPE_NOREGEX and this
    # removes namespaces which starts with double underscore (e.g. std::__1::__libcpp_refstring)

    local_types_by_ordinal = {}  # ordinal: (name, decl, dependencies)
    local_types_by_name = {}  # name: (ordinal, decl, dependencies)

    tinfo = ida_typeinf.tinfo_t()
    for ordinal in range(1, ida_typeinf.get_ordinal_qty(None)):
        if not tinfo.get_numbered_type(None, ordinal):
            continue  # deleted ordinal

        name = tinfo.get_type_name()

        flags = (
            ida_typeinf.PRTYPE_MULTI |  # multiline
            ida_typeinf.PRTYPE_TYPE |  # required to have it named
            ida_typeinf.PRTYPE_PRAGMA |  # include alignment pragmas
            ida_typeinf.PRTYPE_SEMI |  # end with semicolon
            ida_typeinf.PRTYPE_CPP |  # unsure if this is needed, but to be on the safe side...
            ida_typeinf.PRTYPE_DEF |  # required to have a full definition
            ida_typeinf.PRTYPE_NOREGEX  # required to keep the name as-is
        )

        decl = ida_typeinf.print_tinfo(None, 2, 0, flags, tinfo, name, None)
        decl = decl.strip()
        # also strip trailing spaces since they arn't block-encodable in YAML
        decl = "\n".join(line.rstrip() for line in decl.splitlines())

        # TODO @TH: there is an IDA bug where comment-only changes to local types are not updated
        #           when using parse_decls() so for now we just strip all comments until it'll be
        #           fixed
        decl = strip_comments(decl)

        # IDA apparently can't handle templates in parse_decls(), so we we don't bother syncing
        # them at all. hopefully no sane reverser actually uses them and these are only imported
        # from debug symbols and never touched
        if "<" in decl:
            logger.warning(
                f"skipping syncing of local type {name!r} because templates are unsupported"  # noqa: COM812
            )
            continue

        # do a sanity for the extract of the name, since it'll be used when updating
        parsed_name, _ = decl_to_name_and_type(decl)
        assert name == parsed_name

        # see if this type is dependent on others
        dependencies = set()

        udt = ida_typeinf.udt_type_data_t()
        if tinfo.get_udt_details(udt):
            for i in range(udt.size()):
                udm = udt[i]
                udm_ordinal = udm.type.get_ordinal()
                if udm_ordinal:
                    dependencies.add(udm_ordinal)

        # keep it
        local_types_by_ordinal[ordinal] = (name, decl, dependencies)
        assert name not in local_types_by_name
        local_types_by_name[name] = (ordinal, decl, dependencies)

    dep_graph = {}
    for name, _, deps in local_types_by_ordinal.values():
        dep_names = {local_types_by_ordinal[d][0] for d in deps if d in local_types_by_ordinal}
        if len(dep_names) != len(deps):
            logger.warning(
                f"skipping syncing of local type {name!r} because it depends on other non-synced "
                "local types"  # noqa: COM812
            )
            continue

        dep_graph[name] = dep_names

    # generate a normalized header file that is sorted by dependecy order and lexigraphically,
    # to make YAML diffs sane
    nhdr = io.StringIO()

    for name in stable_topological_sort(dep_graph):
        _, decl, _ = local_types_by_name[name]

        # add uuid
        tid = ida_typeinf.get_named_type_tid(name)
        if tid == ida_idaapi.BADADDR:
            msg = f"failed to resolve tid of local type {name!r}:\n{decl}"
            raise LabSyncError(msg)
        decl_uuid = types.get(tid)
        if not decl_uuid:
            types[tid] = decl_uuid = str(uuid.uuid4())

        nhdr.write(LOCAL_TYPES_COMMENT_FMT.format(decl_uuid))
        nhdr.write("\n")
        nhdr.write(decl)
        nhdr.write("\n\n")

    return nhdr.getvalue().strip()


def strip_comments(decl: str) -> str:
    stripped_lines = []

    for line in decl.splitlines():
        try:
            line = line[:line.index("//")].rstrip()
            if not line:
                continue
        except ValueError:
            pass
        stripped_lines.append(line)

    return "\n".join(stripped_lines)


@functools.cache
def decl_to_type_name_pat() -> re.Pattern:
    """this generates a pattern that tries to match the first line of a (stripped) decl to the
    its type (e.g. struct/union/typedef) and name

    the regex matches 2N groups where group 2i is "type" and group 2i+1 is name, for different
    possible subregexes. Only 2 groups (for some i) should be matched

    this regex is a bit more "allowing" than how IDA formats decls, since we also use it to
    match decls from YAMLs in which the user might've changed some whitespacing while manually
    resolving a merge conflict
    """

    # from ida.cfg:TypeNameChars
    name_chars = r"_:$()`'{}" + string.digits + string.ascii_letters
    # from blackbox testing what's allowed from name_chars as the first character
    name_first_chars = r"_$`" + string.ascii_letters

    name_pat = r"([" + re.escape(name_first_chars) + "][" + re.escape(name_chars) + "]*?)"

    type_pat = (
        r"(?!typedef)(\S+)(?=\s)"  # only match non-typedefs
        r".*?\s" +  # everything up to the name (i.e. type + attributes)
        name_pat +
        r"(?:\s*(?<!:):(?!:).*)?"  # optional inheritance or IDA syntax for data types of enums
        r"\s*;?"  # optional semicolon in case of forward declarations
    )

    fptr_typedef_pat = (
        r"(typedef)(?=\s)"  # only match typdefs
        r".*?"  # everything up to the name
        r"\*\s*" +  # the star before the name
        name_pat +
        r"\s*\)\(.*"  # match the first function def in the line, in case there is also a fptr arg
        r".*;"  # everything else
    )

    norm_typedef_pat = (
        r"(typedef)(?=\s)"  # only match typdefs
        r".*?"  # everything up to the name
        r"\**" +  # optional stars before the name
        name_pat +
        r"(?:\s*\[\s*\d*\s*\])?"  # optional array part
        r"\s*;"  # end of the typedef
    )

    pat = (
        r"^(?:"  # start of line
            r"(?:" + type_pat + ")"
        r"|"
            r"(?:" + fptr_typedef_pat + ")"
        r"|"
            r"(?:" + norm_typedef_pat + ")"
        r")$"  # end of line
    )

    return re.compile(pat)


def decl_to_name_and_type(decl: str) -> tuple[str, str]:
    # skip pragma/comment lines
    for first_line in decl.splitlines():
        if not any(first_line.lstrip().startswith(x) for x in ("#", "//")):
            break
    else:
        msg = f"empty local type:\n{decl}"
        raise LabSyncError(msg)

    # strip the first line
    first_line = first_line.strip()

    # extract the name and decl type
    pat = decl_to_type_name_pat()
    m = pat.match(first_line)

    if not m:
        msg = f"failed to parse local type:\n{decl}"
        raise LabSyncError(msg)

    decl_type = m.group(m.lastindex - 1)
    name = m.group(m.lastindex)
        # should never happen according to our regex
    assert decl_type
    assert name

    return name, decl_type


def fix_non_present_arguments(name: str, tinfo: ida_typeinf.tinfo_t, *, add: bool = True) \
        -> tuple[ida_typeinf.tinfo_t, bool]:

    def type_exists(tinfo: ida_typeinf.tinfo_t) -> bool:
        if tinfo.present():
            return True

        # originally we used just tinfo.present(), but for some reason it keeps returning False
        # even after we saved the type (as a forward declaration)
        #
        # then we used tinfo.get_ordinal() > 0 as a test, but on huge IDBs with >10k types, for
        # some reason it kept returning 0 even after we saved the type
        #
        # therefore we moved to checking if we can get the tid for the type name. you have to watch
        # out, however, since for deleted types tinfo.get_type_name() raises UnicodeDecodeError
        try:
            tname = tinfo.get_type_name()
        except UnicodeDecodeError:
            return False

        tid = ida_typeinf.get_named_type_tid(tname)
        return tid != ida_idaapi.BADADDR

    def fix_non_present(tinfo: ida_typeinf.tinfo_t) -> tuple[ida_typeinf.tinfo_t, bool]:
        tinfo_orig = tinfo.copy()

        # deref pointer/array until we reach the actual type
        depth = 0
        while depth < 128:
            if not tinfo.remove_ptr_or_array():
                break
            depth += 1
        else:
            msg = "max pointer depth reached"
            raise LabSyncError(msg)

        # if we're allowed to and this type is missing, add its base to local types
        if not type_exists(tinfo) and add:
            # add the type to local types
            if tinfo.save_type() == 0:
                tname = tinfo.get_type_name()

                logger.warning(
                    f"the prototype for {name} used type {tname} that was not present in the TIL. "
                    "we silently added it to allow syncing"  # noqa: COM812
                )

                return tinfo_orig, False
            else:
                # we can't use tname for deleted types (it raises UnicodeDecodeError)
                tinfo_clean = tinfo.copy()
                tinfo_clean.set_modifiers(0)
                tname = tinfo_clean.dstr()

                logger.warning(
                    f"the prototype for {name} used type {tname} that was not present in the TIL. "
                    "we failed to silently add it to allow syncing (perhaps a deleted type?)"  # noqa: COM812
                )

        # TODO @TH: IDA has a bug where they can't parse _BOOL8 args, so we replace them.
        #           remove this flow after they fix it
        bool8_realtype = ida_typeinf.BT_BOOL | ida_typeinf.BTMT_BOOL8
        if tinfo.get_realtype() == bool8_realtype:
            pass
        elif type_exists(tinfo):
            # if it's real present type, we're good
            return tinfo_orig, False

        # replace the type with an unknown type
        tinfo_generic = ida_typeinf.tinfo_t()
        assert tinfo_generic.create_simple_type(ida_typeinf.BT_UNKNOWN)

            # set the original modifiers
        tinfo_generic.set_modifiers(tinfo.get_modifiers())

            # recrate the pointer depth on top of tinfo_generic
        for _ in range(depth):
            assert tinfo_generic.create_ptr(tinfo_generic)

        return tinfo_generic, True

    ftype = ida_typeinf.func_type_data_t()
    assert tinfo.get_func_details(ftype, ida_typeinf.GTD_NO_ARGLOCS)

    ftype.rettype, fixed = fix_non_present(ftype.rettype)
    for i, argtype in enumerate(ftype):
        ftype[i].type, arg_fixed = fix_non_present(argtype.type)
        fixed |= arg_fixed

    tinfo_new = ida_typeinf.tinfo_t()
    assert tinfo_new.create_func(ftype)

    return tinfo_new, fixed


def prototype(ea: int) -> str:
    tinfo = ida_typeinf.tinfo_t()
    if not ida_nalt.get_tinfo(tinfo, ea):
        return None

    name = ida_name.get_ea_name(ea, ida_name.GN_VISIBLE)

    # replace non-present arguments in the prototype if relevant
    tinfo_new, fixed = fix_non_present_arguments(name, tinfo)

    if fixed:
        ptype = ida_typeinf.print_tinfo(None, 0, 0, ida_typeinf.PRTYPE_1LINE, tinfo, None, None)
        new_ptype = ida_typeinf.print_tinfo(None, 0, 0, ida_typeinf.PRTYPE_1LINE, tinfo_new, None,
                                            None)
        logger.warning(
            f"replacing prototype for {name} because it uses types that are not present in the "
            f"TIL from:\n\t{ptype!r}\nto:\n\t{new_ptype!r}"  # noqa: COM812
        )

        if not ida_nalt.set_tinfo(ea, tinfo_new):
            logger.warning(f"failed setting new prototype for {name}! skipping it")
            return None

        tinfo = tinfo_new

    # generate the prototype to dump
    if NORMALIZE_PROTOTYPES:
        name = "FUNCTION"

    # we have to remove special characters from the name, otherwise we'll have an issue applying
    # the prototype afterwards (e.g. `__Foo.cxx_destruct_`)
    allowed = r"_$" + string.digits + string.ascii_letters
    pname = "".join(c if c in allowed else "_" for c in name)
    return ida_typeinf.print_tinfo(None, 0, 0, ida_typeinf.PRTYPE_1LINE, tinfo, pname, None)


def dump_prototypes(
    binary: SyncedBinary, storage: functioninliner.ClonesStorage) -> dict[int, str]:

    d = {}
    for ea in idautils.Functions():
        if not binary.contains(ea):
            continue

        seg = ida_segment.getseg(ea)
        seg_name = ida_segment.get_segm_name(seg)

        # skip funcs that are in inlined chunks somehow (shouldn't happen)
        if seg_name.startswith("inlined_"):
            continue

        # skip funcs that have been inlined
        if ea in storage:
            continue

        ptype = prototype(ea)
        if ptype:
            dea = binary.ea2dump(ea)
            d[dea] = ptype

    return d


def dump(binary: SyncedBinary, types: netnode.Netnode) -> str:
    storage = functioninliner.ClonesStorage()
    storage.update_from_storage()

    d = {
        "version": 4,
        "names": dump_names(binary, storage),
        "inlined_funcs": dump_inlined_funcs(binary, storage),
        # we have to dump prototypes before we dump local types because this may add new types to
        # the TIL
        "prototypes": dump_prototypes(binary, storage),
        "local_types": dump_local_types(types),
    }

    return yaml.dump(
        d, Dumper=LabSyncYAMLDumper, default_flow_style=False, sort_keys=True,
    )


# IMPORT LOGIC


def update_names(
    binary: SyncedBinary, storage: functioninliner.ClonesStorage, names: dict[int, str]) -> None:

    # delete names if required
    if REMOVE_MISSING_NAMES_AND_PROTOTYPES:
        for dea in dump_names(binary, storage):
            # delete name if unnamed in the new dict
            if dea not in names:
                ea = binary.dump2ea(dea)

                msg = f"removing name from {ea:#x}"
                logger.debug(msg)

                success = ida_name.set_name(ea, "", ida_name.SN_NOWARN)
                if not success:
                    if logger.getEffectiveLevel() > logging.DEBUG:
                        logger.warning("failed " + msg)
                    else:
                        logger.warning("removal failed!")

    # update names
    for dea, name in names.items():
        ea = binary.dump2ea(dea)
        cur_name = ida_name.get_name(ea)

        if cur_name != name:
            msg = f"renaming {ea:#x} from {cur_name!r} to {name!r}"
            logger.debug(msg)

            # check if the new name already exists in the database
            cur_ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name)
            name_changed = False
            try:
                # if the new name is already in use in the IDB --
                if cur_ea != ida_idaapi.BADADDR:
                    # verify that the repo also has a different name for the EA currently holding
                    # the new name
                    #
                    # perhaps we can even assert that this never happens
                    cur_dea = binary.ea2dump(cur_ea)
                    if cur_dea not in names:
                        logger.warning(
                            f"cannot rename {ea:#x} to {name!r} as this name already "
                            f"exists in the IDB for {cur_ea:#x}, and that EA doesn't "
                            "have a different name in the repo"  # noqa: COM812
                        )
                        continue

                    # temporarily rename it to something else
                    msg2 = f"temporarily renaming {cur_ea:#x} away from {name!r}"
                    logger.debug("\t" + msg2)
                    success = ida_name.set_name(
                        cur_ea,
                        name + "_labsync_temp",
                        ida_name.SN_NOWARN | ida_name.SN_FORCE,
                    )

                    # handle temporary rename failure
                    if not success:
                        if logger.getEffectiveLevel() > logging.DEBUG:
                            logger.warning("failed " + msg2)
                        else:
                            logger.warning("\ttemporary rename failed!")
                        continue

                # now do the actual rename
                success = ida_name.set_name(ea, name, ida_name.SN_NOWARN)

                # handle rename failure
                if not success:
                    if logger.getEffectiveLevel() > logging.DEBUG:
                        logger.warning("failed " + msg)
                    else:
                        logger.warning("rename failed!")
                    continue

                name_changed = True
            finally:
                # if we failed, undo the temporary rename if we did any
                if cur_ea != ida_idaapi.BADADDR and not name_changed:
                    msg2 = f"undoing the temporarily rename of {cur_ea:#x}"
                    logger.debug("\t" + msg2)
                    success = ida_name.set_name(cur_ea, name, ida_name.SN_NOWARN)

                    # handle temporary rename undoing failure
                    if not success:
                        if logger.getEffectiveLevel() > logging.DEBUG:
                            logger.warning("failed " + msg2)
                        else:
                            logger.warning("\ttemporary rename undoing failed!")


def update_inlined_funcs(
    binary: SyncedBinary, storage: functioninliner.ClonesStorage, funcs: list[int]) -> None:

    cur = {ea for ea in storage if binary.contains(ea)}
    new = set(map(binary.dump2ea, funcs))

    undo = cur - new
    do = new - cur

    for ea in undo:
        func = sark.Function(ea)
        if func.ea != ea:
            logger.warning(f"\tcannot undo inlining of {ea:#x} since it's not a function start!")
            continue
        msg = f"undoing inlining of {func.name}"
        logger.debug(msg)
        functioninliner.undo_inline_function(func)

    for ea in do:
        func = sark.Function(ea)
        if func.ea != ea:
            logger.warning(
                f"not inlining function @ {ea:#x} since it's not a function start"  # noqa: COM812
            )
        logger.debug(f"inlining {func.name}")
        functioninliner.inline_function(func)


def _rename_local_type(tid: int, name: str) -> tuple[int, str]:
    """note: if name is in use by another local type, that local type will be removed"""

    ordinal = ida_typeinf.get_tid_ordinal(tid)
    assert ordinal

    tinfo = ida_typeinf.tinfo_t()
    assert tinfo.get_numbered_type(None, ordinal)

    # TODO  @TH: perhaps tinfo.rename_type can be used instead? found out about it later
    err = tinfo.set_numbered_type(None, ordinal, ida_typeinf.NTF_REPLACE, name)
    errstr = ida_typeinf.tinfo_errstr(err)
    return err, errstr


def rename_local_type(
    cur_name: str, name: str, types: netnode.Netnode, uuids: dict[int, str],
) -> None:

    msg = f"renaming local type {cur_name!r} to {name!r}"
    logger.debug(msg)

    # resolve the type we're changing
    tid = ida_typeinf.get_named_type_tid(cur_name)
    assert tid != ida_idaapi.BADADDR

    # check if the new name is in use
    cur_tid = ida_typeinf.get_named_type_tid(name)
    name_changed = False
    try:
        # if the new name is already in use in the IDB --
        if cur_tid != ida_idaapi.BADADDR:
            # assert that the repo also has a different name for the type currently holding the new
            # name
            #
            # this is an assertion because if it's missing from the repo we should've already
            # deleted it
            #
            # also, we don't actually verify that the name since we already verified beforehand
            # that there are no duplicate names
            cur_uuid = types[cur_tid]
            assert cur_uuid in uuids

            # temporarily rename it to something else
            msg2 = f"temporarily renaming local type {name!r}"
            logger.debug("\t" + msg2)
            err, errstr = _rename_local_type(cur_tid, name + "_labsync_temp")

            # handle temporary rename failure
            if err:
                if logger.getEffectiveLevel() > logging.DEBUG:
                    logger.warning("failed " + msg2 + f": {errstr}")
                else:
                    logger.warning(f"\ttemporary rename failed: {errstr}")
                return False

        # now do the actual rename
        err, errstr = _rename_local_type(tid, name)

        # handle rename failure
        if err:
            if logger.getEffectiveLevel() > logging.DEBUG:
                logger.warning("failed " + msg + f": {errstr}")
            else:
                logger.warning(f"rename failed: {errstr}")
            return False

        name_changed = True
    finally:
        # if we failed, undo the temporary rename if we did any
        if cur_tid != ida_idaapi.BADADDR and not name_changed:
            msg2 = f"undoing the temporarily rename of local type {name!r}"
            logger.debug("\t" + msg2)
            err, errstr = _rename_local_type(cur_tid, name)

            # handle temporary rename undoing failure
            if err:
                if logger.getEffectiveLevel() > logging.DEBUG:
                    logger.warning("failed " + msg2 + f": {errstr}")
                else:
                    logger.warning(f"\ttemporary rename undoing failed: {errstr}")

    return bool(err)


def parse_local_types(nhdr: str) -> Generator[tuple[str, str, str, str]]:
    # split according to empty lines
    decls = re.split(r"\n\n", nhdr)

    for decl in decls:
        decl = decl.strip()

        # extract the uuid of the type
        uuid_line, decl = decl.split("\n", maxsplit=1)
        r = parse.parse(LOCAL_TYPES_COMMENT_FMT, uuid_line)
        if not r:
            msg = f"failed to extract uuid from local type uuid line:\n{uuid_line}"
            raise LabSyncError(msg)
        decl_uuid = r.fixed[0]

        # extract the name of the type and generate a forward declaration for it
        name, decl_type = decl_to_name_and_type(decl)

        yield name, decl_uuid, decl, decl_type


def update_local_types(nhdr: str, types: netnode.Netnode) -> None:
    # parse type declaration names
    name2decl = {}
    uuids = {}
    decls = []
    fdecls = []
    typedefs = []
    for name, decl_uuid, decl, decl_type in parse_local_types(nhdr):
        # make sure it's unique, mostly for sanity purposes
        if name in name2decl:
            if (name2decl[name] == decl and
                uuids.get(decl_uuid) == name):
                # duplicate local type with same decl and UUID. probably accidentally copied from
                # both sides during conflict resolution. we'll skip the redundant copy
                continue

            msg = f"found two local type declarations with the same name: {name}"
            raise LabSyncError(msg)
        name2decl[name] = decl

        # remember the name for each uuid
        uuids[decl_uuid] = name

        # accumulate
        if decl_type in {"struct", "union", "enum", "class"}:
            decls.append(decl)

            fdecl = f"{decl_type} {name};"
            fdecls.append(fdecl)
        elif decl_type == "typedef":
            typedefs.append(decl)
        else:
            msg = f"found unexpected kind of local type:\n{decl}"
            raise LabSyncError(msg)

    # remove local types if required
    for tid, decl_uuid in list(types.items()):
        if decl_uuid not in uuids:
            name = ida_typeinf.get_tid_name(tid)
            logger.debug(f"removing local type {name!r}")
            ida_typeinf.del_named_type(None, name, ida_typeinf.NTF_TYPE)
            del types[tid]

    # rename local types if required
        # create a mapping from uuid to type name in our IDB
    cur_uuid_to_name = {u: ida_typeinf.get_tid_name(t) for t, u in types.items()}

    for decl_uuid, name in uuids.items():
        # skip if uuid doesn't exist (this is a new type) or name didn't change
        cur_name = cur_uuid_to_name.get(decl_uuid)
        if not cur_name or cur_name == name:
            continue

        # do the renaming
        rename_local_type(cur_name, name, types, uuids)

    # create a reordered header that that should be parsable with regards to forward declarations
    hdr = "\n".join(fdecls) + "\n\n" + "\n".join(typedefs) + "\n\n" + "\n\n".join(decls)

    # iteratively try to load the header as long as dependencies get resolved
    #
    # TODO @TH: originally we didn't move the typedefs to before the decls, so honestly I think
    #           that now there should always be just one iteration here. we should probably verify
    #           that and remove the loop here afterwards
    last_n_errors = float("inf")
    iters = 1
    while iters < len(name2decl) + 1:
        # parse some more local types
        n_errors = ida_typeinf.parse_decls(None, hdr, None, ida_typeinf.HTI_DCL)
        logger.debug(f"loaded local types with {n_errors} errors")

        # add new uuids. we do this on every iteration in case we will eventually bail out --
        # we don't want to have added new types without keeping their uuids
        for decl_uuid, name in uuids.items():
            tid = ida_typeinf.get_named_type_tid(name)
            if tid != ida_idaapi.BADADDR:
                if tid not in types:
                    types[tid] = decl_uuid
                else:
                    assert types[tid] == decl_uuid

        # stop if there are no more errors or if we didn't add anything on this iteration
        if n_errors == 0 or n_errors >= last_n_errors:
            break

        last_n_errors = n_errors
        iters += 1
    else:
        msg = "local type loading took more than it makes sense. aborting"
        raise LabSyncError(msg)

    logger.debug(
        f"finished loading local types with {n_errors} errors after {iters} iterations "
        f"({len(name2decl)} types)"  # noqa: COM812
    )

    if n_errors:
        # IDA SDK doesn't properly export an interface for printer_t so we can't get the actual
        # errors :/
        fd, hdr_path = tempfile.mkstemp(suffix=".h", text=True)
        os.write(fd, hdr.encode("latin1"))
        os.close(fd)

        logger.error(
            "run the following in IDC shell to see the local types parsing errors:\n"
            f'\tparse_decls("{hdr_path}", PT_FILE)'  # noqa: COM812
        )

        msg = f"failed to parse local types ({n_errors} errors)"
        raise LabSyncError(msg)

    # in case loading the local types resulted in a new type being created, we might've encountered
    # a bug where IDA recreates an anonymous local type for an unnamed embedded subtype
    #
    # in that case, look for and delete dangling anonymous local types that should've been left
    new_names = set(local_types())
    if set(name2decl.keys()) != new_names:
        for name in new_names:
            tinfo = ida_typeinf.tinfo_t()
            assert tinfo.get_named_type(None, name)

            # check if anonymous
            if not tinfo.is_anonymous_udt():
                continue

            # check if it has any xrefs
            tid = tinfo.get_tid()
            if (ida_xref.get_first_cref_to(tid) != ida_idaapi.BADADDR or
                ida_xref.get_first_dref_to(tid) != ida_idaapi.BADADDR):
                continue

            # check if it's referenced by any typedef
            used_by_typedef = False

            any_tinfo = ida_typeinf.tinfo_t()
            for any_name in local_types():
                assert any_tinfo.get_named_type(None, any_name)

                if not any_tinfo.is_typedef():
                    continue

                if any_tinfo.get_next_type_name() != name:
                    continue

                used_by_typedef = True
                break

            if used_by_typedef:
                continue

            # remove it
            logger.debug(f"removing dangling local type {name!r}")
            ida_typeinf.del_named_type(None, name, ida_typeinf.NTF_TYPE)
            del types[tid]


def update_prototypes(
    binary: SyncedBinary, storage: functioninliner.ClonesStorage, d: dict[int, str]) -> None:

    # delete prototypes if required
    if REMOVE_MISSING_NAMES_AND_PROTOTYPES:
        # dump_prototypes() is a bit heavy, so we duplicate some code here instead of calling it
        # because we don't need the actual prototypes
        for ea in idautils.Functions():
            seg = ida_segment.getseg(ea)
            seg_name = ida_segment.get_segm_name(seg)

            # skip funcs that are in inlined chunks somehow (shouldn't happen)
            if seg_name.startswith("inlined_"):
                continue

            # skip funcs that have been inlined
            if ea in storage:
                continue

            # delete type if untyped in the new dict
            dea = binary.ea2dump(ea)
            if dea not in d:
                logger.debug(f"removing prototype from {ea:#x}")
                ida_nalt.del_tinfo(ea)

    # update prototypes
    for dea, ptype in d.items():
        ea = binary.dump2ea(dea)
        cur_ptype = prototype(ea)

        # TODO @TH: because we had numerous issues with prototypes syncing, we want to make sure
        #           that we always apply all prototypes so that issues will be raised early.
        #           we can probably undo it after we feel more confident
        # if cur_ptype != ptype:
        if True:
            msg = f"changing {ea:#x} prototype from:\n\t{cur_ptype!r}\nto\n\t{ptype!r}"
            if cur_ptype != ptype:
                logger.debug(msg)

            success = ida_typeinf.apply_cdecl(
                None, ea, ptype + ";", ida_typeinf.TINFO_DEFINITE,
            )

            if not success:
                if logger.getEffectiveLevel() > logging.DEBUG or cur_ptype == ptype:
                    logger.warning("failed " + msg)
                else:
                    logger.warning("prototype change failed!")


def migrate(d: dict, types: netnode.Netnode) -> dict:  # noqa: ARG001
    ver = d["version"]
    if ver not in {1, 2, 3, 4}:
        msg = f"data file is of unexpected version: {ver}"
        raise LabSyncError(msg)

    # version 4 is latest
    if ver == 4:
        return d

    # version 3 was missing local type uuids
        # we dropped the logic here since it was non-trivial to update per newer changes in code
        # and this version is very old anyway
    msg = ("migrating from YAML version 3 is not supported. use an older version of LabSync to do "
           "so")
    raise NotImplementedError(msg)

    if ver == 3:
        return d

    # version 2 was missing prototypes
    d["prototypes"] = dump_prototypes()

    if ver == 2:
        return d

    # version 1 was missing local types
    d["local_types"] = dump_local_types()

    return d


def parse_data(data: str, types: netnode.Netnode) -> dict:
    d = yaml.load(data, Loader=yaml.CLoader)  # noqa: S506
    return migrate(d, types)


def update(binary: SyncedBinary, data: str, types: netnode.Netnode) -> None:
    storage = functioninliner.ClonesStorage()
    storage.update_from_storage()

    d = parse_data(data, types)

    update_names(binary, storage, d["names"])
    update_inlined_funcs(binary, storage, d["inlined_funcs"])
    update_local_types(d["local_types"], types)
    update_prototypes(binary, storage, d["prototypes"])


def adopt_uuids(data: str, types: netnode.Netnode) -> None:
    # extract local types normalized header
    d = parse_data(data, types)
    nhdr = d["local_types"]

    # parse uuids
    name_to_uuid = {}
    for name, decl_uuid, _, _ in parse_local_types(nhdr):
        if name in name_to_uuid:
            msg = (
                f"found two local type declarations with the same name: {name} (uuids "
                f"{name_to_uuid[name]} and {decl_uuid})"
            )
            raise LabSyncError(msg)

        name_to_uuid[name] = decl_uuid

    # adopt uuids for local types we don't already have uuids for, if their names match
    for name in local_types():
        tid = ida_typeinf.get_named_type_tid(name)
        if tid == ida_idaapi.BADADDR:
            msg = f"failed to resolve tid of local type {name!r}"
            raise LabSyncError(msg)

        # skip local types that already have a uuid
        if tid in types:
            continue

        # adopt the uuid from the repo if it has a local type with the same name (best effort
        # heuristic)
        if name in name_to_uuid:
            decl_uuid = name_to_uuid[name]
            logger.debug(f"adopting repo uuid for local type {name!r}: {decl_uuid}")
            types[tid] = decl_uuid


# REPO


class LabSyncRepo:
    _path: str
    _repo: git.Repo

    def __init__(self, path: str):
        # TODO @TH: we should add a mechanism to lock the repo while we're doing git operations on
        #           it, in case multiple IDBs are using it at the same time

        self._path = path
        self._repo = git.Repo(path)

        with self._repo.config_reader() as cr:
            try:
                mergetool = cr.get("merge", "tool")
                logger.info(f"using merge.tool = {mergetool}")
            except (configparser.NoOptionError, configparser.NoSectionError):
                msg = "merge.tool must be set in git configuration to use this plugin"
                raise LabSyncError(msg)  # noqa: B904

        if self._repo.active_branch.name != "master":
            logger.warning("data repo is not on branch 'master'")

        with self._repo.config_writer() as cw:
            cw.set_value("mergetool", "keepbackup", "false")
            cw.set_value("mergetool", "writetotemp", "true")
            cw.set_value("mergetool", "hideresolved", "true")

    def _is_clean(self) -> bool:
        return not (self._repo.is_dirty() or self._repo.untracked_files)

    def _ensure_clean(self) -> None:
        if self._repo.is_dirty():
            msg = "data repo is dirty. please fix this externally for now"
            raise LabSyncError(msg)

        if self._repo.untracked_files:
            msg = "data repo has untracked files. please fix this externally for now"
            raise LabSyncError(msg)

    @property
    def path(self) -> str:
        return self._path

    def _pull(self) -> None:
        assert self._is_clean()

        # resolve remote tracking branch
        current = self._repo.active_branch
        tracking_branch = current.tracking_branch()
        if not tracking_branch:
            logger.warning("no remote tracking branch. skipping pull")
            return

        # pull
        try:
            self._repo.git.pull(rebase="false", allow_unrelated_histories=True)
        except git.GitCommandError as e:
            if "fix conflicts and then commit" not in e.stdout:
                raise

            logger.warning(f"pull failed with: {e.stdout}")
        else:
            return

        # resolve conflicts
        try:
            self._repo.git.mergetool("--no-prompt", "--gui")
        except git.GitCommandError as e:
            msg = (
                "mergetool returned failure. please fix the conflict and conclude the merge "
                "manually"
            )
            raise LabSyncError(msg) from e

        # conclude the merge
        self._repo.git.commit("--no-edit")

        # make sure that we're done
        if self._repo.is_dirty():
            msg = (
                "data repo is still dirty after supposedly committing the fixed merge. please "
                "fix the repo manually"
            )
            raise LabSyncError(msg)

    def _id_path(self, _id: str) -> pathlib.Path:
        return pathlib.Path(self.path) / f"{_id}.yaml"

    @contextlib.contextmanager
    def _open(self, _id: str, mode: str) -> IO:
        path = self._id_path(_id)
        with path.open(mode) as fp:
            yield fp

    def _commit(self, _id: str) -> None:
        index = self._repo.index
        index.add([_id + ".yaml"])
        # this is broken for some reason: index.commit(f"updated {_id}")
        self._repo.git.commit("-m", f"updated {_id}")

    def _commit_dangling(self, _id: str) -> None:
        index = self._repo.index
        index.add([_id + ".yaml"])
        index.commit(f"initial commit for {_id}", [])

    def _push(self) -> None:
        assert self._is_clean()

        # resolve remote tracking branch
        current = self._repo.active_branch
        tracking_branch = current.tracking_branch()
        if not tracking_branch:
            logger.warning("no remote tracking branch. skipping push")
            return

        # push
        self._repo.git.push()

    def get(self, _id: str) -> tuple[Optional[str], str]:
        self._ensure_clean()

        self._pull()

        commit = self._repo.commit().hexsha

        if self._id_path(_id).exists():
            with self._open(_id, "rt") as fp:
                data = fp.read()
        else:
            data = None

        return data, commit

    def put(self, _id: str, content: str, *, base: Optional[str]) -> tuple[str, bool]:
        self._ensure_clean()

        if base and self._repo.head.commit.hexsha != base:
            # we need to fetch before resetting, because if the IDB was copied from somewhere,
            # it could be that we don't locally have the base commit
            logger.debug("fetching repo")
            self._repo.git.fetch()

            logger.debug(f"resetting HEAD to {base}")
            self._repo.head.reset(base, working_tree=True)

        with self._open(_id, "wt") as fp:
            fp.write(content)

        if self._is_clean():
            logger.debug("no changes")

            changed = False
        else:
            if base:
                self._commit(_id)
            else:
                logger.debug("committing to a danging commit")
                self._commit_dangling(_id)

            changed = True

        return self._repo.commit().hexsha, changed

    def sync(self) -> str:
        self._pull()
        self._push()

        return self._repo.commit().hexsha

    def ping(self) -> bool:
        try:
            self._repo.git.ls_remote(heads=True)
        except git.GitCommandError:
            return False
        else:
            return True


# PLUGIN STUFF


class LabSyncActionBase(ida_kernwin.action_handler_t):
    plugin: ida_idaapi.plugin_t

    def __init__(self, plugin: ida_idaapi.plugin_t):
        super().__init__()
        self.plugin = plugin

    @property
    def name(self) -> str:
        return f"{self.plugin.wanted_name}:{self.__class__.__name__}"

    @property
    def label(self) -> str:
        raise NotImplementedError

    @property
    def shortcut(self) -> Optional[str]:
        return None

    @property
    def tooltip(self) -> Optional[str]:
        return None

    @property
    def icon(self) -> int:
        return 0

    @property
    def flags(self) -> int:
        return 0

    @property
    def path(self) -> str:
        return f"Edit/Plugins/{self.plugin.wanted_name}/"

    def register(self) -> None:
        desc = ida_kernwin.action_desc_t(
            self.name,
            self.label,
            self,
            self.shortcut,
            self.tooltip,
            self.icon,
        )
        ida_kernwin.register_action(desc)

    def unregister(self) -> None:
        ida_kernwin.unregister_action(self.name)

    def activate(self, ctx: Any) -> None:
        raise NotImplementedError

    def update(self, ctx: Any) -> None:
        raise NotImplementedError


class LabSyncEnableAction(LabSyncActionBase):
    @property
    def label(self) -> str:
        return "Enable"

    def activate(self, ctx: Any) -> None:
        self.plugin.enable()

    def update(self, ctx: Any) -> None:
        if self.plugin.enabled:
            return ida_kernwin.AST_DISABLE
        else:
            return ida_kernwin.AST_ENABLE


class LabSyncDisableAction(LabSyncActionBase):
    @property
    def label(self) -> str:
        return "Disable"

    def activate(self, ctx: Any) -> None:
        self.plugin.disable()

    def update(self, ctx: Any) -> None:
        if self.plugin.enabled:
            return ida_kernwin.AST_ENABLE
        else:
            return ida_kernwin.AST_DISABLE


class LabSyncResetAction(LabSyncActionBase):
    @property
    def label(self) -> str:
        return "Reset"

    def activate(self, ctx: Any) -> None:
        self.plugin.reset()

    def update(self, ctx: Any) -> None:
        if self.plugin.enabled:
            return ida_kernwin.AST_DISABLE
        else:
            return ida_kernwin.AST_ENABLE


class LabSyncHooks(ida_kernwin.UI_Hooks):
    plugin: ida_idaapi.plugin_t
    menu_actions: list[LabSyncActionBase]

    def __init__(self, plugin: ida_idaapi.plugin_t, menu_actions: list[LabSyncActionBase]):
        super().__init__()

        self.plugin = plugin
        self.menu_actions = menu_actions

    def ready_to_run(self) -> None:
        for action in self.menu_actions:
            ida_kernwin.attach_action_to_menu(
                action.path, action.name, ida_kernwin.SETMENU_APP,
            )

    def saved(self) -> None:
        if not self.plugin.enabled:
            return

        if self.plugin.sync_in_progress:
            return

        # TODO @TH: we need to figure out a way not to have this run on "Save as...".
        #           or if we do want it to run -- get_path(PATH_TYPE_IDB) here will still hold the
        #           "old" path so reinvoke the 2nd save differently
        changed = self.plugin.sync()
        if changed:
            logger.debug("resaving the idb with synced changes")
            idb_name = ida_loader.get_path(ida_loader.PATH_TYPE_IDB)
            # TODO @TH: can we somehow find the flags the original save was invoked with and reuse
            #           them?
            ida_loader.save_database(idb_name, 0)


class LabSyncPlugin(ida_idaapi.plugin_t):
    version: int = ida_idp.IDP_INTERFACE_VERSION
    flags: int = ida_idaapi.PLUGIN_MOD | ida_idaapi.PLUGIN_HIDE

    comment: str = "helps to partially synchronize IDBs over git"
    help: str = ""
    wanted_name: str = "LabSync"
    wanted_hotkey: str = ""

    menu_actions_types: list[type[LabSyncActionBase], ...] = \
        (LabSyncEnableAction, LabSyncDisableAction, LabSyncResetAction)

    CFGFILE: str = "labsync.cfg"
    NETNODE: str = "$ labsync.plugin"
    TYPES_NETNODE: str = "$ labsync.types"
    BINARIES_NETNODE: str = "$ labsync.binaries"

    COMMIT_KEY_PREFIX = "commit."

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.menu_actions = []
        self.hooks = None

        self._sync_in_progress = False

    @classmethod
    @property
    @functools.cache
    def netnode(cls) -> netnode.Netnode:
        """this netnode holds:
            commit.<idb_id>: Optional[str]  # the git commit this binary is synced to
            enabled: bool  # whether LabSync is enabled
            custom_idb_id: Optional[str]  # override of the IDB id to use
        """
        return netnode.Netnode(LabSyncPlugin.NETNODE)

    @classmethod
    @functools.cache
    def types(cls, idb_id: str) -> netnode.Netnode:
        """this netnode matches:
                tid [int]  # of a local type
            to
                uuid [str]  # of it in the YAML
        """
        name = f"{LabSyncPlugin.TYPES_NETNODE}.{idb_id}"
        return netnode.Netnode(name)

    @classmethod
    @property
    @functools.cache
    def _binaries_netnode(cls) -> netnode.Netnode:
        """this netnode matches:
                seg_prefix [str]  # prefix of segment names
            to
                tuple[str, int]  # of (<idb id>, <base ea>) to use for mapping segments to separate
                                 # YAMLs
        """
        return netnode.Netnode(LabSyncPlugin.BINARIES_NETNODE)

    @classmethod
    @property
    @functools.cache
    def cfg(cls) -> dict:
        path = (
            pathlib.Path(ida_diskio.get_user_idadir()) / ida_diskio.CFG_SUBDIR /
            LabSyncPlugin.CFGFILE
        )
        try:
            cfg_raw = path.open("rt").read()
        except OSError as e:
            logger.error(f"failed to read configuration file from {path!r}: {e}")  # noqa: TRY400
            return None

        cfg_parser = configparser.ConfigParser()
        try:
            cfg_parser.read_string("[section]\n" + cfg_raw)
        except Exception:
            logger.exception("failed to parse configuration file")
            return None

        return dict(cfg_parser["section"].items())

    @classmethod
    @property
    @functools.cache
    def repo(cls) -> LabSyncRepo:
        # validate/apply basic configuration
        if "repo_path" not in cls.cfg:
            logger.error("repo_path not found in configuration file")
            return None

        try:
            repo = LabSyncRepo(cls.cfg["repo_path"])
        except Exception:
            logger.exception("configured repo_path does not point to a valid git repo")
            return None

        logger.info(f"using repo: {repo.path}")

        return repo

    @classmethod
    @property
    def idb_id(cls) -> str:
        default_id = ida_nalt.retrieve_input_file_md5().hex()
        return cls.netnode.get("custom_idb_id", default_id)

    @classmethod
    def commit_key(cls, idb_id: str) -> str:
        return f"{cls.COMMIT_KEY_PREFIX}{idb_id}"

    def _register(self) -> None:
        for t in LabSyncPlugin.menu_actions_types:
            a = t(self)
            a.register()
            self.menu_actions.append(a)

        self.hooks = LabSyncHooks(self, self.menu_actions)
        self.hooks.hook()

    def _deregister(self) -> None:
        if self.hooks:
            self.hooks.unhook()

        for a in self.menu_actions:
            a.unregister()

    @staticmethod
    def is_compatible() -> bool:
        info = ida_idaapi.get_inf_structure()
        return info.procname == "ARM" and info.is_64bit()

    @staticmethod
    def _init_logging() -> None:
        logger_formatter = logging.Formatter(
            fmt="{name}.{levelname:<5s}: {message}", style="{",
        )

        logger_hdlr = logging.StreamHandler()
        logger_hdlr.setFormatter(logger_formatter)

        logger.addHandler(logger_hdlr)

        # may be overridden with the 'log' config entry
        logger.setLevel(logging.INFO)

    def migrate(self) -> None:
        # migrate types to be per binary
        old_types = netnode.Netnode(LabSyncPlugin.TYPES_NETNODE)
        new_types = self.types(self.idb_id)
        for tid, _uuid in list(old_types.items()):
            new_types[tid] = _uuid
            del old_types[tid]

        # migrate commit to be per binary
        old_commit = self.netnode.get("commit")
        if old_commit is not None:
            self.netnode[self.commit_key(self.idb_id)] = old_commit
            del self.netnode["commit"]

    def init(self) -> int:
        LabSyncPlugin._init_logging()

        if not self.cfg or not self.repo:
            return ida_idaapi.PLUGIN_SKIP

        for binary in self.binaries:
            if binary.seg_prefix is None:
                name = "main"
            else:
                name = f"{binary.seg_prefix}*"
            logger.info(f"{name} idb id: {binary.idb_id}")

        if self.enabled:
            logger.info("syncing enabled")
        else:
            logger.info("syncing disabled")

        self.migrate()

        for binary in self.binaries:
            commit = self.netnode.get(self.commit_key(binary.idb_id))
            if self.enabled or commit:
                logger.info(
                    f"{binary.idb_id} last synced to commit: {commit}",
                )

        log_level = self.cfg.get("log")
        if log_level:
            logger.setLevel(log_level)

        self._register()

        logger.info("initialized successfully")

        return ida_idaapi.PLUGIN_KEEP

    def term(self) -> None:
        self._deregister()

    def run(self, arg: int = 0) -> None:
        pass

    @classmethod
    @property
    def _lockfile_path(cls) -> pathlib.Path:
        return pathlib.Path(__file__).parent / LOCKFILE

    @contextlib.contextmanager
    def _lock(
        self, wait_timeout: Optional[float] = None,
    ) -> contextlib.AbstractContextManager[None]:
        """it makes more sense to logically lock the LabSyncRepo and not LabSyncPlugin, but it's
        problematic because we want to use a lockfile which orignally wasn't in the .gitignore of
        that repo (and so syncing older IDBs will fail because we'll create the lockfile and that
        will make the repo dirty)
        """

        deadline = time.monotonic() + wait_timeout
        while True:  # we want to try once even if the timeout is 0
            logger.debug("taking lock on the repo")
            try:
                fd = os.open(self._lockfile_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
                break
            except FileExistsError:
                pass

            if time.monotonic() >= deadline:
                logger.error(
                    "failed to get lock on the data repository within the configured timeout. "
                    "either another IDA instance is still syncing, or something bad happened "
                    "during a previous sync. if you're sure it's safe, you can delete the "
                    f"lockfile manually: {self._lockfile_path}"  # noqa: COM812
                )
                msg = "failed to get repo lock"
                raise LabSyncLockError(msg)

        try:
            yield
        finally:
            logger.debug("releasing lock on the repo")
            os.close(fd)
            self._lockfile_path.unlink()

    @property
    def sync_in_progress(self) -> bool:
        return self._sync_in_progress

    @classmethod
    @property
    def binaries(cls) -> tuple[SyncedBinary, ...]:
        rules = dict(cls._binaries_netnode.items())
        return cls._binaries_from_mapping_rules(rules)

    @classmethod
    def _binaries_from_mapping_rules(
        cls, rules: dict[str, tuple[str, Optional[int]]]) -> tuple[SyncedBinary, ...]:

        if not rules:
            return (SyncedBinary(idb_id=cls.idb_id),)

        # generate a synced binary per every segment matching rule and the default
        seg2bin = {}
        for seg_prefix, (idb_id, base_ea) in rules.items():
            # start with start_ea/end_ea pointing at max/min, and we'll have update them in the
            # next loop
            seg2bin[seg_prefix] = SyncedBinary(
                idb_id=idb_id, start_ea=ida_idaapi.BADADDR, end_ea=0, base_ea=base_ea,
                seg_prefix=seg_prefix,
            )
        default_bin = SyncedBinary(
            idb_id=cls.idb_id, start_ea=ida_idaapi.BADADDR, end_ea=0, seg_prefix=None,
        )

        # find the start/end EAs for each synced binary
        for i in range(ida_segment.get_segm_qty()):
            seg = ida_segment.getnseg(i)
            name = ida_segment.get_segm_name(seg)

            # skip inlined chunks segments
            if name.startswith("inlined_"):
                continue

            matched = False
            for seg_prefix, binary in seg2bin.items():
                if name.startswith(seg_prefix):
                    binary.start_ea = min(binary.start_ea, seg.start_ea)
                    binary.end_ea = max(binary.end_ea, seg.end_ea)
                    matched = True

            if not matched:
                default_bin.start_ea = min(default_bin.start_ea, seg.start_ea)
                default_bin.end_ea = max(default_bin.end_ea, seg.end_ea)

        # remove synced binaries that didn't match any segment
        for seg, _bin in list(seg2bin.items()):
            if _bin.start_ea == ida_idaapi.BADADDR:
                del seg2bin[seg]
        if default_bin.start_ea == ida_idaapi.BADADDR:
            assert default_bin.end_ea == 0
            default_bin.start_ea = 0  # make it empty

        # make sure that none of the synced binaries overlap
        last_bin = None
        last_seg = None
        for _bin, seg in sorted(
                [(_bin, seg) for (seg, _bin) in seg2bin.items()] + [(default_bin, "<default>")],
            ):

            assert _bin.start_ea <= _bin.end_ea
            if last_bin:
                assert last_bin.start_ea <= _bin.start_ea
                if last_bin.end_ea > _bin.start_ea:
                    msg = f"Found overlapping segments for {last_seg!r} and {seg!r}"
                    raise LabSyncBinaryMatchingError(msg)

            last_bin = _bin
            last_seg = seg

        # validate/set base_ea-s
        for seg, _bin in seg2bin.items():
            if _bin.base_ea is None:
                _bin.base_ea = _bin.start_ea
            elif _bin.base_ea > _bin.start_ea:
                msg = f"Base EA for {seg!r} is greater than its start EA"
                raise LabSyncBinaryMatchingError(msg)

        binaries = (default_bin, *seg2bin.values())

        # make sure the we don't have overlapping IDB ids
        assert len(binaries) == len({b.idb_id for b in binaries})

        return binaries

    @classmethod
    def map_segments_to_idb_id(
        cls, seg_prefix: str, idb_id: str, *, base_ea: Optional[int] = None) -> None:

        rules = dict(cls._binaries_netnode.items())

        new_rules = rules.copy()
        seg_rule = (idb_id, base_ea)
        new_rules[seg_prefix] = seg_rule

        # check that the new rules are valid
        try:
            cls._binaries_from_mapping_rules(new_rules)
        except LabSyncBinaryMatchingError:
            logger.exception(f"failed mapping segment prefix {seg_prefix!r} to idb id {idb_id}")
            return

        # add the rule to netnode
        cls._binaries_netnode[seg_prefix] = seg_rule

        # invalidate previous types if we had them (since they might be outdated)
        types = cls.types(idb_id)
        for tid in list(types.keys()):
            del types[tid]

    def sync(self) -> None:
        assert not self.sync_in_progress

        try:
            self._sync_in_progress = True

            logger.debug("checking that the remote repository is reachable")
            if not self.repo.ping():
                logger.warning("skipping sync because the remote repository is unreachable")
                return False

            timeout = float(self.cfg.get("lock_timeout_sec", DEFAULT_LOCK_TIMEOUT))
            with self._lock(timeout):
                binaries = self.binaries
                for binary in self.binaries:
                    if len(binaries) > 1:
                        logger.debug(f"syncing {binary.idb_id}")

                    # in case this is the initial commit, try to adopt local types from the repo
                    # (in case this IDB was already synced)
                    base = self.netnode.get(self.commit_key(binary.idb_id))
                    if not base:
                        latest_data, latest_commit = self.repo.get(binary.idb_id)
                        if latest_data:
                            logger.debug(f"adopting local type uuids from commit: {latest_commit}")
                            adopt_uuids(latest_data, self.types(binary.idb_id))

                    # sync IDB -> repo
                    logger.debug("saving changes to repo")

                    data = dump(binary, self.types(binary.idb_id))
                    commit, changed = self.repo.put(binary.idb_id, data, base=base)

                    if changed:
                        # save the IDB with the updated types and commit, in case something will
                        # break later during the merge
                        #
                        # TODO @TH: this path will be wrong if we're under "Save as..." flow.
                        #           see comments under LabSyncHooks.saved
                        self.netnode[self.commit_key(binary.idb_id)] = commit

                        idb_name = ida_loader.get_path(ida_loader.PATH_TYPE_IDB)
                        ida_loader.save_database(idb_name, 0)

                    # sync repo with upstream
                    self.repo.sync()

                    # sync repo -> IDB
                    new_data, new_commit = self.repo.get(binary.idb_id)

                        # if there's a new commit with new data for our IDB -- update
                    if new_data != data:
                        logger.debug("updating with changes from repo")
                        with wait_box("updating with changes from repo...", hide_cancel=True):
                            update(binary, new_data, self.types(binary.idb_id))
                            changed = True

                        # save the commit we're synced to in netnode (even if just the commit
                        # changed)
                    if new_commit != commit:
                        self.netnode[self.commit_key(binary.idb_id)] = new_commit
                        changed = True

                    p = "" if len(binaries) == 1 else f"{binary.idb_id} "
                    s = "" if changed else " (no changes)"
                    logger.info(f"{p}synced to commit: {new_commit}{s}")

                return changed
        except BaseException as exc:  # noqa: BLE001
            ida_kernwin.warning("LabSync synchronization failed. See output window for more "
                                "details")
            if not isinstance(exc, LabSyncLockError):
                raise
        finally:
            self._sync_in_progress = False

    @classmethod
    def enable(cls) -> None:
        cls.netnode["enabled"] = True
        logger.info("syncing enabled")

        for binary in cls.binaries:
            commit = cls.netnode.get(cls.commit_key(binary.idb_id))
            logger.info(
                f"{binary.idb_id} last synced to commit: {commit}",
            )

    @classmethod
    def disable(cls) -> None:
        cls.netnode["enabled"] = False
        logger.info("syncing disabled")

    @classmethod
    def reset(cls) -> None:
        # make sure that we're already disabled, to make sure that no one resets by mistake
        assert not cls.enabled

        idb_ids = {}

        # delete the commit states
        for k in cls.netnode.keys():  # noqa: SIM118
            if k.startswith(cls.COMMIT_KEY_PREFIX):
                idb_ids.add(k[len(cls.COMMIT_KEY_PREFIX):])
                del cls.netnode[k]

        for binary in cls.binaries:
            idb_ids.add(binary.idb_id)

        # delete the tid->uuid mapping for all of the IDB ids we noted
        for idb_id in idb_ids:
            types = cls.types(idb_id)
            for tid in list(types.keys()):
                del types[tid]

        logger.info("sync data reset")

    @classmethod
    @property
    def enabled(cls) -> bool:
        return cls.netnode.get("enabled", False)


def PLUGIN_ENTRY() -> ida_idaapi.plugin_t:  # noqa: N802
    return LabSyncPlugin()


repository: Mec4nico/Radio_IC-980Pro
file: stm32 ida v0_1/firmloader.py
wanted_name: FirmLoader
comment: FirmLoader
created_at: 2025-01-19 14:04:07+00:00
pushed_at: 2025-02-11 12:42:20+00:00
description:
Reversing project of chinese IC-980Pro Radio to understand all the hardware and try to add some functionalities.

url:
https://github.com/Mec4nico/Radio_IC-980Pro/blob/e2c77fb832324ba2c7851fb895f708cc9eb9ad1d/stm32%20ida%20v0_1/firmloader.py

content:
import os
import json
import idaapi
import idc
import ida_funcs
import ida_segment
import ida_segment
import ida_xref
import ida_bytes
import idautils
import ida_kernwin


class FirmwareChooser(ida_kernwin.Choose):
    def __init__(self, firmwares):
        ida_kernwin.Choose.__init__(self, "Select firmware", [["Name", 30 | ida_kernwin.CHCOL_PLAIN]], flags = ida_kernwin.CH_QFLT)
        self.firmware_list = firmwares

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

    def OnGetLine(self, n):
        return [self.firmware_list[n]["name"]]


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

    def activate(self, ctx):
        firmware_list = []
        for file_name in sorted(os.listdir(os.path.join(os.path.dirname(os.path.realpath(__file__)),"firmloader_data"))):
            with open(os.path.join(os.path.dirname(os.path.realpath(__file__)),"firmloader_data",file_name)) as data_file:
                data = json.load(data_file)
                key = file_name.replace(".json","")
                if data["family"]:
                    display_name = data["brand"] + "/" + data["family"] + "/"+ data["name"]
                else:
                    display_name = data["brand"] + "/" + data["name"]

                firmware_list.append({"key": key, "name": display_name})

        chooser = FirmwareChooser(firmware_list)
        sel = chooser.Show(True)
        #if (sel := chooser.Show(True)) == -1:
        #    return

        chosen_key = firmware_list[sel]["key"]

        found_rom = False
        rom_segment = 0
        for segment_ea in idautils.Segments():
                # Default ROM segment
                if idc.get_segm_name(segment_ea) == "ROM":
                    found_rom = True
                    rom_segment = segment_ea
        # If there is no ROM segment warn user and close
        if not found_rom:
            ida_kernwin.warning("Segment with name \"ROM\" does not exist. Please, rename the main code segment to \"ROM\" and run the plugin again.")
            return
        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "firmloader_data", chosen_key + ".json")) as data_file:
            current_mcu = json.load(data_file)
            # Create segments
            for segment in current_mcu["segments"]:
                start = int(segment["start"],16)
                end = int(segment["end"],16)
                ida_segment.add_segm(0,start,end,segment["name"],segment["type"],0)
                # Set mode
            for segment_ea in idautils.Segments():
                # Set mode Thumb/ARM if ARM architecture
                if idaapi.get_inf_structure().procname == 'ARM' and found_rom:
                    idaapi.split_sreg_range(rom_segment, idaapi.str2reg("T"), current_mcu["mode"], idaapi.SR_user)
            # Create peripherals
            if ida_kernwin.ask_yn(1, "Would you like to load peripherals?"):
                for peripheral in current_mcu["peripherals"]:
                    # Start of the peripheral struct
                    start = int(peripheral["start"],16)
                    end = int(peripheral["end"],16)
                    ida_segment.add_segm(0,start,end,peripheral["name"],"DATA",0)
                    if peripheral["registers"]:
                        for register in peripheral["registers"]:
                            offset = int(register["offset"],16)
                            idc.set_name(start + offset, f'{peripheral["name"]}_{register["name"]}', idc.SN_NOCHECK)
                    # Add comment if any
                    idc.set_cmt(start,peripheral["comment"],False)
            # Crate vector table
            if current_mcu["vector_table"]:
                if ida_kernwin.ask_yn(1, "Would you like to also populate the vector table?"):
                    vector_address = ida_kernwin.ask_str(hex(rom_segment),0,"Set vector table at address:")
                    try:
                        if "x" in vector_address:
                            vector_address = int(vector_address,16)
                        else:
                            vector_address = int(vector_address)
                    except:
                        ida_kernwin.info("Invalid address!")
                        return
                    if not idc.get_segm_name(vector_address):
                        ida_kernwin.info("Invalid address!")
                        return
                    else:
                        for vector in current_mcu["vector_table"]:
                            # 
                            addr = vector_address + int(vector["addr"],16)
                            # Create a data value
                            if current_mcu["bits"] == 8:
                                ida_bytes.create_byte(addr,1)
                            elif current_mcu["bits"] == 16:
                                ida_bytes.create_qword(addr,2)
                            elif current_mcu["bits"] == 32:
                                ida_bytes.create_dword(addr,4)
                            elif current_mcu["bits"] == 64:
                                ida_bytes.create_qword(addr,8)

                            # Set name of the address
                            idc.set_name(addr, vector["name"]+"_vector", idc.SN_NOCHECK)
                            # Set comment
                            idc.set_cmt(addr,vector["comment"],False)
                            # Get XREF destination
                            destination = ida_bytes.get_32bit(addr)
                            # If ARM account for Thumb mode
                            if idaapi.get_inf_structure().procname == 'ARM' and destination % 2 == 1: # Thumb mode
                                destination = destination - 1
                            # If the segment is CODE
                            if ida_segment.get_segm_class(ida_segment.getseg(destination)) == "CODE":
                                # Create code reference
                                ida_xref.add_cref(addr,destination,1)
                                # If the name does not exist or is unk_*
                                if idc.get_name(destination) == "" or idc.get_name(destination).startswith("unk_"):
                                    ida_funcs.add_func(destination)
                                    idc.set_name(destination, vector["name"] + "_handler", idc.SN_NOCHECK | idc.SN_PUBLIC)
            # All is done at this point, mark the ROM segment for auto-analysis
            if found_rom:
                idc.auto_mark_range(rom_segment,idc.get_segm_end(rom_segment),idaapi.AU_WEAK) # Let IDA decide what is code and what is not

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

class firmloader_t(idaapi.plugin_t):
    comment = "FirmLoader"
    help = "This script helps when working with various microcontroller firmware files."
    wanted_name = "FirmLoader"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        load_action_id = "firmloader:load"
        load_action = idaapi.action_desc_t(
            load_action_id,
            "Apply firmware info...",
            Firmloader(),
            '',
            "Apply known firmware info to the database."
        )

        idaapi.register_action(load_action)
        idaapi.attach_action_to_menu("Edit/FirmLoader/Apply firmware info...", load_action_id, idaapi.SETMENU_APP)

    def run(self):
        pass

    def term(self):
        pass

def PLUGIN_ENTRY():
    return firmloader_t()


repository: Thehepta/D810Learn
file: D810_Extern/genmy.py
comment: Display microcode
created_at: 2024-06-02 06:36:20+00:00
pushed_at: 2025-05-11 10:47:50+00:00
url:
https://github.com/Thehepta/D810Learn/blob/52ad29245037f279411a2ddd6b8b6e3f15e9bc9a/D810_Extern/genmy.py

content:
# genmc - Display Hex-Rays Microcode
#
# Requires IDA and decompiler(s) >= 7.3
#
# Based on code/ideas from:
# - vds13.py from Hexrays SDK
# - https://github.com/RolfRolles/HexRaysDeob
# - https://github.com/NeatMonster/MCExplorer

__author__ = "Dennis Elser"

# -----------------------------------------------------------------------------
import os, shutil, errno, re

import ida_idaapi
import ida_bytes
import ida_range
import ida_kernwin as kw
import ida_hexrays as hr
import ida_funcs
import ida_diskio
import ida_ida
import ida_graph
import ida_lines
import ida_moves
from graphviz import Digraph

PLUGIN_NAME = "genmc"
global microcode_viewer_instance
microcode_viewer_instance = None
# -----------------------------------------------------------------------------
def get_mcode_name(mcode):
    """returns the name of the mcode_t passed in parameter."""
    for x in dir(hr):
        if x.startswith('m_'):
            if mcode == getattr(hr, x):
                return x
    return None


# -----------------------------------------------------------------------------
def get_mopt_name(mopt):
    """returns the name of the mopt_t passed in parameter."""
    for x in dir(hr):
        if x.startswith('mop_'):
            if mopt == getattr(hr, x):
                return x
    return None


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


# -----------------------------------------------------------------------------
def get_target_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_target_filename())


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

def install_plugin():
    """Installs script to IDA userdir as a plugin."""
    if is_plugin():
        kw.msg("Command not available. Plugin already installed.\n")
        return False

    src = __file__
    if is_installed():
        btnid = kw.ask_yn(kw.ASKBTN_NO,
                          "File exists:\n\n%s\n\nReplace?" % get_target_filename())
        if btnid is not kw.ASKBTN_YES:
            return False
    dst = get_target_filename()
    usrdir = os.path.dirname(dst)
    kw.msg("Copying script from \"%s\" to \"%s\" ..." % (src, usrdir))
    if not os.path.exists(usrdir):
        try:
            os.path.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 IDA instance now.\n"))
    return True


# -----------------------------------------------------------------------------
def is_ida_version(requested):
    """Checks minimum required IDA version."""
    rv = requested.split(".")
    kv = kw.get_kernel_version().split(".")

    count = min(len(rv), len(kv))
    if not count:
        return False

    for i in range(count):
        if int(kv[i]) < int(rv[i]):
            return False
    return True

# -----------------------------------------------------------------------------
class printer_t(hr.vd_printer_t):
    """Converts microcode output to an array of strings."""

    def __init__(self, *args):
        hr.vd_printer_t.__init__(self)
        self.mc = []

    def get_mc(self):
        return self.mc

    def _print(self, indent, line):    # 被东调用，一行的传入microcode反编译结果
        self.mc.append(line)
        return 1

def dominanceFlow(dispatch_block):
    # import pydevd_pycharm
    # pydevd_pycharm.settrace('localhost', port=31235, stdoutToServer=True, stderrToServer=True)

    print("dispatch_block serial:",hex(dispatch_block.serial))
    blk_preset_list = [x for x in dispatch_block.predset]
    print("dispatch_block father list:",blk_preset_list)

    def dfs( current_node, target_node, path, paths, visited):
        path.append(current_node.serial)
        visited.add(current_node.serial)

        for neighbor in current_node.succs():
            if neighbor.serial == target_node.serial and len(path) > 1:
                paths.append(list(path))
            elif neighbor.serial not in visited:
                dfs(neighbor, target_node, path, paths, visited)

        path.pop()
        visited.remove(current_node.serial)
    paths = []
    dfs(dispatch_block,dispatch_block,[],paths,set())
    for path in paths:
        print("path:",path)


def graphviz(mba,output_path):
    dot = Digraph()
    dot.attr(splines='ortho')
    for blk_idx in range(mba.qty):
        blk = mba.get_mblock(blk_idx)
        if blk.head == None:
            continue
        lines = []

        lines.append("{0}:{1}".format(blk_idx, hex(blk.head.ea)))
        insn = blk.head
        while insn:
            lines.append(insn.dstr())
            if insn == blk.tail:
                break
            insn = insn.next
        label = "\n".join(lines)
        dot.node(str(blk_idx), label=label, shape="rect", style="filled", fillcolor="lightblue")

    for blk_idx in range(mba.qty):
        blk = mba.get_mblock(blk_idx)
        succset = [x for x in blk.succset]
        for succ in succset:
            blk_succ = mba.get_mblock(succ)
            if blk_succ.head is None:
                continue
            if blk.head is None:
                continue
            dot.edge(str(blk_idx), str(succ))

    # dot.render("/home/chic/graph_with_contentgraph_with_content", format="png")
    with open(output_path, "w") as f:
        f.write(dot.source)
    # dot.render("/home/chic/graph_with_content", format="png")
    print("dot已保存到 :",output_path)


# -----------------------------------------------------------------------------
# class microcode_insnviewer_t(ida_graph.GraphViewer):
#     """Displays the graph view of Hex-Rays microcode."""
#     def __init__(self, mba, mmat_name, fn_name, block, serial):
#         title = "Microinstruction: %s - %d.%d (%s)" % (fn_name, block, serial, mmat_name)
#         ida_graph.GraphViewer.__init__(self, title, True)
#         self.mblock = mba.get_mblock(block)
#         self.minsn = self.get_minsn(serial)
#
#     def get_minsn(self, serial):
#         curr = self.mblock.head
#         for i in range(serial):
#             curr = curr.next
#         return curr
#
#     def _insert_mop(self, mop, parent):
#         if mop.t == 0:
#             return -1
#
#         text = get_mopt_name(mop.t) + '\n' + mop._print()
#         node_id = self.AddNode(text)
#         self.AddEdge(parent, node_id)
#
#         if mop.t == hr.mop_d: # result of another instruction
#             dst_id = self._insert_minsn(mop.d)
#             if dst_id:
#                 self.AddEdge(node_id, dst_id)
#         elif mop.t == hr.mop_f: # list of arguments
#             for arg in mop.f.args:
#                 self._insert_mop(arg, node_id)
#         elif mop.t == hr.mop_a: # mop_addr_t: address of operand
#             self._insert_mop(mop.a, node_id)
#         elif mop.t == hr.mop_p: # operand pair
#             self._insert_mop(mop.pair.lop, node_id)
#             self._insert_mop(mop.pair.hop, node_id)
#         return node_id
#
#     def _insert_minsn(self, minsn):
#         if minsn:
#             text = get_mcode_name(minsn.opcode) + '\n' + minsn._print()
#             node_id = self.AddNode(text)
#
#             self._insert_mop(minsn.l, node_id)
#             self._insert_mop(minsn.r, node_id)
#             self._insert_mop(minsn.d, node_id)
#             return node_id
#         return None
#
#     def OnRefresh(self):
#         self.Clear()
#         self._insert_minsn(self.minsn)
#         return True
#
#     def OnGetText(self, node_id):
#         return self._nodes[node_id]

# -----------------------------------------------------------------------------
class microcode_graphviewer_t(ida_graph.GraphViewer):
    """Displays the graph view of Hex-Rays microcode."""

    def __init__(self, mba, title, lines):
        title = "Microcode graph: %s" % title
        ida_graph.GraphViewer.__init__(self, title, True)
        self._mba = mba
        self._mba.set_mba_flags(hr.MBA_SHORT)
        self._process_lines(lines)
        if mba.maturity == hr.MMAT_GENERATED or mba.maturity == hr.MMAT_PREOPTIMIZED:
            mba.build_graph()

    def _process_lines(self, lines):
        self._blockcmts = {}
        curblk = "-1"
        self._blockcmts[curblk] = []
        for i, line in enumerate(lines):
            plain_line = ida_lines.tag_remove(line).lstrip()
            if plain_line.startswith(';'):
                # msg("%s" % plain_line)
                re_ret = re.findall("BLOCK ([0-9]+) ", plain_line)
                if len(re_ret) > 0:
                    curblk = re_ret[0]
                    self._blockcmts[curblk] = [line]
                else:
                    self._blockcmts[curblk].append(line)
        if "0" in self._blockcmts:
            self._blockcmts["0"] = self._blockcmts["-1"] + self._blockcmts["0"]
        del self._blockcmts["-1"]

    def OnRefresh(self):
        self.Clear()
        qty = self._mba.qty
        for src in range(qty):
            self.AddNode(src)
        for src in range(qty):
            mblock = self._mba.get_mblock(src)
            for dest in mblock.succset:
                self.AddEdge(src, dest)
        return True

    def OnGetText(self, node):
        mblock = self._mba.get_mblock(node)
        vp = hr.qstring_printer_t(None, True)
        mblock._print(vp)

        node_key = "%d" % node
        if node_key in self._blockcmts:
            return ''.join(self._blockcmts[node_key]) + vp.s
        else:
            return vp.s


class dominance_graphviewer_t(microcode_graphviewer_t):

    def __init__(self, *args):
        microcode_graphviewer_t.__init__(self, *args)
        self.dom_command_id = self.AddCommand("Show Dominance Graph", "D")
        self.full_command_id = self.AddCommand("Show Full Graph", "F")
        self.back_command_id = self.AddCommand("Show Previous Graph", "P")
        self.save_graphviz_id = self.AddCommand("save Graph to graphviz ", "S")
        self.show_dom_log_id = self.AddCommand("show current Dominance log ", "A")
        self.state = "cfg"
        self.back_stack = []
        self.select_block = None
        self.select_node = -1
        self.dom = {}

        self.compute_dominates()

    def OnCommand(self, cmd_id):
        if self.select_node != -1:
            node = self[self.select_node]
            if isinstance(node, hr.mblock_t):
                self.select_block = node
            elif isinstance(node, int):
                self.select_block = self._mba.get_mblock(node)
        if cmd_id == self.dom_command_id and self.select_node != -1:
            self.state = "dom"
            self.Refresh()
            self.Select(self.select_node)
            self.back_stack.append(self.select_block.serial)
        elif cmd_id == self.full_command_id:
            self.state = "cfg"
            last_select = self.select_block.serial if self.select_block else -1
            self.select_node = -1
            self.select_block = None
            self.Refresh()
            if last_select >= 0:
                self.Select(last_select)
        elif cmd_id == self.back_command_id and len(self.back_stack) > 0:
            self.back_stack.pop()
            if not self.back_stack:
                return self.OnCommand(self.full_command_id)
            else:
                serial = self.back_stack[-1]
                self.select_block = self._mba.get_mblock(serial)
                self.state = "dom"
                self.Refresh()
                self.Select(self.select_node)
        elif cmd_id == self.show_dom_log_id:
            dominanceFlow(self.select_block)
        elif cmd_id == self.save_graphviz_id:
            file_path = kw.ask_file(True, "*.graphviz", "Please select a file")
            if file_path:
                kw.msg("Selected file: {}\n".format(file_path))
                graphviz(self._mba,file_path)
            else:
                kw.msg("No file selected\n")
            print("save_graphviz_id")

    def OnClick(self, node_id):
        self.select_node = node_id

    def compute_dominates(self):
        nodes = set(list(range(self._mba.qty)))
        self.dom = {}
        for node in nodes:
            self.dom[node] = set(nodes)

        self.dom[0] = set([0])
        todo = set(nodes)

        while todo:
            node = todo.pop()

            if node == 0:
                continue

            new_dom = None
            mblock = self._mba.get_mblock(node)
            for pred in mblock.predset:
                if not pred in nodes:
                    continue

                if new_dom is None:
                    new_dom = set(self.dom[pred])
                new_dom &= self.dom[pred]
            if new_dom is None:
                new_dom = set([node])
            else:
                new_dom |= set([node])

            if new_dom == self.dom[node]:
                continue

            self.dom[node] = new_dom
            for succ in mblock.succset:
                todo.add(succ)

    def _get_dominates(self, blk):
        result = []
        for dom in self.dom:
            if blk.serial in self.dom[dom]:
                result.append(self._mba.get_mblock(dom))
        return result

    def OnRefresh(self):
        if self.state == "dom" and self.select_block:
            self.Clear()
            node_ids = {}
            dominates = self._get_dominates(self.select_block)
            for block in dominates:
                node_id = self.AddNode(block)
                if block.serial == self.select_block.serial:
                    self.select_node = node_id
                node_ids[block.serial] = node_id
            for block in dominates:
                for dest in block.succset:
                    if dest in node_ids:
                        self.AddEdge(node_ids[block.serial], node_ids[dest])
            return True
        return microcode_graphviewer_t.OnRefresh(self)

    def OnGetText(self, node_id):
        if isinstance(self[node_id], hr.mblock_t):
            node_id = self[node_id].serial
        return microcode_graphviewer_t.OnGetText(self, node_id)


# -----------------------------------------------------------------------------
# 显示microcode 文本内容的 ui界面类
class microcode_viewer_t(kw.simplecustviewer_t):
    """Creates a widget that displays Hex-Rays microcode."""

    def Create(self, mba, title, mmat_name, fn_name, lines=[]):
        self.title = "Microcode: %s" % title
        self._mba = mba
        self.mmat_name = mmat_name
        self.fn_name = fn_name
        if not kw.simplecustviewer_t.Create(self, self.title):
            return False
        self.lines = lines
        for line in lines:
            self.AddLine(line)
        return True


    def OnPopup(self, widget, popup_handle):
        """
        重写 OnPopup 方法，向右键菜单添加选项
        """
        # 将动作附加到右键菜单
        kw.attach_action_to_popup(widget, popup_handle, "my:action1", None)
        kw.attach_action_to_popup(widget, popup_handle, "my:action2", None)
        return super().OnPopup(widget, popup_handle)

    # 调整界面，内容剧中
    def _fit_graph(self, graph):
        if graph:
            gv = graph.GetWidget()
            ida_graph.viewer_fit_window(gv)
            ida_graph.refresh_viewer(gv)
            return True
        return False

    # 新界面靠右并排显示
    def _dock_widgets(self, graph, dockpos=kw.DP_RIGHT):
        if graph:
            gv = graph.GetWidget()
            kw.set_dock_pos(kw.get_widget_title(gv), self.title, dockpos)

            gli = ida_moves.graph_location_info_t()
            if ida_graph.viewer_get_gli(gli, gv):
                gli.zoom = 0  # auto-position
                ida_graph.viewer_set_gli(gv, gli, ida_graph.GLICTL_CENTER)
                ida_graph.refresh_viewer(gv)
            return True
        return False

    """TODO: it's better to handle keyboard input by
    registering an "action" and assigning it a hotkey"""

    def show_microcode_graph(self,shift):
        print("show_microcode_graph,",shift)
        g = show_microcode_graph(self._mba, self.title, self.lines)
        if g:
            self._fit_graph(g)
            self._dock_widgets(g, dockpos=kw.DP_FLOATING if shift else kw.DP_RIGHT)

    def OnKeydown(self, vkey, shift):
        if vkey == ord("G"):
            self.show_microcode_graph(shift)
            return True
        # elif vkey == ord("I"):
        #     """TODO: at some point, the textual representation of the mba
        #          should manually be created.
        #       -> we would no longer have to parse the textual output
        #          that is created by the gen_microcode() function
        #       .> we may insert COLOR_ADDR tags which would allow us to
        #          contextually link different viewers"""
        #     widget = self.GetWidget()
        #     line = kw.get_custom_viewer_curline(widget, False)
        #     line = ida_lines.tag_remove(line)
        #     p = line.find(" ")
        #     if p != -1 and '.' in line[:p]:
        #         block, serial = line.split('.')[:2]
        #         serial = serial.strip().split(' ')[0]
        #         g = microcode_insnviewer_t(self._mba, self.mmat_name, self.fn_name, int(block), int(serial))
        #         if g:
        #             g.Show()
        #             self._fit_graph(g)
        #             self._dock_widgets(g,
        #                                dockpos=kw.DP_FLOATING if shift else kw.DP_TAB)
        #     else:
        #         message = ("There is something wrong with the output generated by gen_microcode()!\n"
        #                    "Please rerun '%s.py'!" % PLUGIN_NAME)
        #         if line.startswith(";") or not (len(line)):
        #             message = "Please position the cursor on a microcode instruction."
        #         kw.warning(message)
        #     return True
        # return False


# -----------------------------------------------------------------------------
def ask_desired_maturity():
    """Displays a dialog which lets the user choose a maturity level
    of the microcode to generate."""
    maturity_levels = [
        ["MMAT_GENERATED", hr.MMAT_GENERATED],
        ["MMAT_PREOPTIMIZED", hr.MMAT_PREOPTIMIZED],
        ["MMAT_LOCOPT", hr.MMAT_LOCOPT],
        ["MMAT_CALLS", hr.MMAT_CALLS],
        ["MMAT_GLBOPT1", hr.MMAT_GLBOPT1],
        ["MMAT_GLBOPT2", hr.MMAT_GLBOPT2],
        ["MMAT_GLBOPT3", hr.MMAT_GLBOPT3],
        ["MMAT_LVARS", hr.MMAT_LVARS]]

    class MaturityForm(kw.Form):
        def __init__(self):
            self.title = "Display Microcode"
            form = ("STARTITEM {id:mat_lvl}\n"
                    "%s\n"
                    " \n"
                    "<Maturity level:{mat_lvl}>\n\n"
                    "<##Options##Output includes comments:{flags_short}>{chkgroup_flags}>\n\n" %
                    self.title)

            dropdown_ctl = kw.Form.DropdownListControl(
                [text for text, _ in maturity_levels])
            chk_ctl = kw.Form.ChkGroupControl(("flags_short",))

            controls = {"mat_lvl": dropdown_ctl,
                        "chkgroup_flags": chk_ctl}

            kw.Form.__init__(self, form, controls)

    form = MaturityForm()
    form, args = form.Compile()
    form.flags_short.checked = True
    ok = form.Execute()

    mmat = None
    text = None
    flags = 0
    if ok == 1:
        text, mmat = maturity_levels[form.mat_lvl.value]
    flags |= 0 if form.flags_short.checked else hr.MBA_SHORT
    form.Free()
    return (text, mmat, flags)


# -----------------------------------------------------------------------------
def show_microcode():
    # import pydevd_pycharm
    # pydevd_pycharm.settrace('localhost', port=31235, stdoutToServer=True, stderrToServer=True)

    """Generates and displays microcode for an address range.
    An address range can be a selection of code or that of
    the current function."""
    sel, sea, eea = kw.read_range_selection(None)
    pfn = ida_funcs.get_func(kw.get_screen_ea())
    if not sel and not pfn:
        return (False, "Position cursor within a function or select range")

    if not sel and pfn:
        sea = pfn.start_ea
        eea = pfn.end_ea

    addr_fmt = "%016x" if ida_ida.inf_is_64bit() else "%08x"
    fn_name = (ida_funcs.get_func_name(pfn.start_ea)
               if pfn else "0x%s-0x%s" % (addr_fmt % sea, addr_fmt % eea))
    F = ida_bytes.get_flags(sea)
    if not ida_bytes.is_code(F):
        return (False, "The selected range must start with an instruction")

    text, mmat, mba_flags = ask_desired_maturity()  # dialog
    if text is None and mmat is None:
        return (True, "Cancelled")

    if not sel and pfn:
        mbr = hr.mba_ranges_t(pfn)
    else:
        mbr = hr.mba_ranges_t()
        mbr.ranges.push_back(ida_range.range_t(sea, eea))

    hf = hr.hexrays_failure_t()
    ml = hr.mlist_t()
    mba = hr.gen_microcode(mbr, hf, ml, hr.DECOMP_WARNINGS, mmat)
    if not mba:
        return (False, "0x%s: %s" % (addr_fmt % hf.errea, hf.desc()))
    vp = printer_t()
    mba.set_mba_flags(mba_flags)
    mba._print(vp)     # printer_t 这个是ida系统提供的获取所有microcode 文本内容的类，
    global microcode_viewer_instance
    microcode_viewer_instance = microcode_viewer_t()
    if not microcode_viewer_instance.Create(mba, "0x%s-0x%s (%s)" % (addr_fmt % sea, addr_fmt % eea, text), text, fn_name, vp.get_mc()):
        return (False, "Error creating viewer")

    microcode_viewer_instance.Show()
    return (True,
            "Successfully generated microcode for 0x%s..0x%s" % (addr_fmt % sea, addr_fmt % eea))

def show_microcode_graph(mba,fn_name,lines=None):
    if lines == None:
        vp = printer_t()
        mba.set_mba_flags(mba.get_mba_flags())
        mba._print(vp)
        g = dominance_graphviewer_t(mba, fn_name, vp.get_mc())
        if g:
            g.Show()
        return g
    else:
        g = dominance_graphviewer_t(mba, fn_name, lines)
        if g:
            g.Show()
        return g

#-------------------------------------------------------------------------------
class ActionHandler(kw.action_handler_t):
    def __init__(self, callback):
        super().__init__()
        self.callback = callback

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

    def update(self, ctx):
        return 1  # Always enabled

#------------------------------------------------------------------------------
# 添加菜单
def register_actions():
    global microcode_viewer_instance

    action1 = kw.action_desc_t(
        "my:action1",  # 唯一标识符
        "show Code FlowChart   G！",    # 菜单显示名称
        ActionHandler(lambda: microcode_viewer_instance.show_microcode_graph(0)),
        None,          # 热键
        "show Code FlowChart   G",  # 提示信息
        -1             # 图标
    )
    action2 = kw.action_desc_t(
        "my:action2",
        "Option 2",
        ActionHandler(lambda: print("Option 2 selected")),
        None,
        "Select Option 2",
        -1
    )
    kw.register_action(action1)
    kw.register_action(action2)



# -----------------------------------------------------------------------------
def create_mc_widget():
    """Checks minimum requirements for the script/plugin to be able to run.
    Displays microcode or in case of failure, displays error message.
    This function acts as the main entry point that is invoked if the
    code is run as a script or as a plugin."""
    register_actions()
    success, message = show_microcode()
    output = kw.msg if success else kw.warning
    output("%s: %s\n" % (PLUGIN_NAME, message))
    return success


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

    def init(self):
        return ida_idaapi.PLUGIN_OK

    def run(self, arg):
        create_mc_widget()

    def term(self):
        pass


# -----------------------------------------------------------------------------
def PLUGIN_ENTRY():
    """Entry point of this code if launched as a plugin."""
    return genmc()


def printf(str):
    print(str)


# -----------------------------------------------------------------------------
def SCRIPT_ENTRY():
    """Entry point of this code if launched as a script."""
    printf("SCRIPT_ENTRY")
    if not is_plugin():
        if not is_installed():
            kw.msg(("%s: Available commands:\n"
                    "[+] \"install_plugin()\" - install script to ida_userdir/plugins\n") % (
                       PLUGIN_NAME))
        create_mc_widget()
        return True

    return False


# -----------------------------------------------------------------------------
SCRIPT_ENTRY()


repository: lucasg/idamagnum
file: idamagnum_plugin.py
description: idamagnum is a plugin for integrating MagnumDB requests within IDA
wanted_name: IdaMagnum
comment: search magic numbers using magnumdb.com
created_at: 2018-12-19 21:06:48+00:00
pushed_at: 2020-04-10 08:14:48+00:00
url:
https://github.com/lucasg/idamagnum/blob/1ed487c67ae563c1f2f752ee8d5935e0efea5133/idamagnum_plugin.py

content:
import json

try:
    # Python2
    from urllib2 import urlopen
except ImportError:
    # Python3
    from urllib.request import urlopen

import idc
import idaapi
import idautils
from idaapi import plugin_t


try:
    # ida < 7.4
    from idaapi import Choose2 as Choose
    from idc import OpEnumEx as op_enum
except ImportError:
    # ida 7.4
    from ida_kernwin import Choose
    from idc import op_enum



class ChooseMagicNumber(Choose):
    
    def __init__(self,value,results):

        Choose.__init__(self,
            "[IdaMagnum] Select enum from MagnumDB.com for value : 0x%X" % value,
            [ ["name",   13 | Choose.CHCOL_PLAIN], 
              ["value",  10 | Choose.CHCOL_HEX],
              ["source",  13 | Choose.CHCOL_PLAIN],
            ],
            Choose.CH_MODAL
        )

        self._results = results

    def OnSelectLine(self, n):
        pass

    def OnGetLine(self, n):
        res = self._results[n]
        return [
            res["Title"], 
            res.get("HexValue", ""),
            res["DisplayFilePath"]
        ]  

    def OnRefresh(self, n):
        return n

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

class SearchMagicNumber(idaapi.action_handler_t):

    MAGNUMDB_QUERY = "https://www.magnumdb.com/api.aspx?q=0x{value:X}&key={key:s}"
    MAGNUMDB_KEY = "f344dc86-7796-499f-be38-ec39a5414289"

    def __init__(self, manager):
        idaapi.action_handler_t.__init__(self)
        self._manager = manager

    def shift_bit_length(self, x):
        # https://stackoverflow.com/questions/14267555/find-the-smallest-power-of-2-greater-than-n-in-python
        return 1<<(x-1).bit_length()

    def activate(self, ctx):
        # print("ctx.cur_ea : 0x%x" % ctx.cur_ea)
        # print("ctx.cur_extracted_ea : 0x%x" % ctx.cur_extracted_ea)

        # Extract selected enum
        instruction = idc.GetDisasm(ctx.cur_ea)
        selection = ctx.cur_extracted_ea

        # correctly parse the selected value as int (hex or decimal)
        # since ctx.cur_extracted_ea has a bug (IDA always consider
        # the selected value as hex)
        if instruction.find("{0:X}h".format(selection)) != -1:
            # print("hex value found !")
            selected_value = ctx.cur_extracted_ea
        elif instruction.find("{0:d}".format(selection)) != -1:
            # print("int value found !")
            selected_value = int("%x" % ctx.cur_extracted_ea)
        else:
            # print("nothing selected !")
            return 1

        # next power of two for masking
        selected_value_mask = self.shift_bit_length(selected_value) - 1
        # print("selected_value : 0x%X" % selected_value)
        # print("selected_value mask : 0x%X" % selected_value_mask)

        # query magnum db
        url = SearchMagicNumber.MAGNUMDB_QUERY.format(
            value = selected_value,
            key = SearchMagicNumber.MAGNUMDB_KEY
        )
        answer = urlopen(url)
        results = json.loads(answer.read())
        
        # Let the user select the best answer
        c = ChooseMagicNumber(selected_value, results["Items"])
        selected_index = c.Show(modal=True)
        if selected_index < 0:
            return

        # Apply the newly found enum
        selected_item = results["Items"][selected_index]
        selected_name = selected_item["Title"].encode('ascii')
        selected_value = int(selected_item["Value"])

        # serial is important since several entries can have the same value
        entryid, serial = self._manager.add_magnumdb_entry(
            selected_name, 
            selected_value
        )

        # locate the operand where to apply the enum
        insn = idautils.DecodeInstruction(ctx.cur_ea)
        
        try:
            # ida < 7.4
            operands = insn.Operands
        except AttributeError:
            # ida 7.4
            operands = insn.ops

        for op in filter(lambda o: o.type == idaapi.o_imm,operands):

            # heuristic : the selected immediate is the first in the instruction with 
            # the same exact value (we are using a mask since IDA loves to set FFFFFFFF to high words)
            if op.value & selected_value_mask == selected_value:
                # Apply the enum
                op_enum(ctx.cur_ea, op.n, idaapi.get_enum("_IDA_MAGNUMDB"), serial)
                break

        return 1

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

class ConfigureIdaMagnum(idaapi.action_handler_t):
    def __init__(self, manager):
        idaapi.action_handler_t.__init__(self)
        self._manager = manager

    def activate(self, ctx):
         #self._manager.proc_rop()
        return 1

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

class IdaMagnumManager(object):

    def __init__(self):
        self._attach_to_menu_items()

    def _attach_to_menu_items(self):

        self.search_magic_desc = idaapi.action_desc_t(
            'idamagnum:searchmagic',             
            'search magic number ...',              
            SearchMagicNumber(self),         
            "Shift+M",                     
            'Search this value on MagnumDB', 
        )

        self.configure_plugin_desc = idaapi.action_desc_t(
            'idamagnum:configure',             
            'Configure',              
            ConfigureIdaMagnum(self),         
            "",                     
            'Configure plugin',
        )

        idaapi.register_action(self.search_magic_desc)
        idaapi.register_action(self.configure_plugin_desc)

        idaapi.attach_action_to_menu(
            'Edit/Plugins/IdaMagnum/',
            'idamagnum:searchmagic',
            idaapi.SETMENU_APP
        )

        idaapi.attach_action_to_menu(
            'Edit/Plugins/IdaMagnum/',
            'idamagnum:configure',
            idaapi.SETMENU_APP
        )

        return 0

    def _detach_from_menu_items(self):
        idaapi.detach_action_from_menu('Edit/Plugins/IdaMagnum/', 'idamagnum:searchmagic')
        idaapi.detach_action_from_menu('Edit/Plugins/IdaMagnum/', 'idamagnum:configure')

    def ensure_magnumdb_enum_type(self):
        """Ensure we have a valid MAGNUMDB enum"""

        enum_id = idaapi.get_enum("_IDA_MAGNUMDB")
        if enum_id == 0xffffffffffffffff:
            enum_id = idaapi.add_enum(idaapi.BADADDR, "_IDA_MAGNUMDB", 0)

        return enum_id

    def add_magnumdb_entry(self, name, value):
         
        enum_id = self.ensure_magnumdb_enum_type()

        # idaapi.add_enum_member accept only str (Py3)
        if type(name) == type(b''):
            name = name.decode('utf-8')

        serial = 0
        enum_memberid = idaapi.get_enum_member(enum_id, value, serial, 0)
        while  enum_memberid != 0xffffffffffffffff:

            if idaapi.get_enum_member_name(enum_memberid) == name:
                return enum_memberid, serial

            serial += 1
            enum_memberid = idaapi.get_enum_member(enum_id, value, serial, 0)

        if enum_memberid == 0xffffffffffffffff:
            enum_memberid = idaapi.add_enum_member(enum_id, name, value)

        return enum_memberid, serial
        

class IdaMagnumPlugin(plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = "search magic numbers using magnumdb.com"
    help = "search magic numbers using magnumdb.com"
    wanted_name = "IdaMagnum"
    wanted_hotkey = ""


    def init(self):
        global ida_magnumdb_manager

        if not 'ida_magnumdb_manager' in globals():
            ida_magnumdb_manager = IdaMagnumManager()

            print("[IdaMagnum] Ida plugin for MagnumDB v0.0 initialized")

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        pass

def PLUGIN_ENTRY():
    return IdaMagnumPlugin()


repository: Jinmo/palette-plugins
file: community-headers/plugins/main.py
description: Palette plugins
wanted_name: Community headers
comment: This is a comment
created_at: 2019-03-03 07:45:35+00:00
pushed_at: 2019-04-28 13:39:53+00:00
url:
https://github.com/Jinmo/palette-plugins/blob/7b00707fedbe1031aa9b420ffbe4206db7883ebd/community-headers/plugins/main.py

content:
from __palette__ import Palette, show_palette, Action

import idaapi
import requests
import json
import threading

from functools import partial
from idc import ParseTypes

sess = requests.Session()
sess.trust_env = False

REPO_URL = 'https://api.github.com/repos/Jinmo/headers/git/trees/master'
REPO_RAW_URL = 'https://raw.githubusercontent.com/Jinmo/headers/master/'


def import_headers(name, item):
    r = sess.get(REPO_RAW_URL + name)
    ParseTypes(str(r.content))
    print 'loading', name, 'done'
    return True


def show_online_headers():
    def handler():
        headers = json.loads(sess.get(REPO_URL).content).get('tree')
        headers = map(lambda x: x['path'], headers)
        actions = list(Action(handler=partial(import_headers, x), id=x, description= x)
                       for x in headers)
        show_palette(Palette(myplugin_t.wanted_name, "Enter header path...", actions))
    t = threading.Thread(target=handler)
    t.start()


# Helper for registering actions
def register_action(name, shortcut=''):
    def handler(f):
        # 1) Create the handler class
        class MyHandler(idaapi.action_handler_t):
            def __init__(self):
                idaapi.action_handler_t.__init__(self)

            # Say hello when invoked.
            def activate(self, ctx):
                t = threading.Thread(target=f)
                t.start()
                return 1

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

        # 2) Describe the action
        action_desc = idaapi.action_desc_t(
            name,  # The action name. This acts like an ID and must be unique
            name,  # The action text.
            MyHandler(),  # The action handler.
            shortcut,  # Optional: the action shortcut
            name,  # Optional: the action tooltip (available in menus/toolbar)
            0)  # Optional: the action icon (shows when in menus/toolbars)

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

    return handler


class myplugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE | idaapi.PLUGIN_FIX
    comment = "This is a comment"
    help = "Finds headers"
    wanted_name = "Community headers"
    wanted_hotkey = ""

    def init(self):
        register_action('Community Headers: Find headers online...')(show_online_headers)
        return idaapi.PLUGIN_OK

    def run(self, arg):
        pass

    def term(self):
        pass


def PLUGIN_ENTRY():
    return myplugin_t()


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

content:
##########################################################################
# An IDAPython plugin that generates "fuzzy" function signatures that can
# be shared and applied amongst different IDBs.
#
# There are multiple sets of signatures that are generated:
#
#   o "Formal" signatures, where functions must match exactly
#   o "Fuzzy" signatures, where functions must only resemble each other
#     in terms of data/call references.
#   o String-based signatures, where functions are identified based on
#     unique string references.
#   o Immediate-based signatures, where functions are identified based
#     on immediate value references.
#
# These signatures are applied based on accuracy, that is, formal
# signatures are applied first, then string and immediate based
# signatures, and finally fuzzy signatures.
#
# Further, functions are identified based on call references. Consider,
# for example, two functions, one named 'foo', the other named 'bar'.
# The 'foo' function is fairly unique and a reliable signature is easily
# generated for it, but the 'bar' function is more difficult to reliably
# identify. However, 'foo' calls 'bar', and thus once 'foo' is identified,
# 'bar' can also be identified by association.
#
# Craig Heffner
# @devttys0
##########################################################################
import idc
import idaapi
import idautils

import os
import sys
import time
import pickle       # http://natashenka.ca/pickle/
import collections

class RizzoSignatures(object):
    '''
    Simple wrapper class for storing signature info.
    '''

    SHOW = []

    def __init__(self):
        self.fuzzy = {}
        self.formal = {}
        self.strings = {}
        self.functions = {}
        self.immediates = {}

        self.fuzzydups = set()
        self.formaldups = set()
        self.stringdups = set()
        self.immediatedups = set()

    def show(self):
        if not self.SHOW:
            return

        print "\n\nGENERATED FORMAL SIGNATURES FOR:"
        for (key, ea) in self.formal.iteritems():
            func = RizzoFunctionDescriptor(self.formal, self.functions, key)
            if func.name in self.SHOW:
                print func.name

        print "\n\nGENERATRED FUZZY SIGNATURES FOR:"
        for (key, ea) in self.fuzzy.iteritems():
            func = RizzoFunctionDescriptor(self.fuzzy, self.functions, key)
            if func.name in self.SHOW:
                print func.name

class RizzoStringDescriptor(object):
    '''
    Wrapper class for easily accessing necessary string information.
    '''

    def __init__(self, string):
        self.ea = string.ea
        self.value = str(string)
        self.xrefs = [x.frm for x in idautils.XrefsTo(self.ea)]

class RizzoBlockDescriptor(object):
    '''
    Code block info is stored in tuples, which minimize pickle storage space.
    This class provides more Pythonic (and sane) access to values of interest for a given block.
    '''

    def __init__(self, block):
        self.formal = block[0]
        self.fuzzy = block[1]
        self.immediates = block[2]
        self.functions = block[3]

    def match(self, nblock, fuzzy=False):
        # TODO: Fuzzy matching at the block level gets close, but produces a higher number of
        #       false positives; for example, it confuses hmac_md5 with hmac_sha1.
        #return ((self.formal == nblock.formal or (fuzzy and self.fuzzy == nblock.fuzzy)) and
        return (self.formal == nblock.formal and
                len(self.immediates) == len(nblock.immediates) and
                len(self.functions) == len(nblock.functions))

class RizzoFunctionDescriptor(object):
    '''
    Function signature info is stored in dicts and tuples, which minimize pickle storage space.
    This class provides more Pythonic (and sane) access to values of interest for a given function.
    '''

    def __init__(self, signatures, functions, key):
        self.ea = signatures[key]
        self.name = functions[self.ea][0]
        self.blocks = functions[self.ea][1]

class Rizzo(object):
    '''
    Workhorse class which performs the primary logic and functionality.
    '''

    DEFAULT_SIGNATURE_FILE = "rizzo.sig"

    def __init__(self, sigfile=None):
        if sigfile:
            self.sigfile = sigfile
        else:
            self.sigfile = self.DEFAULT_SIGNATURE_FILE

        # Useful for quickly identifying string xrefs from individual instructions
        self.strings = {}
        for string in idautils.Strings():
            self.strings[string.ea] = RizzoStringDescriptor(string)

        start = time.time()
        self.signatures = self.generate()
        end = time.time()

        print "Generated %d formal signatures and %d fuzzy signatures for %d functions in %.2f seconds." % (len(self.signatures.formal), len(self.signatures.fuzzy), len(self.signatures.functions), (end-start))

    def save(self):
        print ("Saving signatures to %s..." % self.sigfile),
        fp = open(self.sigfile, "wb")
        pickle.dump(self.signatures, fp)
        fp.close()
        print "done."

    def load(self):
        print ("Loading signatures from %s..." % self.sigfile),
        fp = open(self.sigfile, "rb")
        sigs = pickle.load(fp)
        fp.close()
        print "done."
        return sigs

    def sighash(self, value):
        return hash(str(value)) & 0xFFFFFFFF

    def block(self, block):
        '''
        Returns a tuple: ([formal, block, signatures], [fuzzy, block, signatures], set([unique, immediate, values]), [called, function, names])
        '''
        formal = []
        fuzzy = []
        functions = []
        immediates = []

        ea = block.startEA
        while ea < block.endEA:
            idaapi.decode_insn(ea)

            # Get a list of all data/code references from the current instruction
            drefs = [x for x in idautils.DataRefsFrom(ea)]
            crefs = [x for x in idautils.CodeRefsFrom(ea, False)]

            # Add all instruction mnemonics to the formal block hash
            formal.append(idc.GetMnem(ea))

            # If this is a call instruction, be sure to note the name of the function
            # being called. This is used to apply call-based signatures to functions.
            #
            # For fuzzy signatures, we can't use the actual name or EA of the function,
            # but rather just want to note that a function call was made.
            #
            # Formal signatures already have the call instruction mnemonic, which is more
            # specific than just saying that a call was made.
            if idaapi.is_call_insn(ea):
                for cref in crefs:
                    func_name = idc.Name(cref)
                    if func_name:
                        functions.append(func_name)
                        fuzzy.append("funcref")
            # If there are data references from the instruction, check to see if any of them
            # are strings. These are looked up in the pre-generated strings dictionary.
            #
            # String values are easily identifiable, and are used as part of both the fuzzy
            # and the formal signatures.
            #
            # It is more difficult to determine if non-string values are constants or not;
            # for both fuzzy and formal signatures, just use "data" to indicate that some data
            # was referenced.
            elif drefs:
                for dref in drefs:
                    if self.strings.has_key(dref):
                        formal.append(self.strings[dref].value)
                        fuzzy.append(self.strings[dref].value)
                    else:
                        formal.append("dataref")
                        fuzzy.append("dataref")
            # If there are no data or code references from the instruction, use every operand as
            # part of the formal signature.
            #
            # Fuzzy signatures are only concerned with interesting immediate values, that is, values
            # that are greater than 65,535, are not memory addresses, and are not displayed as
            # negative values.
            elif not drefs and not crefs:
                for n in range(0, len(idaapi.cmd.Operands)):
                    opnd_text = idc.GetOpnd(ea, n)
                    formal.append(opnd_text)
                    if idaapi.cmd.Operands[n].type == idaapi.o_imm and not opnd_text.startswith('-'):
                        if idaapi.cmd.Operands[n].value >= 0xFFFF:
                            if idaapi.getFlags(idaapi.cmd.Operands[n].value) == 0:
                                fuzzy.append(str(idaapi.cmd.Operands[n].value))
                                immediates.append(idaapi.cmd.Operands[n].value)

            ea = idc.NextHead(ea)

        return (self.sighash(''.join(formal)), self.sighash(''.join(fuzzy)), immediates, functions)

    def function(self, func):
        '''
        Returns a list of blocks.
        '''
        blocks = []

        for block in idaapi.FlowChart(func):
            blocks.append(self.block(block))

        return blocks

    def generate(self):
        signatures = RizzoSignatures()

        # Generate unique string-based function signatures
        for (ea, string) in self.strings.iteritems():
            # Only generate signatures on reasonably long strings with one xref
            if len(string.value) >= 8 and len(string.xrefs) == 1:
                func = idaapi.get_func(string.xrefs[0])
                if func:
                    strhash = self.sighash(string.value)

                    # Check for and remove string duplicate signatures (the same
                    # string can appear more than once in an IDB).
                    # If no duplicates, add this to the string signature dict.
                    if signatures.strings.has_key(strhash):
                        del signatures.strings[strhash]
                        signatures.stringdups.add(strhash)
                    elif strhash not in signatures.stringdups:
                        signatures.strings[strhash] = func.startEA

        # Generate formal, fuzzy, and immediate-based function signatures
        for ea in idautils.Functions():
            func = idaapi.get_func(ea)
            if func:
                # Generate a signature for each block in this function
                blocks = self.function(func)

                # Build function-wide formal and fuzzy signatures by simply
                # concatenating the individual function block signatures.
                formal = self.sighash(''.join([str(e) for (e, f, i, c) in blocks]))
                fuzzy = self.sighash(''.join([str(f) for (e, f, i, c) in blocks]))

                # Add this signature to the function dictionary.
                signatures.functions[func.startEA] = (idc.Name(func.startEA), blocks)

                # Check for and remove formal duplicate signatures.
                # If no duplicates, add this to the formal signature dict.
                if signatures.formal.has_key(formal):
                    del signatures.formal[formal]
                    signatures.formaldups.add(formal)
                elif formal not in signatures.formaldups:
                    signatures.formal[formal] = func.startEA

                # Check for and remove fuzzy duplicate signatures.
                # If no duplicates, add this to the fuzzy signature dict.
                if signatures.fuzzy.has_key(fuzzy):
                    del signatures.fuzzy[fuzzy]
                    signatures.fuzzydups.add(fuzzy)
                elif fuzzy not in signatures.fuzzydups:
                    signatures.fuzzy[fuzzy] = func.startEA

                # Check for and remove immediate duplicate signatures.
                # If no duplicates, add this to the immediate signature dict.
                for (e, f, immediates, c) in blocks:
                    for immediate in immediates:
                        if signatures.immediates.has_key(immediate):
                            del signatures.immediates[immediate]
                            signatures.immediatedups.add(immediate)
                        elif immediate not in signatures.immediatedups:
                            signatures.immediates[immediate] = func.startEA

        # These need not be maintained across function calls,
        # and only add to the size of the saved signature file.
        signatures.fuzzydups = set()
        signatures.formaldups = set()
        signatures.stringdups = set()
        signatures.immediatedups = set()

        # DEBUG
        signatures.show()

        return signatures

    def match(self, extsigs):
        fuzzy = {}
        formal = {}
        strings = {}
        immediates = {}

        # Match formal function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.formal.iteritems():
            if self.signatures.formal.has_key(extsig):
                newfunc = RizzoFunctionDescriptor(extsigs.formal, extsigs.functions, extsig)
                curfunc = RizzoFunctionDescriptor(self.signatures.formal, self.signatures.functions, extsig)
                formal[curfunc] = newfunc
        end = time.time()
        print "Found %d formal matches in %.2f seconds." % (len(formal), (end-start))

        # Match fuzzy function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.fuzzy.iteritems():
            if self.signatures.fuzzy.has_key(extsig):
                curfunc = RizzoFunctionDescriptor(self.signatures.fuzzy, self.signatures.functions, extsig)
                newfunc = RizzoFunctionDescriptor(extsigs.fuzzy, extsigs.functions, extsig)
                # Only accept this as a valid match if the functions have the same number of basic code blocks
                if len(curfunc.blocks) == len(newfunc.blocks):
                    fuzzy[curfunc] = newfunc
        end = time.time()
        print "Found %d fuzzy matches in %.2f seconds." % (len(fuzzy), (end-start))

        # Match string based function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.strings.iteritems():
            if self.signatures.strings.has_key(extsig):
                curfunc = RizzoFunctionDescriptor(self.signatures.strings, self.signatures.functions, extsig)
                newfunc = RizzoFunctionDescriptor(extsigs.strings, extsigs.functions, extsig)
                strings[curfunc] = newfunc
        end = time.time()
        print "Found %d string matches in %.2f seconds." % (len(strings), (end-start))

        # Match immediate baesd function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.immediates.iteritems():
            if self.signatures.immediates.has_key(extsig):
                curfunc = RizzoFunctionDescriptor(self.signatures.immediates, self.signatures.functions, extsig)
                newfunc = RizzoFunctionDescriptor(extsigs.immediates, extsigs.functions, extsig)
                immediates[curfunc] = newfunc
        end = time.time()
        print "Found %d immediate matches in %.2f seconds." % (len(immediates), (end-start))

        # Return signature matches in the order we want them applied
        # The second tuple of each match is set to True if it is a fuzzy match, e.g.:
        #
        #   ((match, fuzzy), (match, fuzzy), ...)
        return ((formal, False), (strings, False), (immediates, False), (fuzzy, True))

    def rename(self, ea, name):
        # Don't rely on the name in curfunc, as it could have already been renamed
        curname = idc.Name(ea)
        # Don't rename if the name is a special identifier, or if the ea has already been named
        # TODO: What's a better way to check for reserved name prefixes?
        if curname.startswith('sub_') and name.split('_')[0] not in set(['sub', 'loc', 'unk', 'dword', 'word', 'byte']):
            # Don't rename if the name already exists in the IDB
            if idc.LocByName(name) == idc.BADADDR:
                if idc.MakeName(ea, name):
                    idc.SetFunctionFlags(ea, (idc.GetFunctionFlags(ea) | idc.FUNC_LIB))
                    #print "%s  =>  %s" % (curname, name)
                    return 1
            #else:
            #    print "WARNING: Attempted to rename '%s' => '%s', but '%s' already exists!" % (curname, name, name)
        return 0

    def apply(self, extsigs):
        count = 0

        start = time.time()

        # This applies formal matches first, then fuzzy matches
        for (match, fuzzy) in self.match(extsigs):
            # Keeps track of all function names that we've identified candidate functions for
            rename = {}

            for (curfunc, newfunc) in match.iteritems():
                if not rename.has_key(newfunc.name):
                    rename[newfunc.name] = []

                # Attempt to rename this function
                rename[newfunc.name].append(curfunc.ea)

                bm = {}
                duplicates = set()

                # Search for unique matching code blocks inside this function
                for nblock in newfunc.blocks:
                    nblock = RizzoBlockDescriptor(nblock)
                    for cblock in curfunc.blocks:
                        cblock = RizzoBlockDescriptor(cblock)

                        if cblock.match(nblock, fuzzy):
                            if bm.has_key(cblock):
                                del bm[cblock]
                                duplicates.add(cblock)
                            elif cblock not in duplicates:
                                bm[cblock] = nblock

                # Rename known function calls from each unique identified code block
                for (cblock, nblock) in bm.iteritems():
                    for n in range(0, len(cblock.functions)):
                        ea = idc.LocByName(cblock.functions[n])
                        if ea != idc.BADADDR:
                            if rename.has_key(nblock.functions[n]):
                                rename[nblock.functions[n]].append(ea)
                            else:
                                rename[nblock.functions[n]] = [ea]

                # Rename the identified functions
                for (name, candidates) in rename.iteritems():
                    if candidates:
                        winner = collections.Counter(candidates).most_common(1)[0][0]
                        count += self.rename(winner, name)

        end = time.time()
        print "Renamed %d functions in %.2f seconds." % (count, (end-start))

def RizzoBuild(sigfile=None):
    print "Building Rizzo signatures, this may take a few minutes..."
    start = time.time()
    r = Rizzo(sigfile)
    r.save()
    end = time.time()
    print "Built signatures in %.2f seconds" % (end-start)

def RizzoApply(sigfile=None):
    print "Applying Rizzo signatures, this may take a few minutes..."
    start = time.time()
    r = Rizzo(sigfile)
    s = r.load()
    r.apply(s)
    end = time.time()
    print "Signatures applied in %.2f seconds" % (end-start)

def rizzo_produce(self, arg):
    fname = idc.AskFile(1, "*.riz", "Save signature file as")
    if fname:
        if '.' not in fname:
            fname += ".riz"
        RizzoBuild(fname)
    return None

def rizzo_load(self, arg):
    fname = idc.AskFile(0, "*.riz", "Load signature file")
    if fname:
        RizzoApply(fname)
    return None


class RizzoPlugin(idaapi.plugin_t):
    flags = 0
    comment = "Function signature"
    help = ""
    wanted_name = "Rizzo"
    wanted_hotkey = ""

    NAME = "rizzo.py"

    # Actions
    action_rizzo_loadfile = 'rizzo:loadfile'
    action_rizzo_producefile = 'rizzo:producefile'

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.register_action(idaapi.action_desc_t(
                self.action_rizzo_loadfile,
                'Rizzo signature file...',
                self.rizzo_loadHandler(),
                None,
                'Rizzo signature file...'))
            idaapi.attach_action_to_menu("File/Load file/",
                self.action_rizzo_loadfile, idaapi.SETMENU_APP)

            idaapi.register_action(idaapi.action_desc_t(
                self.action_rizzo_producefile,
                'Rizzo signature file...',
                self.rizzo_produceHandler(),
                None,
                'Rizzo signature file...'))
            idaapi.attach_action_to_menu("File/Produce file/",
                self.action_rizzo_producefile, idaapi.SETMENU_APP)
        else:
            self.menu_context_load = idaapi.add_menu_item("File/Load file/", "Rizzo signature file...", "", 0, self.rizzo_load, (None,))
            self.menu_context_produce = idaapi.add_menu_item("File/Produce file/", "Rizzo signature file...", "", 0, self.rizzo_produce, (True,))
        return idaapi.PLUGIN_KEEP

    def term(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.unregister_action(self.action_rizzo_loadfile)
            idaapi.unregister_action(self.action_rizzo_producefile)
        else:
            idaapi.del_menu_item(self.menu_context_load)
            idaapi.del_menu_item(self.menu_context_produce)
        return None

    def run(self, arg):
        return None

    def rizzo_script(self):
        idaapi.IDAPython_ExecScript(self.script, globals())

    class rizzo_loadHandler(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)
            
        def activate(self, ctx):
            rizzo_load(self, 0)

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

    class rizzo_produceHandler(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)
            
        def activate(self, ctx):
            rizzo_produce(self, 0)

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

def PLUGIN_ENTRY():
    return RizzoPlugin()


repository: doyaGu/GooFramework
file: scripts/ClassDumper.py
description: The WIP SDK for World of Goo 2.
wanted_name: ClassDumper
comment: Dump class infomations
created_at: 2024-08-10 16:12:12+00:00
pushed_at: 2024-08-10 16:34:26+00:00
url:
https://github.com/doyaGu/GooFramework/blob/fa5f72a9ca4888ec79deec3f64d352780c5a2f2c/scripts/ClassDumper.py

content:
from typing import List, Tuple, Optional

import idaapi
import idc
import ida_bytes
import ida_name
import ida_funcs
import ida_segment
import idautils
import pyperclip

# Define masks for flags in __vmi_class_type_info
NON_DIAMOND_REPEAT_MASK = 1
DIAMOND_SHAPED_MASK = 2

# Define masks for offset flags in __base_class_type_info
VIRTUAL_MASK = 1
PUBLIC_MASK = 2
OFFSET_SHIFT = 8

CLASS_TYPE_INFO = 0
SI_CLASS_TYPE_INFO = 1
VMI_CLASS_TYPE_INFO = 2

ACTION_DUMP = "classdumper:dump"
ACTION_COPY_VTABLE = "classdumper:copyvtable"
ACTION_PASTE_VTABLE = "classdumper:pastevtable"


def output_lines(lines: List[str], filename: str) -> None:
    """
    Output the lines to a specified file.

    Args:
        lines (List[str]): The lines to write to the file.
        filename (str): The name of the file to write to.
    """
    if filename.find("::") != -1:
        filename = filename.replace("::", "_")

    with open(filename, "w", encoding="utf-8") as f:
        for line in lines:
            f.write(line + "\n")


def get_pointer_size() -> int:
    """
    Get the size of a pointer on the current platform (4 bytes for 32-bit, 8 bytes for 64-bit).

    Returns:
        int: The size of a pointer in bytes.
    """
    return 8 if idaapi.get_inf_structure().is_64bit() else 4


def get_strlit(ea: int) -> Optional[bytes]:
    """
    Retrieve the string literal contents at the specified address.

    Args:
        ea (int): The address from which to retrieve the string literal.

    Returns:
        Optional[bytes]: The string literal contents at the given address, or None if not found.
    """
    return ida_bytes.get_strlit_contents(ea, -1, idc.STRTYPE_C)


def get_string(ea: int) -> str:
    """
    Retrieve a null-terminated string from the specified address.

    Args:
        ea (int): The address to read the string from.

    Returns:
        str: The retrieved string.
    """
    strlit = get_strlit(ea)
    return strlit.decode("utf-8") if strlit is not None else ""


def get_function_name(ea: int) -> str:
    """
    Retrieve the function name at the given address.

    Args:
        ea (int): The address to get the function name from.

    Returns:
        str: The function name.
    """
    return idc.get_func_name(ea)


def set_function_name(ea: int, name: str) -> bool:
    """
    Set the name of a function at the specified address.

    Parameters:
        ea (int): The effective address of the function to rename.
        name (str): The new name to assign to the function.

    Returns:
        bool: True if the function name was successfully set, False otherwise.
    """
    return idc.set_name(ea, name, idc.SN_CHECK)


def get_real_function_ea(ea: int) -> int:
    """
    Retrieve the function name at the given address.

    Args:
        ea (int): The address to get the function name from.

    Returns:
        str: The function name.
    """

    if idc.is_code(idc.get_full_flags(ea)) and idc.print_insn_mnem(ea) == "jmp":
        real_ea = idc.get_operand_value(ea, 0)
    elif idc.is_data(idc.get_full_flags(ea)):
        real_ea = read_pointer(ea)
    else:
        real_ea = ea

    return real_ea


def demangle_name(mangled_name: str) -> str:
    """
    Demangle a name to get the human-readable form.

    Args:
        mangled_name (str): The mangled name to demangle.

    Returns:
        str: The demangled name, or the original name if demangling fails.
    """
    if str is None or str == "":
        return ""

    if idaapi.get_inf_structure().cc.id == idc.COMP_GNU:
        if mangled_name.startswith("___"):
            return mangled_name
        if not mangled_name.startswith("__Z"):
            mangled_name = "__Z" + mangled_name

    demangled_name = ida_name.demangle_name(mangled_name, ida_name.DQT_COMPILER)
    return demangled_name if demangled_name is not None else mangled_name


def read_pointer(ea: int) -> int:
    """
    Read a pointer-sized value from the given address.

    Args:
        ea (int): The address to read from.

    Returns:
        int: The pointer-sized value.
    """
    size = get_pointer_size()
    if size == 8:
        return ida_bytes.get_qword(ea)
    else:
        return ida_bytes.get_dword(ea)


def is_valid_pointer(ea: int) -> bool:
    """
    Check if the specified address contains a valid pointer.

    Args:
        ea (int): The address to check.

    Returns:
        bool: True if the address contains a valid pointer, False otherwise.
    """
    if ea == idc.BADADDR:
        return False
    value = read_pointer(ea)
    seg = ida_segment.getseg(value)
    return seg is not None


def is_valid_function_ea(ea: int) -> bool:
    """
    Check if the specified address is a valid function pointer.

    Args:
        ea (int): The address to check.

    Returns:
        bool: True if the address is a valid function pointer, False otherwise.
    """
    func = ida_funcs.get_func(ea)
    if func is None:
        seg = ida_segment.getseg(ea)
        if seg is not None and idaapi.get_segm_class(seg) == "XTRN":
            return True
    else:
        return True
    return False


def is_typeinfo(ea: int) -> bool:
    """
    Check if the specified address is a typeinfo structure.

    A typeinfo structure typically contains:
    - A pointer to a vtable for typeinfo
    - A pointer to a type name

    Args:
        ea (int): The address to check.

    Returns:
        bool: True if the address is a typeinfo structure, False otherwise.
    """
    # Check if vtable pointer points to a vtable for typeinfo
    # We assume the vtable for typeinfo has a known pattern or known address
    vtable_name = ida_name.get_ea_name(ea, ida_name.GN_DEMANGLED)
    if "typeinfo" in vtable_name:
        return True

    # Read vtable pointer
    if not is_valid_pointer(ea):
        return False

    # Read name pointer
    name_ptr = ea + get_pointer_size()
    if not is_valid_pointer(name_ptr):
        return False

    # Validate the name pointer
    type_name = get_string(read_pointer(name_ptr))
    if len(type_name) == 0:
        return False

    return True


def is_vtable_entry(ea: int) -> bool:
    """
    Check if the specified address is a valid vtable entry.

    Args:
        ea (int): The address to check.

    Returns:
        bool: True if the address is a valid vtable entry, False otherwise.
    """
    if is_valid_pointer(ea):
        ptr = read_pointer(ea)
        return is_valid_function_ea(ptr)
    return False


def parse_vtable_entries(vtblptr: int, count: int = -1) -> List[Tuple[int, int]]:
    """
    Parse the vtable to retrieve function pointers.

    Args:
        vtblptr (int): The address of the vtable.
        count (int): The number of entries to parse (default is -1 for all entries).

    Returns:
        List[Tuple[int, int]]: A list of tuples containing the vtable entry address and function pointer.
    """
    pointer_size = get_pointer_size()
    vtable_entries = []

    if count == -1:
        i = 0
        while True:
            entry = read_pointer(vtblptr + i * pointer_size)
            if is_vtable_entry(vtblptr + i * pointer_size):
                vtable_entries.append((vtblptr + i * pointer_size, entry))
            else:
                break
            i += 1
    else:
        for i in range(count):
            entry = read_pointer(vtblptr + i * pointer_size)
            if is_vtable_entry(vtblptr + i * pointer_size):
                vtable_entries.append((vtblptr + i * pointer_size, entry))
            else:
                break

    return vtable_entries


def get_class_name(type_info_ea: int, demangle: bool = True) -> str:
    """
    Retrieve the class name from the __class_type_info structure.

    Args:
        type_info_ea (int): The address of the __class_type_info structure.
        demangle (bool): Whether to demangle the class name (default is True).

    Returns:
        str: The class name.
    """
    name_pointer = read_pointer(type_info_ea + get_pointer_size())
    name = get_string(name_pointer)
    if not demangle:
        return name
    else:
        return demangle_name(name)


def is_si_class_type_info(type_info_ea: int) -> bool:
    """
    Check if the given address points to a __si_class_type_info structure.

    Args:
        type_info_ea (int): The address to check.

    Returns:
        bool: True if the address points to a __si_class_type_info structure, False otherwise.
    """
    if not is_typeinfo(type_info_ea):
        return False

    si_class_type_info_offset = 16 if idaapi.get_inf_structure().is_64bit() else 8

    base_type_pointer = read_pointer(type_info_ea + si_class_type_info_offset)
    return is_typeinfo(base_type_pointer)


def is_vmi_class_type_info(type_info_ea: int) -> bool:
    """
    Check if the given address points to a __vmi_class_type_info structure.

    Args:
        type_info_ea (int): The address to check.

    Returns:
        bool: True if the address points to a __vmi_class_type_info structure, False otherwise.
    """
    if not is_typeinfo(type_info_ea):
        return False

    if idaapi.get_inf_structure().is_64bit():
        vmi_class_type_info_offset = 16
        flags = ida_bytes.get_dword(type_info_ea + vmi_class_type_info_offset)
        base_count = ida_bytes.get_dword(type_info_ea + vmi_class_type_info_offset + 4)
        return flags >> 2 == 0 and base_count != 0
    else:
        vmi_class_type_info_offset = 8
        flags = ida_bytes.get_word(type_info_ea + vmi_class_type_info_offset)
        base_count = ida_bytes.get_word(type_info_ea + vmi_class_type_info_offset + 4)
        return flags >> 2 == 0 and base_count != 0


def identify_rtti_type(type_info_ea: int) -> int:
    """
    Identify the RTTI type of the structure at the given address.

    Args:
        type_info_ea (int): The address to check.

    Returns:
        int: The RTTI type (CLASS_TYPE_INFO, SI_CLASS_TYPE_INFO, or VMI_CLASS_TYPE_INFO).
    """
    if is_si_class_type_info(type_info_ea):
        return SI_CLASS_TYPE_INFO
    elif is_vmi_class_type_info(type_info_ea):
        return VMI_CLASS_TYPE_INFO
    else:
        return CLASS_TYPE_INFO


def parse_si_class_type_info(type_info_ea: int) -> int:
    """
    Parse the __si_class_type_info structure to retrieve the base class information.

    Args:
        type_info_ea (int): The address of the __si_class_type_info structure.

    Returns:
        int: The address of the base class typeinfo structure.
    """
    si_class_type_info_offset = 16 if idaapi.get_inf_structure().is_64bit() else 8
    return read_pointer(type_info_ea + si_class_type_info_offset)


def get_base_class_info(base_info_ea: int) -> Tuple[int, int]:
    """
    Retrieve the base class type information from the __base_class_type_info structure.

    Args:
        base_info_ea (int): The address of the __base_class_type_info structure.

    Returns:
        tuple: A tuple containing the base type pointer and the offset_flags.
    """
    base_type_pointer = read_pointer(base_info_ea)
    pointer_size = get_pointer_size()
    if pointer_size == 8:
        offset_flags = ida_bytes.get_qword(base_info_ea + pointer_size)
    else:
        offset_flags = ida_bytes.get_dword(base_info_ea + pointer_size)
    return base_type_pointer, offset_flags


def parse_offset_flags(offset_flags: int) -> Tuple[int, int]:
    """
    Parse the __offset_flags to extract the signed offset and flags.

    Args:
        offset_flags (int): The offset_flags value to parse.

    Returns:
        tuple: A tuple containing the offset and the flags.
    """
    flags = offset_flags & 0xFF
    offset = offset_flags >> OFFSET_SHIFT

    # Convert offset to signed if necessary
    # Check if the sign bit is set (for 64-bit signed)
    if offset & (1 << (64 - OFFSET_SHIFT - 1)):
        offset -= (1 << (64 - OFFSET_SHIFT))

    return offset, flags


def parse_vmi_class_type_info(type_info_ea: int) -> Tuple[int, list]:
    """
    Parse the __vmi_class_type_info structure to retrieve base class information.

    Args:
        type_info_ea (int): The address of the __vmi_class_type_info structure.

    Returns:
        tuple: A tuple containing the flags and a list of base class information (base class name, offset, flags).
    """
    pointer_size = get_pointer_size()
    vmi_class_type_info_offset = 16 if idaapi.get_inf_structure().is_64bit() else 8

    if pointer_size == 8:
        flags = ida_bytes.get_dword(type_info_ea + vmi_class_type_info_offset)
    else:
        flags = ida_bytes.get_word(type_info_ea + vmi_class_type_info_offset)

    if pointer_size == 8:
        base_count = ida_bytes.get_dword(type_info_ea + vmi_class_type_info_offset + 4)
    else:
        base_count = ida_bytes.get_word(type_info_ea + vmi_class_type_info_offset + 2)

    vmi_class_base_info_offset = vmi_class_type_info_offset + pointer_size
    base_info_array = type_info_ea + vmi_class_base_info_offset

    base_classes = []
    for i in range(base_count):
        base_info_ea = base_info_array + i * (2 * pointer_size)
        base_type_pointer, offset_flags = get_base_class_info(base_info_ea)
        base_name = get_class_name(base_type_pointer)
        offset, base_flags = parse_offset_flags(offset_flags)
        base_classes.append((base_name, offset, base_flags))

    return flags, base_classes


def is_marked_as_itanium_vtable(ea: int) -> bool:
    """
    Check if the specified address is a vtable by examining its name.

    Args:
        ea (int): The address to check.

    Returns:
        bool: True if the address is a vtable, False otherwise.
    """
    name = ida_name.get_ea_name(ea)
    if name == "" or not name.startswith("__ZTV"):
        return False
    return True


def is_marked_as_itanium_typeinfo(ea: int) -> bool:
    """
    Check if the specified address is a typeinfo structure by examining its name.

    Args:
        ea (int): The address to check.

    Returns:
        bool: True if the address is a typeinfo structure, False otherwise.
    """
    name = ida_name.get_ea_name(ea)
    if name == "" or not name.startswith("__ZTI"):
        return False
    return True


def generate_class_declaration(name: str, methods: List[str], base_classes: List[Tuple[str, int, int]], size: int = 0) -> List[str]:
    """
    Generate a C++ class declaration from the class name, methods, and optionally base classes and size.

    Args:
        name (str): The name of the class.
        methods (List[str]): A list of method names.
        base_classes (List[Tuple[str, int, int]]): A list of base classes with their information.
        size (int): The size of the class (default is 0).

    Returns:
        List[str]: The lines of the class declaration.
    """
    lines = []
    if len(base_classes) == 0:
        lines.append(f"class {name} {{")
    elif len(base_classes) == 1:
        lines.append(f"class {name} : public {base_classes[0][0]} {{")
    else:
        parents = []
        for base in base_classes:
            parent = "virtual " if base[2] == VIRTUAL_MASK else ""
            if base[2] == PUBLIC_MASK:
                parent += "public "
            parent += base[0]
            parents.append(parent)

        parent_list = ", ".join(parents)
        lines.append(f"class {name} : {parent_list} {{")

    lines.append("public:")
    offset = 0
    for method in methods:
        if method == "___cxa_pure_virtual":
            decl = f"    virtual void* pure_virtual_{offset:x}();"
        elif method.endswith("D1Ev"):
            decl = f"    virtual ~{name}();"
        elif method.endswith("D0Ev"):
            continue
        else:
            method_name = demangle_name(method)
            decl = f"    virtual void* {method_name};"
            if decl.endswith("const;"):
                decl = decl.replace("const;", " /* const */;")
            decl = decl.replace("(void)", "()")

        lines.append(decl)
        offset += get_pointer_size()

    if size != 0:
        lines.append("public:")
        for offset in range(get_pointer_size(), size, get_pointer_size()):
            decl = f"    void* field_{offset:x};"
            lines.append(decl)

    lines.append("};")
    return lines


def dump_class_info_from_vtblptr(ea: int) -> Tuple[str, List[Tuple[str, int, int]], List[str]]:
    """
    Dump the class information from the vtable pointer at the specified address.

    Args:
        ea (int): The address of the vtable.

    Returns:
        Tuple[str, List[Tuple[str, int, int]], List[str]]: The class name, base classes, and methods.
    """
    # vtable_name = ida_name.get_ea_name(ea, ida_name.GN_DEMANGLED)
    # if "vtable" not in vtable_name:
    #     return []
    # class_name = vtable_name.split("'")[1]

    pointer_size = get_pointer_size()

    type_info_ea = read_pointer(ea - pointer_size)
    if not is_typeinfo(type_info_ea):
        return "", [], []

    class_name = get_class_name(type_info_ea)

    rtti_type = identify_rtti_type(type_info_ea)
    if rtti_type == SI_CLASS_TYPE_INFO:
        base_type_info = parse_si_class_type_info(type_info_ea)
        base_class_name = get_class_name(base_type_info)
        flags = 0
        base_classes = [(base_class_name, 0, 0)]
    elif rtti_type == VMI_CLASS_TYPE_INFO:
        flags, base_classes = parse_vmi_class_type_info(type_info_ea)
    else:
        flags = 0
        base_classes = []

    entries = parse_vtable_entries(ea)

    methods = []
    for entry in entries:
        func_name = ida_name.get_ea_name(entry[1])
        if func_name.startswith("__ZN") or func_name == "___cxa_pure_virtual":
            methods.append(func_name)
        else:
            break

    return class_name, base_classes, methods


def dump_class_from_constructor(ctor_ea: int) -> bool:
    """
    Dump the class information starting from a constructor function.

    Args:
        ea (int): The address of the constructor function.

    Returns:
        bool: True if the class was successfully dumped, False otherwise.
    """
    func = ida_funcs.get_func(ctor_ea)
    if func is None:
        return False

    vtblptrs = []
    for addr in idautils.FuncItems(func.start_ea):
        if idc.is_code(idc.get_full_flags(addr)) and idc.print_insn_mnem(addr) == "lea":
            value = idc.get_operand_value(addr, 1)
            seg = ida_segment.getseg(value)
            if seg is not None:
                vtblptrs.append(value)

    for vtblptr in vtblptrs:
        class_name, base_classes, methods = dump_class_info_from_vtblptr(
            vtblptr)
        if class_name != "":
            lines = generate_class_declaration(class_name, methods, base_classes)
            lines.insert(0, "#pragma once\n")
            output_lines(lines, f"{class_name}.h")

    return True


def dump_class(ea: int = idc.BADADDR) -> bool:
    """
    Dump the class information from the specified address.

    Args:
        ea (int): The address to start dumping from (default is the current screen address).

    Returns:
        bool: True if the class was successfully dumped, False otherwise.
    """
    if ea == idc.BADADDR:
        ea = idc.get_screen_ea()
    seg = ida_segment.getseg(ea)
    if seg is None:
        return False

    if seg.perm & idaapi.SEGPERM_EXEC != 0:
        dump_class_from_constructor(ea)
    else:
        class_name, base_classes, methods = dump_class_info_from_vtblptr(ea)
        if class_name != "":
            lines = generate_class_declaration(class_name, methods, base_classes)
            lines.insert(0, "#pragma once\n")
            output_lines(lines, f"{class_name}.h")

    return True


def dump_vtable_entries(ea_start: int, ea_end: int, demangle: bool = False, pair: bool = False) -> List[str]:
    """
    Generate a list of function names (or address-name pairs) for the specified address range.

    Args:
        ea_start (int): The starting address of the range.
        ea_end (int): The ending address of the range.
        demangle (bool): Whether to demangle the names (default is False).
        pair (bool): Whether to return a list of (address, name) tuples (default is False).

    Returns:
        Union[List[str], List[Tuple[int, str]]]: A list of function names if `pair` is False,
        or a list of (address, name) tuples if `pair` is True.
    """
    entries = []
    for ea in range(ea_start, ea_end, get_pointer_size()):
        func_ea = read_pointer(ea)
        if not is_valid_function_ea(func_ea):
            break
        func_name = ida_name.get_ea_name(func_ea)
        if demangle:
            func_name = demangle_name(func_name)
        if pair:
            entries.append(f"{func_ea}:{func_name}")
        else:
            entries.append(func_name)
    return entries


def rename_vtable_entries(ea: int, names: List[str]) -> List[str]:
    """
    Rename vtable entries based on the provided list of names.

    Args:
        ea (int): The address of the vtable.
        names (List[str]): The list of new function names.

    Returns:
        List[str]: A list of old function names that were replaced.
    """
    if not is_vtable_entry(ea) or names is None:
        return []

    pointer_size = get_pointer_size()
    i = 0
    old_names = []
    for name in names:
        if name.endswith("D0Ev"):
            continue
        func_ea = read_pointer(ea + i * pointer_size)
        i += 1
        if not is_valid_function_ea(func_ea):
            continue
        func_ea = get_real_function_ea(func_ea)
        if func_ea == idc.BADADDR:
            continue
        func_name = ida_name.get_ea_name(func_ea)
        if func_name not in ["_purecall", "___cxa_pure_virtual"]:
            set_function_name(func_ea, name)
        old_names.append(func_name)

    return old_names


def copy_vtable_entries(ea_start: int = idc.BADADDR, ea_end: int = idc.BADADDR) -> bool:
    """
    Copy vtable entries from the specified address range to the clipboard.

    Args:
        ea_start (int): The starting address of the range (default is the current selection start).
        ea_end (int): The ending address of the range (default is the current selection end).

    Returns:
        bool: True if the entries were successfully copied, False otherwise.
    """
    if ea_start == idc.BADADDR or ea_end == idc.BADADDR:
        ea_start = idc.read_selection_start()
        ea_end = idc.read_selection_end()

    entries = dump_vtable_entries(ea_start, ea_end)
    if len(entries) != 0:
        pyperclip.copy("\n".join(entries))
        return True
    return False


def paste_vtable_entries(ea: int = idc.BADADDR) -> bool:
    """
    Paste vtable entries from the clipboard and rename the functions at the specified address.

    Args:
        ea (int): The address of the vtable (default is the current screen address).

    Returns:
        bool: True if the entries were successfully pasted, False otherwise.
    """
    if ea == idc.BADADDR:
        ea = idc.get_screen_ea()

    if not is_vtable_entry(ea):
        return False

    text = pyperclip.paste()
    if text is None or len(text) == 0:
        return False

    names = text.split("\n")
    old_names = rename_vtable_entries(ea, names)
    if len(old_names) == 0:
        return False

    pyperclip.copy("\n".join(old_names))
    return True


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

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

    def activate(self, ctx) -> int:
        # Handle the action activation.
        if self.action == ACTION_DUMP:
            if dump_class(ctx.cur_ea):
                return 1
            else:
                print("No class dumped.")
                return 0
        elif self.action == ACTION_COPY_VTABLE:
            if copy_vtable_entries(ctx.cur_ea):
                return 1
            else:
                print("Failed to copy vtable")
                return 0
        elif self.action == ACTION_PASTE_VTABLE:
            if paste_vtable_entries(ctx.cur_ea):
                return 1
            else:
                print("Failed to paste vtable")
                return 0
        return 1

    def update(self, ctx) -> int:
        # This method is called to update the action's state.
        # Return idaapi.AST_ENABLE_ALWAYS to keep the action enabled.
        return idaapi.AST_ENABLE_ALWAYS


class UI_Hook(idaapi.UI_Hooks):
    """
    UI hook for attaching actions to the context menu.
    """

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

    def finish_populating_widget_popup(self, widget, popup_handle, ctx) -> None:
        # Populate the widget's context menu.
        widget_type = idaapi.get_widget_type(widget)

        if widget_type == idaapi.BWN_DISASM:
            t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(), idaapi.get_current_viewer()
            if idaapi.read_selection(view, t0, t1) or idc.get_item_size(idc.get_screen_ea()) > 1:
                idaapi.attach_action_to_popup(widget, popup_handle, ACTION_DUMP, None)
                idaapi.attach_action_to_popup(widget, popup_handle, ACTION_COPY_VTABLE, None)
                idaapi.attach_action_to_popup(widget, popup_handle, ACTION_PASTE_VTABLE, None)


class class_dumper_t(idaapi.plugin_t):
    """
    Main plugin class for class dumping.
    """
    flags = 0
    comment = "Dump class infomations"
    help = ""
    wanted_name = "ClassDumper"
    wanted_hotkey = ""

    def __init__(self):
        self.registered_actions = []
        self.ui_hook = None

    def init(self) -> int:
        # Initialize the plugin.
        print("ClassDumper plugin has been loaded.")
        # Register menu actions
        menu_actions = (
            idaapi.action_desc_t(ACTION_DUMP, "Dump class", menu_action_handler_t(ACTION_DUMP), "Ctrl+Shift+D"),
            idaapi.action_desc_t(ACTION_COPY_VTABLE, "Copy vtable", menu_action_handler_t(ACTION_COPY_VTABLE), "Ctrl+Shift+C"),
            idaapi.action_desc_t(ACTION_PASTE_VTABLE, "Paste vtable", menu_action_handler_t(ACTION_PASTE_VTABLE), "Ctrl+Shift+V"),
        )
        for action in menu_actions:
            idaapi.register_action(action)
            self.registered_actions.append(action.name)

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

        return idaapi.PLUGIN_KEEP

    def run(self, arg: int) -> None:
        # Run the plugin.
        pass

    def term(self) -> None:
        # Terminate the plugin.
        for action in self.registered_actions:
            idaapi.unregister_action(action)


def PLUGIN_ENTRY() -> class_dumper_t:
    # Plugin entry point.
    return class_dumper_t()


repository: danigargu/heap-viewer
file: heap_viewer.py
description: IDA Pro plugin to examine the glibc heap, focused on exploit development
created_at: 2018-05-07 16:46:37+00:00
pushed_at: 2022-07-25 09:21:47+00:00
url:
https://github.com/danigargu/heap-viewer/blob/f2c752cf482d199866e22b8fdf2c71109f4ef287/heap_viewer.py

content:
#!/usr/bin/python
# coding: utf-8
#
# HeapViewer - by @danigargu
#

import os
import sys
import idaapi

from heap_viewer import PLUGNAME, plugin_gui
from heap_viewer.misc import is_process_suspended, log

# -----------------------------------------------------------------------
class HeapViewPlugin(idaapi.plugin_t):
    flags = 0
    comment = ""
    help = ""
    wanted_name = PLUGNAME
    wanted_hotkey = "Ctrl-H"

    def init(self):
        self.icon_id = 0
        self.add_menus()
        return idaapi.PLUGIN_KEEP

    def run(self, arg=0):
        try:
            if "ELF" not in idaapi.get_file_type_name():
                raise Exception("Executable must be ELF fomat")

            if not idaapi.is_debugger_on() or not is_process_suspended():
                raise Exception("The debugger must be active and suspended before using this plugin")

            f = plugin_gui.HeapPluginForm()
            f.Show()

        except Exception as e:
            idaapi.warning("[%s] %s" % (PLUGNAME, str(e)))

    def add_menus(self):
        # To avoid creating multiple plugin_t instances
        this = self
        class StartHandler(idaapi.action_handler_t):
            def __init__(self):
                idaapi.action_handler_t.__init__(self)
                
            def activate(self, ctx):
                this.run()
                return 1

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

        act_name = '%s:start' % PLUGNAME
        act_desc = idaapi.action_desc_t(
            act_name,       # The action name. Must be unique
            PLUGNAME,       # Action Text
            StartHandler(), # Action handler
            None,           # Optional shortcut
            'Start plugin', # Action tooltip
            122             # Icon
        )
        idaapi.register_action(act_desc)
        idaapi.attach_action_to_menu(
            'Debugger/Debugger windows/',
            act_name,
            idaapi.SETMENU_APP
        )

    def term(self):
        idaapi.msg("[%s] terminated" % (PLUGNAME))

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

# -----------------------------------------------------------------------
log("Plugin loaded")


repository: WPeace-HcH/SysNR-FuncFinder
file: SysNR-FuncFinder.py
description: A plugin for IDA that renames functions by system call numbers.
wanted_name: SysNR-FuncFinder
comment: SysNR-FuncFinder Plugin for IDA
created_at: 2022-11-30 06:54:32+00:00
pushed_at: 2023-09-20 08:33:36+00:00
url:
https://github.com/WPeace-HcH/SysNR-FuncFinder/blob/9678c39ba392a2885a0baeb2531f73562009bfa8/SysNR-FuncFinder.py

content:
import idc
import idaapi
import sys, os
path = os.path.dirname(os.path.abspath(__file__)) + "\\LffPlugDir_WPeace\\"
sys.path.append(path)
import LinuxFuncFinder_x64
import LinuxFuncFinder_x86
import LinuxFuncFinder_Mips32
import LinuxFuncFinder_Arm32
import eabiFuncFinder_Arm32
import LinuxFuncFinder_PPC32


class myplugin_sysnr(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "SysNR-FuncFinder Plugin for IDA"
    help = "Find more information at https://github.com/wpeace-hch"
    wanted_name = "SysNR-FuncFinder"
    wanted_hotkey = "Ctrl-Alt-F"
    def init(self):
        print("\nSysNR-FuncFinder By WPeace.")
        try:
            WPe_Patcher.register(self, "SysNR-FuncFinder    (Ctrl-Alt-F)")
            WPe_About.register(self, "About")
        except:
            pass
        if idaapi.IDA_SDK_VERSION >= 740:
            idaapi.attach_action_to_menu("Edit/WPeace_Plugins/SysNR-FuncFinder    (Ctrl-Alt-F)", WPe_Patcher.get_name(), idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/WPeace_Plugins/About", WPe_About.get_name(), idaapi.SETMENU_APP)
        else:
            print("Your IDA version needs to be greater than 7.4! :(@WPeace")
        return idaapi.PLUGIN_OK
    def run(self, arg):
        print("SysNR-FuncFinder v2.5 start running...")
        self.patcher()
    def term(self):
        print("SysNR-FuncFinder v2.5 works fine! :)@WPeace\n")
    def patcher(self):
        elf_magic = idc.get_wide_dword(idc.get_first_seg())
        e_type = idc.get_wide_word(idc.get_first_seg() + 0x10)
        if elf_magic == 0x464c457f or elf_magic == 0x7f454c46:
            if e_type == 2:
                e_flags = idc.get_wide_dword(idc.get_first_seg() + 0x24)
                e_machine = idc.get_wide_word(idc.get_first_seg() + 0x12)
                # eabi_syscall
                if e_flags > 0x4000000:
                    # ARM32
                    if e_machine == 40:
                        eabiFuncFinder_Arm32.main()
                # oabi_syscall
                else:
                    # AMD x86-64 architecture
                    if e_machine == 62:
                        LinuxFuncFinder_x64.main()
                    # Intel 80386
                    elif e_machine == 3:
                        LinuxFuncFinder_x86.main()
                    # MIPS I Architecture
                    elif e_machine == 8:
                        LinuxFuncFinder_Mips32.main()
                    # Advanced RISC Machines ARM
                    elif e_machine == 40:
                        LinuxFuncFinder_Arm32.main()
                    # PowerPC32
                    elif e_machine == 20:
                        LinuxFuncFinder_PPC32.main()
                    else:
                        print("请确认插件版本是否支持当前文件架构。")
            else:
                 print("当前插件仅支持EXEC可执行ELF文件。")
        else:
            print("当前插件仅支持ELF文件格式。")
    def about(self):
        f = About_Form()
        f.Execute()
        f.Free()


class 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(),
                instance.get_label(),
                instance
            ))
            
        @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):
            return 1
            
        @classmethod
        def update(self, ctx):
            try:
                return idaapi.AST_ENABLE_FOR_WIDGET
            except Exception as e:
                return idaapi.AST_ENABLE_ALWAYS


class About_Form(idaapi.Form):
    def __init__(self):
        super(About_Form, self).__init__(r"""STARTITEM 0
BUTTON YES* Open author's github
ABOUT
{FormChangeCb}
Plugins for IDA.
Written BY WPeace.
      
""", {
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })
        self.Compile()

    def OnFormChange(self, fid):
        if fid == -2:
            import webbrowser
            webbrowser.open("https://github.com/wpeace-hch", new = 2)
        return 1


class WPe_Patcher(Menu_Context):
        def activate(self, ctx):
            print("\nSysNR-FuncFinder v2.5 start running...")
            self.plugin.patcher()
            return 1


class WPe_About(Menu_Context):
    def activate(self, ctx):
        self.plugin.about()
        return 1


def PLUGIN_ENTRY():
    return myplugin_sysnr()


repository: tintinweb/vscode-decompiler
file: scripts/ida_batch_decompile.py
description: Decompile things directly from VSCode
wanted_name: Ida Batch Decompile
comment: Batch Decompile
created_at: 2020-05-12 19:15:11+00:00
pushed_at: 2023-01-09 00:59:33+00:00
url:
https://github.com/tintinweb/vscode-decompiler/blob/fb499bd90e8bb0607e5581659d602cc7af5d90e3/scripts/ida_batch_decompile.py

content:
#! /usr/bin/env python
# -*- coding: UTF-8 -*-
# Author : <github.com/tintinweb>
"""
IdaBatchDecompile Plugin and Script adds annotation and batch decompilation functionality to IDA Pro

* requires hexrays decompiler plugin

Usage:

* as idascript in ida gui mode: IDA Pro -> File/Script file... -> IdaDecompileBatch ...
* as idascript in ida cmdline mode: ida(w|w64) -B -M -S"<path_to_this_script> \"--option1\" \"--option2\"", "<target>"
 * see --help for options
* as Plugin: follow ida documentation on how to add python plugins

"""
import sys
import json
import glob
import subprocess
import shutil
import os
import tempfile
from optparse import OptionParser

import idaapi
import idautils
from idc import *

if idaapi.IDA_SDK_VERSION >= 700:
    import ida_idaapi
    import ida_kernwin
    from idaapi import *
    from idc import *

import logging

logger = logging.getLogger(__name__)


class IdaLocation(object):
    """ Wrap idautils Function
    """

    def __init__(self, location):
        self.at = location
        # self.name = GetFunctionName(location)
        self.name = GetFuncOffset(location)
        self.start = 0
        self.end = 0
        self.func_offset = 0
        try:
            _func = idaapi.get_func(location)
            self.start = _func.startEA
            self.end = _func.endEA  # ==FindFuncEnd(location)
            self.func_offset = self.start - self.at
        except Exception, e:
            logger.exception(e)
        if not self.name:
            self.indirect = True
        else:
            self.indirect = False

    def __repr__(self, *args, **kwargs):
        return "<Function %r at 0x%x (0x%x::0x%x)>" % (self.name, self.at,
                                                       self.start, self.end)

    def get_xrefs(self):
        return (IdaLocation(x.frm) for x in idautils.XrefsTo(self.at))

    def get_coderefs(self):
        return (IdaLocation(frm) for frm in idautils.CodeRefsTo(self.at, 0))

    def as_dict(self):
        return {'at': self.at, 'name': self.name}

    def decompile(self):
        """ decompile function
        """
        try:
            return idaapi.decompile(self.at)
        except idaapi.DecompilationFailure, e:
            return repr(str(e))
        text = str(idaapi.decompile(self.at)).strip()
        '''
        sprintf:
        Python>for w in idaapi.decompile(0x00001578 ).lvars: print w.name
            s
            format
            result
        '''
        # decompile.arguments
        # for w in idaapi.decompile(0x00001EF0 ).lvars: print w.name
        if not grep:
            return text.split('\n')
        # return all lines
        return [line.strip() for line in text.split('\n') if grep in line]

    def get_function_args(self):
        # find the stack frame
        stack = GetFrame(self.start)
        stack_size = GetStrucSize(stack)
        # figure out all of the variable names
        # base is either ' s' ... saved register or ' r' ... return address
        base = GetMemberOffset(stack, ' s')
        if base == -1:
            base = GetMemberOffset(stack, ' r')
        if base == -1:
            # no ' s' no ' r' assume zero
            base == 0
        stack_vars = []

        for memberoffset in xrange(stack_size):
            previous = stack_vars[-1] if len(stack_vars) else None
            var_name = GetMemberName(stack, memberoffset)
            if not var_name or (previous and var_name == previous.get("name")):
                # skip that entry, already processed
                continue

            offset = GetMemberOffset(stack, var_name) - base
            size = GetMemberSize(stack, memberoffset)
            if previous:
                diff = offset - previous['offset']
                previous['diff_size'] = diff
            stack_vars.append({'name': var_name,
                               'offset': offset,
                               'offset_text': '[bp%Xh]' % offset if offset < 0 else '[bp+%Xh]' % offset,
                               'size': size,
                               'diff_size': size})
        return stack_size, stack_vars


class IdaHelper(object):
    """ Namespace for ida helper functions
    """

    @staticmethod
    def get_functions():
        return (IdaLocation(f) for f in idautils.Functions())

    @staticmethod
    def get_imports():
        for i in xrange(0, idaapi.get_import_module_qty()):
            name = idaapi.get_import_module_name(i)
            if name:
                yield name

    @staticmethod
    def decompile_full(outfile):
        return idaapi.decompile_many(outfile, None, 0)

    @staticmethod
    def annotate_xrefs():
        stats = {'annotated_functions': 0, 'errors': 0}
        for f in IdaHelper.get_functions():
            try:
                function_comment = GetFunctionCmt(f.start, 0)
                if '**** XREFS ****' in function_comment:
                    logger.debug(
                        "[i] skipping function %r, already annotated." % f.name)
                    continue
                xrefs = [x.name for x in f.get_coderefs()]
                comment = []
                if function_comment:
                    comment.append(function_comment)
                comment.append("***** XREFS *****")
                comment.append("* # %d" % len(xrefs))
                comment.append(', '.join(xrefs))
                comment.append("*******************")
                SetFunctionCmt(f.start, '\n'.join(comment), 0)
                stats['annotated_functions'] += 1
            except Exception as e:
                print("Annotate XRefs: %r" % e)
                stats['errors'] += 1
        print "[+] stats: %r" % stats
        print "[+] Done!"

    @staticmethod
    def annotate_functions_with_local_var_size():
        stats = {'annotated_functions': 0, 'errors': 0}
        for f in IdaHelper.get_functions():
            try:
                function_comment = GetFunctionCmt(f.start, 0)
                if '**** Variables ****' in function_comment:
                    logger.debug(
                        "[i] skipping function %r, already annotated." % f.name)
                    continue
                size, stack_vars = f.get_function_args()
                comment = []
                if function_comment:
                    comment.append(function_comment)
                comment.append("**** Variables ****")
                comment.append("* stack size: %s" % size)
                for s in stack_vars:
                    comment.append(json.dumps(s))
                comment.append("*******************")
                SetFunctionCmt(f.start, '\n'.join(comment), 0)
                stats['annotated_functions'] += 1
            except Exception, e:
                print("Annotate Funcs: %r" % e)
                stats['errors'] += 1
        print "[+] stats: %r" % stats
        print "[+] Done!"


class IdaDecompileBatchController(object):
    def __init__(self):
        self.is_windows = sys.platform.startswith('win')
        # hackhackhack - check if we're ida64 or ida32
        self.is_ida64 = GetIdbPath().endswith(".i64")
        logger.debug("[+] is_windows: %r" % self.is_windows)
        logger.debug("[+] is_ida64: %r" % self.is_ida64)
        self.my_path = os.path.abspath(__file__)
        self.temp_path = None
        self._init_target()
        # settings (form)
        # todo: load from configfile if available.
        self.output_path = None
        self.chk_annotate_stackvar_size = False
        self.chk_annotate_xrefs = False
        self.chk_decompile_imports = False
        self.chk_decompile_imports_recursive = False
        self.chk_decompile_alternative = False
        # self.ida_home = idaapi.idadir(".")
        self.ida_home = GetIdaDirectory()
        # wait for ida analysis to finish
        self.wait_for_analysis_to_finish()
        if not idaapi.init_hexrays_plugin():
            logger.warning("forcing hexrays to load...")
            self.load_plugin_decompiler()
        if not idaapi.init_hexrays_plugin():
            raise Exception("hexrays decompiler is not available :(")

    def _init_target(self):
        self.target_path = idc.GetInputFilePath()
        self.target_file = idc.GetInputFile()
        self.target_dir = os.path.split(self.target_path)[0]
        logger.debug("reinitializing target: %r" % self.target_file)

    def init_tempdir(self):
        self.temp_path = self.temp_path or tempfile.mkdtemp(prefix="idbc_")
        logger.debug("[i] using tempdir: %r" % self.temp_path)

    def remove_tempdir(self):
        if not self.temp_path:
            return
        logger.debug("[i] removing tempdir: %r" % self.temp_path)
        shutil.rmtree(self.temp_path)
        self.temp_path = None

    def wait_for_analysis_to_finish(self):
        logger.debug("[+] waiting for analysis to finish...")
        idaapi.autoWait()
        idc.Wait()
        logger.debug("[+] analysis finished.")

    def load_plugin_decompiler(self):
        # load decompiler plugins (32 and 64 bits, just let it fail)
        logger.debug("[+] trying to load decompiler plugins")
        if self.is_ida64:
            # 64bit plugins
            idc.RunPlugin("hexx64", 0)
        else:
            # 32bit plugins
            idc.RunPlugin("hexrays", 0)
            idc.RunPlugin("hexarm", 0)
        logger.debug("[+] decompiler plugins loaded.")

    def run(self):
        files_decompiled = []
        self._init_target()

        if self.chk_decompile_imports:
            self.init_tempdir()
            if self.chk_decompile_imports_recursive:
                pass
            for image_type, image_name, image_path in self.enumerate_import_images():
                try:
                    self.exec_ida_batch_decompile(target=image_path, output=self.output_path,
                                                  annotate_stackvar_size=self.chk_annotate_stackvar_size,
                                                  annotate_xrefs=self.chk_annotate_xrefs,
                                                  imports=self.chk_decompile_imports,
                                                  recursive=self.chk_decompile_imports_recursive,
                                                  experimental_decomile_cgraph=self.chk_decompile_alternative)
                    files_decompiled.append(image_path)
                except subprocess.CalledProcessError, cpe:
                    logger.warning(
                        "[!] failed to decompile %r - %r" % (image_path, cpe))

            self.remove_tempdir()

        if self.chk_annotate_stackvar_size:
            self.annotate_stack_variable_size()
        if self.chk_annotate_xrefs:
            self.annotate_xrefs()

        if self.chk_decompile_alternative:
            raise NotImplemented("Not yet implemented")
            pass
        else:
            self.decompile_all(self.output_path)
            files_decompiled.append(self.target_file)

        logger.info("[+] finished decompiling: %r" % files_decompiled)
        logger.info("    output dir: %s" %
                    self.output_path if self.output_path else self.target_dir)
        return files_decompiled

    def annotate_stack_variable_size(self):
        logger.debug("[+] annotating function stack variables")
        IdaHelper.annotate_functions_with_local_var_size()
        logger.debug("[+] done.")

    def annotate_xrefs(self):
        logger.debug("[+] annotating function xrefs")
        IdaHelper.annotate_xrefs()
        logger.debug("[+] done.")

    def file_is_decompilable(self, path):
        with open(path, 'rb') as ftest:
            magic = ftest.read(4)
            if magic == 'MZ\x90\x00':
                return 'pe/dos'
            elif magic == "\x7fELF":
                return 'elf'
        return None

    def enumerate_import_images(self):
        for import_name in IdaHelper.get_imports():
            logger.debug("[i] trying to find image for %r" % import_name)
            for image_path in glob.glob(os.path.join(self.target_dir, import_name) + '*'):
                image_type = self.file_is_decompilable(image_path)
                if image_type:
                    logger.debug("[i] got image %r as %r" %
                                 (image_path, image_type))
                    yield image_type, os.path.split(image_path)[1], image_path
                    # I do not think there's any need to check other files with the same name ?!
                    break

    def enumerate_files(self, recursive=False):
        for root, dirs, files in os.walk(self.target_dir):
            for name in files:
                fpath = os.path.join(root, name)
                logger.debug("[+] checking %r" % fpath)
                try:
                    ftype = self.file_is_decompilable(fpath)
                    if ftype:
                        logger.debug("[+] is candidate %r" % [fpath, ftype])
                        yield ftype, name, fpath
                except IOError:
                    pass

    def decompile_all(self, outfile=None):
        outfile = self._get_suggested_output_filename(
            outfile or self.target_path)
        logger.warning(outfile)
        logger.debug("[+] trying to decompile %r as %r" % (self.target_file,
                                                           os.path.split(outfile)[1]))
        IdaHelper.decompile_full(outfile)
        logger.debug("[+] finished decompiling %r as %r" % (self.target_file,
                                                            os.path.split(outfile)[1]))

    def _get_suggested_output_filename(self, target):
        # /a/b/c/d/e/bin.ext
        # target is a directory
        if os.path.isdir(target):
            fname, fext = os.path.splitext(self.target_file)
            return '%s.c' % os.path.join(target, fname)
        # target is not a directory
        root, fname = os.path.split(target)
        if fname:
            fname, fext = os.path.splitext(fname)  # bin,ext
        else:
            fname, fext = os.path.splitext(self.target_file)

        # obsolete
        # suggested_outpath = '%s.c'%os.path.join(root,fname)
        # if not os.path.exists(suggested_outpath):
        #    return suggested_outpath
        return '%s.c' % os.path.join(root, fname)

    def exec_ida_batch_decompile(self, target, output, annotate_stackvar_size, annotate_xrefs, imports, recursive,
                                 experimental_decomile_cgraph):
        logger.debug("[+] batch decompile %r" % target)
        # todo: pass commandlines,
        # todo parse commandline
        script_args = ['--output=%s' % output]
        if annotate_stackvar_size:
            script_args.append("--annotate-stackvar-size")
        if annotate_xrefs:
            script_args.append("--annotate-xrefs")
        if imports:
            script_args.append("--imports")
        if recursive:
            script_args.append("--recursive")
        if experimental_decomile_cgraph:
            script_args.append("--experimental-decompile-cgraph")

        script_args = ['\\"%s\\"' % a for a in script_args]
        command = "%s %s" % (self.my_path, ' '.join(script_args))
        self._exec_ida_batch(target, command)

    def _exec_ida_batch(self, target, command):
        # build exe path
        if self.is_windows:
            ida_exe = os.path.join(
                self.ida_home, 'idaw64.exe' if self.is_ida64 else 'idaw.exe')
        else:
            ida_exe = os.path.join(
                self.ida_home, 'idal64' if self.is_ida64 else 'idal')
        '''
        https://www.hex-rays.com/products/ida/support/idadoc/417.shtml
        -B  ..  Batch mode
        -M  ..  disable mouse
        -c  ..  create new database
        -o  ..  database output path
        -S  ..  execute script
        '''
        #temp_path = os.path.join(self.temp_path, os.path.splitext(os.path.split(target)[1])[0] + '.idb')
        cmd = [ida_exe, '-B', '-M', '-c', '-o"%s"' %
               self.temp_path if self.temp_path else '', '-S"%s"' % command, '"' + target + '"']
        logger.debug(' '.join(cmd))
        logger.debug('[+] executing: %r' % cmd)
        # return 0
        # TODO: INSECURE!
        return subprocess.check_call(' '.join(cmd), shell=True)


class TestEmbeddedChooserClass(Choose, Choose2):
    """
    A simple chooser to be used as an embedded chooser
    """

    def __init__(self, title, nb=5, flags=0):
        Choose.__init__(self,
                        title,
                        [["Type", 10], ["Name", 10], ["Path", 30]],
                        flags=flags)
        Choose2.__init__(self,
                         title,
                         [["Type", 10], ["Name", 10],  ["Path", 30]],
                         embedded=True, width=50, height=10, flags=flags)
        self.n = 0
        self.items = []
        self.icon = 5
        self.selcount = 0

        self.selected = []

    def OnClose(self):
        pass

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

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

    def OnRefresh(self, n):
        print "refresh %s" % n

    def OnSelectionChange(self, sel_list):
        self.selected = sel_list

    def getSelected(self):
        for idx in self.selected:
            yield self.items[idx-1]

    def addItem(self, e):
        if e not in self.items:
            self.items.append(e)


class DecompileBatchForm(Form):
    """
    Form to prompt for target file, backup file, and the address
    range to save patched bytes.
    """

    def __init__(self, idbctrl, enumerate_imports=True, enumerate_other=False):
        self.idbctrl = idbctrl
        self.EChooser = TestEmbeddedChooserClass(
            "Batch Decompile", flags=Choose2.CH_MULTI)
        self.propagateItems(enumerate_imports=enumerate_imports,
                            enumerate_other=enumerate_other)
        Form.__init__(self,
                      r"""Ida Batch Decompile ...
{FormChangeCb}
<##Target    :{target}>
<##OutputPath:{outputPath}>
<##Annotate StackVar Size:{chkAnnotateStackVars}>
<##Annotate Func XRefs   :{chkAnnotateXrefs}>
<##Process Imports       :{chkDecompileImports}>
<##Cgraph (experimental) :{chkDecompileAlternative}>{cGroup1}>


<##Scan Target Directory:{btnLoad}> <##Recursive:{chkDecompileImportsRecursive}>{cGroup2}>
<##Decompile!:{btnProcessFiles}>
<Please select items to decompile:{cEChooser}>


""", {
                          'target': Form.FileInput(swidth=50, open=True, value=idbctrl.target_path),
                          'outputPath': Form.DirInput(swidth=50, value=idbctrl.output_path),
                          'cGroup1': Form.ChkGroupControl(("chkAnnotateStackVars", "chkAnnotateXrefs",
                                                           "chkDecompileImports",
                                                           "chkDecompileAlternative")),
                          'cGroup2': Form.ChkGroupControl(("chkDecompileImportsRecursive", )),
                          'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
                          'btnLoad':  Form.ButtonInput(self.OnButtonLoad),
                          'btnProcessFiles': Form.ButtonInput(self.OnButtonProcess),
                          'cEChooser': Form.EmbeddedChooserControl(self.EChooser),
                      })
        self.Compile()

    def propagateItems(self, enumerate_imports=False, enumerate_other=False):
        self.EChooser.addItem([self.idbctrl.file_is_decompilable(self.idbctrl.target_path),
                               os.path.split(self.idbctrl.target_path)[1],
                               self.idbctrl.target_path])

        if enumerate_imports:
            for candidate in self.idbctrl.enumerate_import_images():
                self.EChooser.addItem(list(candidate))
        if enumerate_other:
            for candidate in self.idbctrl.enumerate_files(recursive=self.chkDecompileImportsRecursive.checked):
                self.EChooser.addItem(list(candidate))

    def OnButtonProcess(self, code=0):
        # process selected files
        if not len(list(self.EChooser.getSelected())):
            logger.warning(
                "[!] Aborting. Please select at least one item from the list!")
            return

        self.idbctrl.target = self.target.value
        outputPath = self.GetControlValue(self.outputPath)
        if outputPath == '' or os.path.exists(outputPath):
            self.idbctrl.output_path = outputPath
        else:
            logger.warning("[!!] output path not valid! %r" % outputPath)
            self.idbctrl.output_path = None

        self.idbctrl.chk_annotate_stackvar_size = self.chkAnnotateStackVars.checked
        self.idbctrl.chk_decompile_imports = self.chkDecompileImports.checked
        self.idbctrl.chk_decompile_imports_recursive = self.chkDecompileImportsRecursive.checked
        self.idbctrl.chk_annotate_xrefs = self.chkAnnotateXrefs.checked
        self.idbctrl.chk_decompile_alternative = self.chkDecompileAlternative.checked
        logger.debug("[+] config updated")

        files_decompiled = []
        decompile_main_binary = False

        self.idbctrl.init_tempdir()
        for _type, name, image_path in self.EChooser.getSelected():
            if image_path is self.idbctrl.target_path:
                decompile_main_binary = True
                continue
            try:
                self.idbctrl.exec_ida_batch_decompile(target=image_path, output=outputPath,
                                                      annotate_stackvar_size=self.idbctrl.chk_annotate_stackvar_size,
                                                      annotate_xrefs=self.idbctrl.chk_annotate_xrefs,
                                                      imports=self.idbctrl.chk_decompile_imports,
                                                      recursive=self.idbctrl.chk_decompile_imports_recursive,
                                                      experimental_decomile_cgraph=self.idbctrl.chk_decompile_alternative)
                files_decompiled.append(image_path)
            except subprocess.CalledProcessError, cpe:
                logger.warning("[!] failed to decompile %r - %r" %
                               (image_path, cpe))

        self.idbctrl.remove_tempdir()
        # process current file
        if decompile_main_binary:
            # well, loop here even though we know it can only
            logger.debug("[+] decompiling current file...")
            files_decompiled += self.idbctrl.run()  # decompile main binary
            logger.info("[+] finished decompiling: %r" % files_decompiled)
            logger.info("    output dir: %s" %
                        self.idbctrl.output_path if self.idbctrl.output_path else self.idbctrl.target_dir)

    def OnButtonLoad(self, code=0):
        self.Close(0)
        self.propagateItems(enumerate_other=True, enumerate_imports=True)
        self.Execute()

    def OnFormChange(self, fid):
        # Set initial state
        INIT = -1
        BTN_OK = -2

        if fid == INIT:
            self.EnableField(self.target, False)
            self.EnableField(self.outputPath, True)
            self.EnableField(self.chkDecompileAlternative, False)

        elif fid == BTN_OK:
            # just return
            return True

        # Toggle backup checkbox
        elif fid == self.chkAnnotateStackVars.id:
            self.chkAnnotateStackVars.checked = not self.chkAnnotateStackVars.checked
        elif fid == self.chkDecompileImports.id:
            self.chkDecompileImports.checked = not self.chkDecompileImports.checked
        elif fid == self.chkDecompileImportsRecursive.id:
            self.chkDecompileImportsRecursive.checked = not self.chkDecompileImportsRecursive.checked
        elif fid == self.chkDecompileAlternative.id:
            self.chkDecompileAlternative.checked = not self.chkDecompileAlternative.checked
        elif fid == self.chkAnnotateXrefs.id:
            self.chkAnnotateXrefs.checked = not self.chkAnnotateXrefs.checked

        return False


if idaapi.IDA_SDK_VERSION >= 700:
    class IdaDecompileUiActionHandler(idaapi.action_handler_t):

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

        def activate(self, ctx):
            self.caller.menu_config()
            return 1

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

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


class IdaDecompileBatchPlugin(idaapi.plugin_t):
    """ IDA Plugin Base"""
    flags = idaapi.PLUGIN_FIX
    comment = "Batch Decompile"
    help = "github.com/tintinweb"
    wanted_name = "Ida Batch Decompile"
    wanted_hotkey = ""
    wanted_menu = "File/Produce file/", "{} ...".format(wanted_name)
    wanted_menu_id = 'tintinweb:batchdecompile'

    def init(self):
        NO_HOTKEY = ""
        SETMENU_INS = 0
        NO_ARGS = tuple()

        logger.debug("[+] %s.init()" % self.__class__.__name__)
        self.menuitems = []

        logger.debug("[+] setting up menus for ida version %s" %
                     idaapi.IDA_SDK_VERSION)

        if idaapi.IDA_SDK_VERSION >= 700:
            # >= 700
            action_desc = idaapi.action_desc_t(
                "tintinweb:batchdecompile:load", self.wanted_name, IdaDecompileUiActionHandler(self))
            idaapi.register_action(action_desc)
            idaapi.attach_action_to_menu(
                ''.join(self.wanted_menu), "tintinweb:batchdecompile:load", idaapi.SETMENU_APP)

        else:
            menu = idaapi.add_menu_item(self.wanted_menu[0],
                                        self.wanted_menu[1],
                                        NO_HOTKEY,
                                        SETMENU_INS,
                                        self.menu_config,
                                        NO_ARGS)

            self.menuitems.append(menu)

        return idaapi.PLUGIN_KEEP

    def run(self, arg=None):
        logger.debug("[+] %s.run()" % self.__class__.__name__)

    def term(self):
        logger.debug("[+] %s.term()" % self.__class__.__name__)
        if idaapi.IDA_SDK_VERSION < 700:
            for menu in self.menuitems:
                idaapi.del_menu_item(menu)

    def menu_config(self):
        logger.debug("[+] %s.menu_config()" % self.__class__.__name__)
        self.idbctrl._init_target()  # force target reinit
        DecompileBatchForm(self.idbctrl).Execute()

    def set_ctrl(self, idbctrl):
        logger.debug("[+] %s.set_ctrl(%r)" %
                     (self.__class__.__name__, idbctrl))
        self.idbctrl = idbctrl


def PLUGIN_ENTRY(mode=None):
    """ check execution mode:

        a) as Plugin, return plugin object
        b) as script as part of a batch execution, do not spawn plugin object
     """
    logging.basicConfig(level=logging.DEBUG,
                        format="[%(name)s/%(process)s][%(levelname)-10s] [%(module)s.%(funcName)-14s] %(message)s")
    logger.setLevel(logging.DEBUG)
    # always wait for analysis to finish
    logger.debug("[+] initializing IdaDecompileBatchPlugin")
    # create our controller interface
    idbctrl = IdaDecompileBatchController()
    # parse cmdline
    if mode == '__main__':
        # cmdline mode
        if len(idc.ARGV) > 1:
            # cmdline batch mode
            logger.debug("[+] Mode: commandline")
            parser = OptionParser()
            parser.add_option("-o", "--output", dest="output",
                              help="output path")
            parser.add_option("-S", "--annotate-stackvar-size",
                              action="store_true", default=False,
                              help="Generate stack variable size annotations")
            parser.add_option("-X", "--annotate-xrefs",
                              action="store_true", default=False,
                              help="Generate xref annotations")
            parser.add_option("-I", "--imports",
                              action="store_true", default=False,
                              help="try to decompile files referenced in IAT")
            parser.add_option("-R", "--recursive",
                              action="store_true", default=False,
                              help="Recursive decompile files/imports")
            parser.add_option("-Z", "--experimental-decompile-cgraph",
                              action="store_true", default=False,
                              help="[experimental] decompile funcs referenced in calltree manually")

            options, args = parser.parse_args(idc.ARGV[1:])
            # set options
            idbctrl.output_path = options.output
            idbctrl.chk_annotate_stackvar_size = options.annotate_stackvar_size
            idbctrl.chk_annotate_xrefs = options.annotate_xrefs
            idbctrl.chk_decompile_imports = options.imports
            idbctrl.chk_decompile_imports_recursive = options.recursive
            idbctrl.chk_decompile_alternative = options.experimental_decompile_cgraph
            # set all the idbctrl checkboxes and files
            idbctrl.run()
            idc.Exit(0)
            # return

        logger.debug("[+] Mode: commandline w/o args")
        # PluginMode
        plugin = IdaDecompileBatchPlugin()
        plugin.set_ctrl(idbctrl=idbctrl)
        plugin.init()
        logger.info("[i] %s loaded, see Menu: %s" % (IdaDecompileBatchPlugin.wanted_name,
                                                     IdaDecompileBatchPlugin.wanted_menu))
        # plugin.menu_config()
        return plugin

    else:
        logger.debug("[+] Mode: plugin")
        # PluginMode
        plugin = IdaDecompileBatchPlugin()
        plugin.set_ctrl(idbctrl=idbctrl)
        return plugin


if __name__ == '__main__':
    PLUGIN_ENTRY(mode=__name__)


repository: binsync/libbs
file: examples/template_plugin_entry.py
created_at: 2023-04-20 19:13:21+00:00
pushed_at: 2025-04-25 18:35:23+00:00
description:
A library for writing plugins in any decompiler: includes API lifting, common data formatting, and GUI abstraction!

url:
https://github.com/binsync/libbs/blob/015a1fbea4d066a251050e5fb940e66cdc9e3741/examples/template_plugin_entry.py

content:
# REPLACE_ME: with the description of the plugin you want displayed in Ghidra, and update below items
# @author YourNameHere
# @category YourCategoryHere
# @menupath Tools.MyPlugin.Replace me with short desc shown in Tools>MyPlugin menu

# REPLACE_ME: replace the command to run your plugin from Ghidra Python2 side
plugin_command = "my_library_name --run"

def create_plugin(*args, **kwargs):
    # REPLACE_ME this import with an import of your plugin's create_plugin function
    from my_library_name import create_plugin as _create_plugin
    return _create_plugin(*args, **kwargs)

# =============================================================================
# LibBS generic plugin loader (don't touch things below this)
# =============================================================================

import sys
# Python 2 has special requirements for Ghidra, which forces us to use a different entry point
# and scope for defining plugin entry points
if sys.version[0] == "2":
    # Do Ghidra Py2 entry point
    import subprocess
    from libbs_vendored.ghidra_bridge_server import GhidraBridgeServer

    GhidraBridgeServer.run_server(background=True)
    process = subprocess.Popen(plugin_command.split(" "))
    if process.poll() is not None:
        raise RuntimeError(
            "Failed to run the Python3 backed. It's likely Python3 (and its scripts) is not in your Path inside Ghidra."
        )
else:
    # Try plugin discovery for other decompilers
    try:
        import idaapi
        has_ida = True
    except ImportError:
        has_ida = False
    try:
        import angrmanagement
        has_angr = True
    except ImportError:
        has_angr = False

    if not has_ida and not has_angr:
        create_plugin()
    elif has_angr:
        from angrmanagement.plugins import BasePlugin
        class AngrBSPluginThunk(BasePlugin):
            def __init__(self, workspace):
                super().__init__(workspace)
                globals()["workspace"] = workspace
                self.plugin = create_plugin()

            def teardown(self):
                pass


def PLUGIN_ENTRY(*args, **kwargs):
    """
    This is the entry point for IDA to load the plugin.
    """
    return create_plugin(*args, **kwargs)


repository: feicong/ida_for_mac_green
file: py/hexrays_hlight.py
description: IDA Pro for macOS绿化
wanted_name: HexLight
comment: highlights the matching brace in Pseudocode-View
created_at: 2017-08-15 10:19:47+00:00
pushed_at: 2017-10-09 08:37:02+00:00
url:
https://github.com/feicong/ida_for_mac_green/blob/5c256398100688125e957fac45630da0802ee1f7/py/hexrays_hlight.py

content:
# highlighting plugin for Hex-Rays Decompiler
# Copyright (c) 2016
# Milan Bohacek <milan.bohacek+hexlight@gmail.com>
# All rights reserved.
# 
# ==============================================================================
# 
# This file is part of Hexlight.
# 
# Hexlight is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# 
# ==============================================================================


import idautils
import idaapi
import idc


import traceback

hexlight_cb_info = None
hexlight_cb = None

posledni = 0

def jump(custom_viewer, line):
    (pl, x, y)=idaapi.get_custom_viewer_place(custom_viewer, False)
    pl2 = idaapi.place_t_as_simpleline_place_t(pl.clone())
    oldline = pl2.n
    pl2.n = line
    idaapi.jumpto(custom_viewer, pl2, x, y)
    return oldline

class hexrays_callback_info(object):
    
    def __init__(self):
        self.vu = None
        self.highlights = {}
        self.hicolor = 0xF2E8BF #0x00ffff00
        self.theotherline = None
        self.safe = False
        return
    
    def clearall(self, ps, refresh=True):
        ctr = 0
        for i in self.highlights:
            try:            
                ps[i].bgcolor = self.highlights[i]
                ctr += 1
            except:
                pass
#               place_t_as_simpleline_place_t
#              ps[i].line.replace("\x04{\x04", "{")
#              ps[i].line.replace("\x04}\x04", "}")
            
        self.highlights = {}
        self.theotherline = None
        if((ctr > 0) and refresh):
            idaapi.refresh_idaview_anyway()



    def event_callback(self, event, *args):
        try:
#            print "event: %d"%event
            if event == idaapi.hxe_keyboard:
                vu, keycode, shift = args

                if idaapi.lookup_key_code(keycode, shift, True) == idaapi.get_key_code("B") and shift == 0:
                    if self.theotherline:
                        self.theotherline = jump(vu.ct, self.theotherline)
                    return 0


            if event <= idaapi.hxe_print_func:
                self.safe = False

            if event == idaapi.hxe_switch_pseudocode:
                self.safe = False

            if event == idaapi.hxe_func_printed:
                self.safe = True

            if event == idaapi.hxe_text_ready:
                self.safe = True

            if event == idaapi.hxe_curpos:
                if not self.safe:
                    return 0
                #print "1"
                self.vu = args[0]

                if not self.vu:
                    return 0
                #print "2"

                if self.vu.cfunc.maturity != idaapi.CMAT_FINAL:
                    return 0
                #print "3"

                if not self.vu.visible():
                    return 0
                #print "4"
                if not self.vu.refresh_cpos(idaapi.USE_KEYBOARD):
                 #   print "refresh_cpos failed"
                    return 0
                pos = self.vu.cpos
                ypos = pos.lnnum
                #print "cursor click %d %d %d" % (pos.x, pos.y, pos.lnnum)

                if self.highlights.has_key(ypos):
                    return 0
                #print "5"

                ps = self.vu.cfunc.get_pseudocode()
                #print "6"
                #print "ypos:%d"%ypos
                #print "ps[ypos].line: %s"%(ps[ypos].line)

                #line = [idaapi.COLSTR("[%02d]"%i, chr(i)) for i in range(1,0x40) ]
                #ps[0].line = ''.join(line);
                #ps[1].line = '\x04'.join(line);
                #line = [idaapi.COLSTR(  idaapi.COLSTR("[  \x04%02d\x04  ]"%i, chr(i)), chr(i+1)) for i in range(1,0x40) ]
                #ps[2].line = ''.join(line);
                #ps[3].line = '\x04'.join(line);
                ln = ps[ypos].line[:]
                #print "ln"
                curline = idaapi.tag_remove(ln)
                #print "7"

                idxO = curline.find('{')
                idxC = curline.find('}')
                #print "O:", idxO, " C: ",idxC
                #there is no need to highlight first and last {
                #print "8"

                if (idxO >= 0) or (idxC >= 0):
                #   print "9"
                    self.clearall(ps, False)

                    self.highlights[ypos] = ps[ypos].bgcolor

                    ps[ypos].bgcolor = self.hicolor
                    
                    dir = 1
                    bracechar = '}'
                    idx = idxO

                    if (idxC >= 0):
                        dir = -1
                        bracechar = '{'
                        idx = idxC

                    j = ypos + dir

                    max = len(ps)
                 #   print "max: ",max

                    while (j>=0) and (j<max):
                #       print "10"
                        #print "j:", j
                        ln = idaapi.tag_remove(ps[j].line)
                        if ln.find(bracechar) == idx:
                            if not(self.highlights.has_key(j) ):
                                self.highlights[j] = ps[j].bgcolor
                            #ps[j].line = ps[j].line.replace(bracechar, idaapi.COLSTR("\x04"+bracechar+"\x04", "\x27"))
                            #ps[j].line = ps[j].line.replace(bracechar, idaapi.COLSTR(bracechar, chr(52)))
                            ps[j].bgcolor = self.hicolor
                            self.theotherline = j
                            break
                        j+=dir
                    
                    idaapi.refresh_idaview_anyway()
                else:
                    self.clearall(ps)
                #print "11"
                return 0
        except:
            traceback.print_exc()
        
        return 0

def remove():
    if hexlight_cb:
        idaapi.remove_hexrays_callback(hexlight_cb)

class HexHLightPlugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = "highlights the matching brace in Pseudocode-View"
    help = "press B to jump to the matching brace"
    wanted_name = "HexLight"
    wanted_hotkey = ""

    def init(self):
        # Some initialization
        global hexlight_cb_info, hexlight_cb

        if idaapi.init_hexrays_plugin():
            hexlight_cb_info = hexrays_callback_info()
            hexlight_cb = hexlight_cb_info.event_callback
            if not idaapi.install_hexrays_callback(hexlight_cb):
            #    print "could not install hexrays_callback"
                return idaapi.PLUGIN_SKIP
            print "Hexlight plugin installed"
            addon = idaapi.addon_info_t();
            addon.id = "milan.bohacek.hexlight";
            addon.name = "Hexlight";
            addon.producer = "Milan Bohacek";
            addon.url = "milan.bohacek+hexlight@gmail.com";
            addon.version = "6.95";
            idaapi.register_addon( addon );
            return idaapi.PLUGIN_KEEP
        #print "init_hexrays_plugin failed"
        return idaapi.PLUGIN_SKIP

    def run(self, arg=0):
        return

    def term(self):
        remove()

def PLUGIN_ENTRY():
    return HexHLightPlugin_t()


repository: o0Starley0o/HexRaysPlugins
file: XrefTable.py
wanted_name: XrefTable
comment: This is a comment
created_at: 2019-09-03 20:45:11+00:00
pushed_at: 2019-09-09 20:32:59+00:00
url:
https://github.com/o0Starley0o/HexRaysPlugins/blob/62aeb09a647d269ed01e1936d141dac2c1059124/XrefTable.py

content:
import sys
import os

current_path = os.path.dirname(__file__)

egg_loc = os.path.join(current_path, "pycharm-debug.egg")
sys.path.append(egg_loc)
print egg_loc
import pydevd



import PyQt5.QtGui as QtGui
import PyQt5.QtCore as QtCore
import PyQt5.QtWidgets as QtWidgets
import idaapi
import idautils
import idc
import json
import sip

with open(current_path + "/dict.json", "r") as read_file:
    DICT = json.load(read_file)
ID_NA = "."




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



class TakeAttributes():
    def __init__(self, ea, name, dst, risk):
        self.ea = ea
        self.name = name
        self.dst = dst
        self.risk = risk


# -----------------------------------------------------------------------------
class XrefTable(Choose2):
    def __init__(self, title):

        columns = [["Caller", 20 | 327680], ["Function", 30 | 327680], ["NumOfArgs", 8], ["Risk", 8]]
        # columns += [[i, 8] for i in ARG_NAMES]
        Choose2.__init__(self, title, columns)
        self.items = []
        twidget = idaapi.find_widget("Functions window")
        self.widget = sip.wrapinstance(long(twidget), QtWidgets.QWidget)  # NOTE: LOL
        self.table = self.widget.findChild(QtWidgets.QTableView)
        self.table.selectionModel().selectionChanged.connect(self.my_event)

    def OnClose(self):
        self.items = []

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

    def OnGetLine(self, n):
        return self._make_choser_entry(n)

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

    def feed(self, data, caller=None):
        for item in data:
            call = idc.get_func_off_str(item.ea)
            for temp in caller:
                if temp in call:
                    full_list = [item.ea, item.name, str(item.dst), str(item.risk)]
                    self.items.append(full_list)
                    break
        self.Refresh()
        return

    def my_event(self):
        func_list = []
        self.items = []
        for name in DICT:
            func_list += get_callers(str(name))
            func_list += get_callers("_" + str(name))

        func_list = set(func_list)

        CALLER_LIST = self.get_selected_funcs()
        for ea in func_list:
            try:
                cfunc = idaapi.decompile(ea)
            except idaapi.DecompilationFailure:
                print "Error decompiling function @ 0x%x" % ea
                cfunc = None

            if cfunc:
                fp = func_parser_t(cfunc)
                fp.apply_to(cfunc.body, None)
                self.feed(fp.data, CALLER_LIST)

    def _make_choser_entry(self, n):
        dst = []
        dst += ["%s" % idc.get_func_off_str(self.items[n][0])]
        dst += ["%s" % self.items[n][1]]
        for i in range(2, len(self.items[n])):
            dst += [self.items[n][i]]
        #print dst
        return dst

    def get_selected_funcs(self):
        selected_funcs = [str(s.data()) for s in self.table.selectionModel().selectedRows()]

        return selected_funcs


# -----------------------------------------------------------------------------
class func_parser_t(idaapi.ctree_visitor_t):
    def __init__(self, cfunc):
        idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST)
        self.cfunc = cfunc
        self.data = []
        return

    def _add_func_call(self, func_info):
        self.data.append(func_info)
        return

    def _parse_function(self, e):
        args = e.a  # carglist_t

        dst_name = idaapi.tag_remove(e.print1(None))

        numOfArgs = args.size()
        risk = ""
        for trisk in DICT[idaapi.tag_remove(e.x.print1(None))]:
            if str(numOfArgs) == trisk:
                risk = DICT[idaapi.tag_remove(e.x.print1(None))][trisk]
                break

        self._add_func_call(TakeAttributes(e.ea, dst_name, numOfArgs, risk))
        return True

    def visit_expr(self, e):
        op = e.op
        # if expression type is call
        if op == idaapi.cot_call:
            name = idaapi.tag_remove(e.x.print1(None))
            # print "name func = %s" % name
            # and if the function name is supported
            if name in DICT:
                # parse
                self._parse_function(e)
        return 0



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


def is_ida_version(requested):
    rv = requested.split(".")
    kv = idaapi.get_kernel_version().split(".")

    count = min(len(rv), len(kv))
    if not count:
        return False

    for i in xrange(count):
        if int(kv[i]) < int(rv[i]):
            return False
    return True


# -----------------------------------------------------------------------------
def get_callers(name):
    for xr in idautils.CodeRefsTo(idaapi.get_name_ea(idaapi.BADADDR, name), True):
        fn = idaapi.get_func(xr)
        if fn:
            yield fn.startEA


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

class Xrefs(idaapi.plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = "This is a comment"

    help = "This is help"
    wanted_name = "XrefTable"
    wanted_hotkey = ""#"Alt-F8"

    def init(self):
        #
        idaapi.load_plugin('hexrays')
        idaapi.msg("init() called!\n")


        return idaapi.PLUGIN_OK

    def run(self, arg):
        idaapi.msg("run() called with %d!\n" % arg)
        if not idaapi.init_hexrays_plugin():
            print "This script requires the HexRays decompiler plugin."
        else:

            #dict = {'myfunc':[{'args':'2', 'risk':'2 (medium)'}, {'args':'3', 'risk':'1 (low)'}], 'memset': {{'args':'2', 'risk':'2 (medium)'}}}
            Table = XrefTable("XrefTable")
            Table.Show()


    def term(self):
        return



def PLUGIN_ENTRY():
    return Xrefs()
'''
dict_template = {'dword ptr [esp]': [0x401500]}

for reg in dict_template.keys():
    for addr in dict_template.get(reg):
        print ("Addr log func:0x%08x\n" % addr)
        total = 0
        resolve = 0
        for xref in idautils.XrefsTo(addr, idaapi.XREF_ALL):
            print (xref.frm)
            if idaapi.isCode(GetFlags(xref.frm)):
                print ("cool")
                total +=1
                ea_call = xref.frm

                func_start = idc.GetFunctionAttr(ea_call, idc.FUNCATTR_START)

                if not func_start or func_start == idc.BADADDR:
                    print ("Func don't define: 0x%08x" % (ea_call))
                    continue

                ea = ea_call
                strFname = ""
                #Backtrace
                while ea != idc.BADADDR and ea != func_start:
                    ea = idc.PrevHead(ea, func_start)
                    print("ea = %x " % ea)
                    print (idc.GetMnem(ea))
                    print ("%s = %s" % (idc.GetOpType(ea, 0), idc.o_phrase))
                    print ("%s = %s" % (idc.GetOpnd(ea, 0), reg))
                    print ("%s = %s" % (idc.GetOpType(ea, 1), idc.o_imm))
                    print ("-----------------")

                    if idc.GetMnem(ea) == "mov" and idc.GetOpType(ea,0) ==idc.o_phrase \
                            and idc.GetOpnd(ea,0) == reg and idc.GetOpType(ea, 1) == idc.o_imm:
                        strFname = GetString(GetOperandValue(ea,1))
                        #print ("We found it = %s" % GetString(GetOperandValue(ea,1)))
                        #break
                        print("strFname = %s" % strFname)
                        if strFname is not None:

                            predName = GetFunctionName(func_start)
                            if predName.startswith('sub_'):
                                if idc.MakeName(func_start, strFname):
                                    print ("Addr: 0x%08x Before: %s After: %s" % (xref.frm,predName,strFname))
                                else:
                                    strFname = str(strFname) + "_"+ str(func_start)
                                    idc.MakeName(func_start,strFname)
                                    print ("Addr: 0x%08x Before: %s After: %s" % (xref.frm, predName, strFname))

                                resolve +=1
                                break
                            else:
                                print ("[Error] Conflict! current name: %s new_name: %s" %(predName,strFname))
                                break
                        else:
                            print ("[ERROR] Problem with string!")
                            break
        print ("%d/%d" % (resolve, total))'''


repository: mrexodia/ida-pro-mcp
file: src/ida_pro_mcp/mcp-plugin.py
description: MCP Server for IDA Pro
wanted_name: MCP
comment: MCP Plugin
created_at: 2025-03-25 22:23:41+00:00
pushed_at: 2025-05-09 22:50:39+00:00
url:
https://github.com/mrexodia/ida-pro-mcp/blob/44ebfe0f51960e97a2e061f382ec7bf53dbb2447/src/ida_pro_mcp/mcp-plugin.py

content:
import os
import sys

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

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

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

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

    def mark_unsafe(self, func: Callable) -> Callable:
        self.unsafe.add(func.__name__)
        return func

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

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

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

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

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

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

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

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

rpc_registry = RPCRegistry()

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

def unsafe(func: Callable) -> Callable:
    """Decorator to register mark a function as unsafe"""
    return rpc_registry.mark_unsafe(func)

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

    def do_POST(self):
        global rpc_registry

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

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

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

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

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

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

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

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

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

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

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

class Server:
    HOST = "localhost"
    PORT = 13337

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

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

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

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

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

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

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

import ida_hexrays
import ida_kernwin
import ida_funcs
import ida_gdl
import ida_lines
import ida_idaapi
import idc
import idaapi
import idautils
import ida_nalt
import ida_bytes
import ida_typeinf
import ida_xref
import ida_entry
import idautils
import ida_idd
import ida_dbg
import ida_name
import ida_ida

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

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

class IDASyncError(Exception):
    pass

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

logger = logging.getLogger(__name__)

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

call_stack = queue.LifoQueue()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

DEMANGLED_TO_EA = {}

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

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

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

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

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

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

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

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

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

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

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

T = TypeVar("T")

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

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

def pattern_filter(data: list[T], pattern: str, key: str) -> list[T]:
    def matches(item: T) -> bool:
        # TODO: proper pattern matching
        return pattern.lower() in item[key].lower()
    return list(filter(matches, data))

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

class Global(TypedDict):
    address: str
    name: str

@jsonrpc
@idaread
def list_globals(
    filter: Annotated[str, "Filter to apply to the list. Case-insensitive contains or /regex/ syntax"],
    offset: Annotated[int, "Offset to start listing from (start at 0)"],
    count: Annotated[int, "Number of globals to list (100 is a good default, 0 means remainder)"],
) -> Page[Global]:
    """List all globals in the database (paginated)"""
    globals = []
    for addr, name in idautils.Names():
        # Skip functions
        if not idaapi.get_func(addr):
            globals.append({
                "address": hex(addr),
                "name": name,
            })
    globals = pattern_filter(globals, filter, "name")
    return paginate(globals, offset, count)

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

@jsonrpc
@idaread
def list_strings(
    filter: Annotated[str, "Filter to apply to the list. Case-insensitive contains or /regex/ syntax"],
    offset: Annotated[int, "Offset to start listing from (start at 0)"],
    count: Annotated[int, "Number of strings to list (100 is a good default, 0 means remainder)"],
) -> Page[String]:
    """List all strings in the database (paginated)"""
    strings = []
    for item in idautils.Strings():
        try:
            string = str(item)
            if string:
                strings.append({
                    "address": hex(item.ea),
                    "length": item.length,
                    "string": string,
                })
        except:
            continue
    strings = pattern_filter(strings, filter, "string")
    return paginate(strings, offset, count)

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

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

    return pseudocode

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

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

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

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

@jsonrpc
@idaread
def get_xrefs_to_field(
    struct_name: Annotated[str, "Name of the struct (type) containing the field"],
    field_name: Annotated[str, "Name of the field (member) to get xrefs to"],
) -> list[Xref]:
    """Get all cross references to a named struct field (member)"""
    
    # Get the type library
    til = ida_typeinf.get_idati()
    if not til:
        raise IDAError("Failed to retrieve type library.")

    # Get the structure type info
    tif = ida_typeinf.tinfo_t()
    if not tif.get_named_type(til, struct_name, ida_typeinf.BTF_STRUCT, True, False):
        print(f"Structure '{struct_name}' not found.")
        return []

    # Get The field index
    idx = ida_typeinf.get_udm_by_fullname(None, struct_name + '.' + field_name)
    if idx == -1:
        print(f"Field '{field_name}' not found in structure '{struct_name}'.")
        return []

    # Get the type identifier
    tid = tif.get_udm_tid(idx)
    if tid == ida_idaapi.BADADDR:
        raise IDAError(f"Unable to get tid for structure '{struct_name}' and field '{field_name}'.")

    # Get xrefs to the tid
    xrefs = []
    xref: ida_xref.xrefblk_t
    for xref in idautils.XrefsTo(tid):
        xrefs.append({
            "address": hex(xref.frm),
            "type": "code" if xref.iscode else "data",
            "function": get_function(xref.frm, raise_error=False),
        })
    return xrefs

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

# NOTE: This is extremely hacky, but necessary to get errors out of IDA
def parse_decls_ctypes(decls: str, hti_flags: int) -> tuple[int, str]:
    if sys.platform == "win32":
        import ctypes

        assert isinstance(decls, str), "decls must be a string"
        assert isinstance(hti_flags, int), "hti_flags must be an int"
        c_decls = decls.encode("utf-8")
        c_til = None
        ida_dll = ctypes.CDLL("ida")
        ida_dll.parse_decls.argtypes = [
            ctypes.c_void_p,
            ctypes.c_char_p,
            ctypes.c_void_p,
            ctypes.c_int,
        ]
        ida_dll.parse_decls.restype = ctypes.c_int

        messages = []

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

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

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

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

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

@jsonrpc
@idaread
@unsafe
def dbg_get_registers() -> list[dict[str, str]]:
    """Get all registers and their values. This function is only available when debugging."""
    result = []
    dbg = ida_idd.get_dbg()
    # TODO: raise an exception when not debugging?
    for thread_index in range(ida_dbg.get_thread_qty()):
        tid = ida_dbg.getn_thread(thread_index)
        regs = []
        regvals = ida_dbg.get_reg_vals(tid)
        for reg_index, rv in enumerate(regvals):
            reg_info = dbg.regs(reg_index)
            reg_value = rv.pyval(reg_info.dtype)
            if isinstance(reg_value, int):
                reg_value = hex(reg_value)
            if isinstance(reg_value, bytes):
                reg_value = reg_value.hex(" ")
            regs.append({
                "name": reg_info.name,
                "value": reg_value,
            })
        result.append({
            "thread_id": tid,
            "registers": regs,
        })
    return result

@jsonrpc
@idaread
@unsafe
def dbg_get_call_stack() -> list[dict[str, str]]:
    """Get the current call stack."""
    callstack = []
    try:
        tid = ida_dbg.get_current_thread()
        trace = ida_idd.call_stack_t()

        if not ida_dbg.collect_stack_trace(tid, trace):
            return []
        for frame in trace:
            frame_info = {
                "address": hex(frame.callea),
            }
            try:
                module_info = ida_idd.modinfo_t()
                if ida_dbg.get_module_info(frame.callea, module_info):
                    frame_info["module"] = os.path.basename(module_info.name)
                else:
                    frame_info["module"] = "<unknown>"

                name = (
                    ida_name.get_nice_colored_name(
                        frame.callea,
                        ida_name.GNCN_NOCOLOR
                        | ida_name.GNCN_NOLABEL
                        | ida_name.GNCN_NOSEG
                        | ida_name.GNCN_PREFDBG,
                    )
                    or "<unnamed>"
                )
                frame_info["symbol"] = name

            except Exception as e:
                frame_info["module"] = "<error>"
                frame_info["symbol"] = str(e)

            callstack.append(frame_info)

    except Exception as e:
        pass
    return callstack

def list_breakpoints():
    ea = ida_ida.inf_get_min_ea()
    end_ea = ida_ida.inf_get_max_ea()
    breakpoints = []
    while ea <= end_ea:
        bpt = ida_dbg.bpt_t()
        if ida_dbg.get_bpt(ea, bpt):
            breakpoints.append(
                {
                    "ea": hex(bpt.ea),
                    "type": bpt.type,
                    "enabled": bpt.flags & ida_dbg.BPT_ENABLED,
                    "condition": bpt.condition if bpt.condition else None,
                }
            )
        ea = ida_bytes.next_head(ea, end_ea)
    return breakpoints

@jsonrpc
@idaread
@unsafe
def dbg_list_breakpoints():
    """
    List all breakpoints in the program.
    """
    return list_breakpoints()

@jsonrpc
@idaread
@unsafe
def dbg_start_process() -> str:
    """Start the debugger"""
    if idaapi.start_process("", "", ""):
        return "Debugger started"
    return "Failed to start debugger"

@jsonrpc
@idaread
@unsafe
def dbg_exit_process() -> str:
    """Exit the debugger"""
    if idaapi.exit_process():
        return "Debugger exited"
    return "Failed to exit debugger"

@jsonrpc
@idaread
@unsafe
def dbg_continue_process() -> str:
    """Continue the debugger"""
    if idaapi.continue_process():
        return "Debugger continued"
    return "Failed to continue debugger"

@jsonrpc
@idaread
@unsafe
def dbg_run_to(
    address: Annotated[str, "Run the debugger to the specified address"],
) -> str:
    """Run the debugger to the specified address"""
    ea = parse_address(address)
    if idaapi.run_to(ea):
        return f"Debugger run to {hex(ea)}"
    return f"Failed to run to address {hex(ea)}"

@jsonrpc
@idaread
@unsafe
def dbg_set_breakpoint(
    address: Annotated[str, "Set a breakpoint at the specified address"],
) -> str:
    """Set a breakpoint at the specified address"""
    ea = parse_address(address)
    if idaapi.add_bpt(ea, 0, idaapi.BPT_SOFT):
        return f"Breakpoint set at {hex(ea)}"
    breakpoints = list_breakpoints()
    for bpt in breakpoints:
        if bpt["ea"] == hex(ea):
            return f"Breakpoint already exists at {hex(ea)}"
    return f"Failed to set breakpoint at address {hex(ea)}"

@jsonrpc
@idaread
@unsafe
def dbg_delete_breakpoint(
    address: Annotated[str, "del a breakpoint at the specified address"],
) -> str:
    """del a breakpoint at the specified address"""
    ea = parse_address(address)
    if idaapi.del_bpt(ea):
        return f"Breakpoint deleted at {hex(ea)}"
    return f"Failed to delete breakpoint at address {hex(ea)}"

@jsonrpc
@idaread
@unsafe
def dbg_enable_breakpoint(
    address: Annotated[str, "Enable or disable a breakpoint at the specified address"],
    enable: Annotated[bool, "Enable or disable a breakpoint"],
) -> str:
    """Enable or disable a breakpoint at the specified address"""
    ea = parse_address(address)
    if idaapi.enable_bpt(ea, enable):
        return f"Breakpoint {'enabled' if enable else 'disabled'} at {hex(ea)}"
    return f"Failed to {'' if enable else 'disable '}breakpoint at address {hex(ea)}"

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

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

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

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

def PLUGIN_ENTRY():
    return MCP()


repository: alex-ilgayev/ida-pro-junk-code-removal
file: deobf_plugin.py
wanted_name: Deobfuscate Junk Code
created_at: 2020-02-27 14:32:50+00:00
pushed_at: 2020-02-27 17:16:43+00:00
url:
https://github.com/alex-ilgayev/ida-pro-junk-code-removal/blob/1405a902c692a25ff3bfee5a2995f1faaeb121f3/deobf_plugin.py

content:
import binascii
import hashlib

from idaapi import *
from idautils import *
from idc import *
from collections import deque

from unicorn_emulator.unicorn_emulator import unicorn_emu

def patch_nops(start_ea, end_ea):
    """patches nop instructions (0x90) on all bytes. not including last byte
    (WASN'T USED IN PLUGIN. KEPT FOR REFERENCING)
    """
    for ea in range(start_ea, end_ea):
        patch_byte(ea, 0x90)


def patch_jmp(src_ea, dst_ea):
    """patching src_ea address with a jmp instr into dst_ea
    (WASN'T USED IN PLUGIN. KEPT FOR REFERENCING)
    """
    jmp_inst = b'\xe9'
    jmp_inst += int.to_bytes(dst_ea - (src_ea + 5), 4, byteorder='little')
    
    for i in range(5):
        patch_byte(src_ea + i, jmp_inst[i])


def find_starting_bb(function, flowchart):
    """finds the first basic block of certain function.
    (WASN'T USED IN PLUGIN. KEPT FOR REFERENCING)
    """
    for bb in flowchart:
        if bb.start_ea == function.start_ea:
            return bb
    return None


def bbs_bfs(start_bb):
    """searches for all reachable basic blocks using BFS search 
    (WASN'T USED IN PLUGIN. KEPT FOR REFERENCING)
    """
    q = deque([start_bb])
    found_bbs = []
    while len(q) > 0:
        current_bb = q.popleft()
        #if current_bb.id not in found_bbs:
        if (current_bb.start_ea, current_bb.end_ea) not in [(bb.start_ea, bb.end_ea) for bb in found_bbs]:
            found_bbs.append(current_bb)
            for succ_bb in current_bb.succs():
                q.append(succ_bb)
    return found_bbs


def find_bbs_per_function(ea):
    """retrieves all basic blocks for a function.

    the return format is a list of tuples, where each tuple contains
    (start_ea, last_instruction_in_basic_block_ea)
    """
    function = get_func(ea)
    flowchart = FlowChart(function)

    return [(bb.start_ea, prev_head(bb.end_ea)) for bb in flowchart]
    

def find_potential_junk_codes(ea, exclusion_list):
    """finding potential blocks of code which could be junk code.
    receives exclusion list format (start_ea, end_ea) (exclusive end_ea)

    current heuristic is basic block which contains 5 unpopular assembly instructions.
    could be expanded to any other heuristic.

    return format is a list of the next tuple type: (start_ea, start_ea_of_next_bb, start_ea_of_curr_bb)
    """
    bad_mnem = []
    bad_mnem += ['shr', 'shl', 'aad', 'cbw', 'ror', 'xadd', 'bt', 'btc', 'bswap', 'aaa', 'aas', 'daa', 'clc']
    bad_mnem += ['bsf', 'cdq', 'cmc', 'cwd', 'pushf', 'popf', 'nop', 'stc', 'lahf', 'das', 'bsr']
    reachable = find_bbs_per_function(ea)
    print(f'[+] found reachable {len(reachable)} blocks')

    potent_bbs = []
    OCC_THRESHOLD = 4
    for bb in reachable:
        num_occ = 0
        start_addr = None
        addr = bb[0]
        while addr < bb[1]:
            has_been_removed = False
            for start, end in exclusion_list:
                if start <= addr < end:
                    addr = end
                    has_been_removed = True
                    break
            if has_been_removed:
                continue
            if print_insn_mnem(addr) in bad_mnem:
                num_occ += 1
                if start_addr is None:
                    start_addr = addr
            if num_occ >= OCC_THRESHOLD:
                potent_bbs.append((start_addr, bb[1], bb[0]))
                break
            addr = next_head(addr)
    return potent_bbs


def remove_junk_code(emu, start_ea, end_ea):
    """receives address, and attemps to remove junk code form that address and onwards.
    simulates the instructions on unicorn engine, and stops on first instruction which
    preserves the machine state as is.

    if that happens it calls to add_hidden_range()

    returns the actual range removed or None otherwise.
    """
    print(f'[+] attempting to remove at {hex(start_ea)}:{hex(end_ea)}')

    code = b''
    ea = start_ea
    while ea < end_ea:
        code += get_bytes(ea, get_item_size(ea))
        ea = next_head(ea)

    emu.load_code(code)
    emu.emulate()
    res = emu.find_matching_machine_state()
    if res is not None and res[0] == 0:
        res = None
    if res is not None:
        next_ea = res[0] + start_ea
        ea = prev_head(next_ea)
        
        flags = emu.get_flags_string()
        add_hidden_range(start_ea, next_ea, flags, "", "", 0x0000ff)
        request_refresh(0xffffffff)                    
        print(f'[+] removed {hex(start_ea)}:{hex(next_ea)}')
        return start_ea, next_ea

    print(f'[-] the code is not junk')
    return None


def remove_junk_per_func(emu, ea):
    """main function

    attemps to remove all junk codes of functoin ea belongs to.
    function fails if ea doesn't belongs to a function.

    - look for potential code blocks
    - emulates it in unicorn to check whether it's junk code
    - does the same procedure MAX_TRIES times when going one instruction back (in case heuristics failed)
    - checks whether any more potentials blocks left
    - rinse and repeat untill no blocks left in function.
    """
    if get_func(ea) == None:
        print(f'[-] address given - {hex(ea)} does not belong to a function')
        return False

    MAX_TRIES = 4 # we receiving heuristic starting ea, and trying to go back.
    removal_list = [] # contains all removed code ranges

    print(f'[+] attempting to remove junk code for function {get_func_name(ea)}')
    
    potential_addrs = find_potential_junk_codes(ea, removal_list)
    while len(potential_addrs) > 0:
        print(f'[+] found potential {len(potential_addrs)} blocks')

        counter = 0
        for start_ea, end_ea, bb_start in potential_addrs:
            for _ in range(MAX_TRIES):
                res_range = remove_junk_code(emu, start_ea, end_ea)
                if res_range is not None:
                    removal_list.append(res_range)
                    break
                if start_ea == bb_start:
                    break
                start_ea = prev_head(start_ea)
        potential_addrs = find_potential_junk_codes(ea, removal_list)
    
    print(f'[+] successfully removed {len(removal_list)} blocks')

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

VERSION = "v1.0"
AUTHORS = ['Alex Ilgayev']

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

class deobf_junk_t(plugin_t):

    flags = PLUGIN_FIX
    comment = ""
    help = ""
    wanted_name = "Deobfuscate Junk Code"
    wanted_hotkey = "Shift+D"

    def __init__(self):
        self._emu = unicorn_emu()

    def init(self):
        return PLUGIN_OK

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        ea = here()
        remove_junk_per_func(self._emu, ea)

    def term(self):
        pass


    



repository: BehroozAbbassi/hyperv-research-scripts
file: scripts/IA32-VMX-Helper/IA32_VMX_Helper.py
description: IDA scripts for hypervisor (Hyper-v) analysis and reverse engineering automation
created_at: 2021-12-07 20:03:28+00:00
pushed_at: 2021-12-07 20:03:37+00:00
url:
https://github.com/BehroozAbbassi/hyperv-research-scripts/blob/049b9b749cfec704b56d8d680c5886c1bd657fb1/scripts/IA32-VMX-Helper/IA32_VMX_Helper.py

content:
__author__  = "Behrooz Abbassi @BehroozAbbassi"
__license__ = "MIT"
__version__ = "2.0.0"


"""
You can use this script to find and decode all (I hope!) MSRs and VMCSs in hypervisor binaries,
it also has a siple gui that enables you to search in founded values by Adddress, Name and Const values
In addition to searching and decoding, you can set proper comments or symbolic constant on these values
automatically
"""


import collections
import os

# IDAPython 
import idaapi
import idautils
import idc
import struct
# ------------------------------------------------------------------------------------- #
"""
If you enable this mode, the script will try to find and decode more MSRs and VMCS by appling more
analysis on the binary, BUT because of linear analysis there is some false positives, which can cause 
some problems :-( so this mode is disabled by default


Examples

mov     rbx, 681Ch ; VMCS_GUEST_RSP
vmwrite rbx, rsi

This is the basic form of calling vmwrite this one is easy to figure but what about next examples?

mov     rax, 681Ch ; VMCS_GUEST_RSP
mov     rbx, rax 
vmwrite rbx, rsi


mov     eax, 681Ch ; VMCS_GUEST_RSP
call    Function_0()
mov     ecx, eax 
vmwrite ecx, rsi

"""
cfg_enable_smart_mode        = False
cfg_smart_mode_max_back_step = 10

# ------------------------------------------------------------------------------------- #
g_knowledge_db = None
# ------------------------------------------------------------------------------------- #
def AddSymbolicConstantsEnumsToIda():
    # ------------------------------------------------------------------------------------- #
    # Based on @AmarSaar post [https://msrc-blog.microsoft.com/2018/12/10/first-steps-in-hyper-v-research]
    # Notice these values can change between different builds!
    # ------------------------------------------------------------------------------------- #
    id = idc.AddEnum(0, "HYPERV_STRUCTS", idaapi.hexflag())
    idc.AddConstEx(id, "Self", 0, -1)
    idc.AddConstEx(id, "CpuIndex", 8, -1)
    idc.AddConstEx(id, "CurrentThread", 0x38, -1)
    idc.AddConstEx(id, "LogicalCpuIndex", 0x8080, -1)
    idc.AddConstEx(id, "CurrentVP", 0x82E0, -1)
    idc.AddConstEx(id, "IDT", 0x8340, -1)
    idc.AddConstEx(id, "GDT", 0x8348, -1)
    idc.AddConstEx(id, "CurrentPartition", 0x83A8, -1)
    idc.AddConstEx(id, "CurrentVMCS", 0x186E0, -1)



    #https://www.geoffchappell.com/studies/windows/km/ntoskrnl/api/hvilib/hviintel/cpuid_result.htm
    id = idc.AddEnum(0, "HV_CPUID_FUNCTION_ENUM", idaapi.hexflag())
    idc.AddConstEx(id, "HvCpuIdFunctionVersionAndFeatures", 0x00000001, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionHvVendorAndMaxFunction", 0x40000000, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionHvInterface", 0x40000001, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionMsHvVersion", 0x40000002, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionMsHvFeatures", 0x40000003, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionMsHvEnlightenmentInformation",0x40000004, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionMsHvImplementationLimits", 0x40000005, -1)

    idc.AddConstEx(id, "HvCpuIdFunctionMsHvHardwareFeatures", 0x40000006, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionMsHvCpuManagementFeatures", 0x40000007, -1)
    idc.AddConstEx(id, "HvCpuIdFunctionMsHvSvmFeatures", 0x40000008, -1)
    #https://patchwork.kernel.org/patch/10649047/
    idc.AddConstEx(id, "HvCpuIdFunctionNestedFeatures", 0x4000000A, -1)




    # ------------------------------------------------------------------------------------- #
    # Exit Reasons
    # ------------------------------------------------------------------------------------- #
    id = idc.AddEnum(0, "IA32_VMX_EXIT_REASONS_ENUM", idaapi.hexflag())
    idc.AddConstEx(id, "EXIT_REASON_EXCEPTION_NMI", 0x0, -1)
    idc.AddConstEx(id, "EXIT_REASON_EXTERNAL_INTERRUPT", 0x1, -1)
    idc.AddConstEx(id, "EXIT_REASON_TRIPLE_FAULT", 0x2, -1)
    idc.AddConstEx(id, "EXIT_REASON_INIT", 0x3, -1)
    idc.AddConstEx(id, "EXIT_REASON_SIPI", 0x4, -1)
    idc.AddConstEx(id, "EXIT_REASON_IO_SMI", 0x5, -1)
    idc.AddConstEx(id, "EXIT_REASON_OTHER_SMI", 0x6, -1)
    idc.AddConstEx(id, "EXIT_REASON_PENDING_VIRT_INTR", 0x7, -1)
    idc.AddConstEx(id, "EXIT_REASON_PENDING_VIRT_NMI",  0x8, -1)
    idc.AddConstEx(id, "EXIT_REASON_TASK_SWITCH", 0x9, -1)
    idc.AddConstEx(id, "EXIT_REASON_CPUID", 0x0A, -1)
    idc.AddConstEx(id, "EXIT_REASON_GETSEC", 0x0B, -1)
    idc.AddConstEx(id, "EXIT_REASON_HLT", 0x0C, -1)
    idc.AddConstEx(id, "EXIT_REASON_INVD", 0x0D, -1)
    idc.AddConstEx(id, "EXIT_REASON_INVLPG", 0x0E, -1)
    idc.AddConstEx(id, "EXIT_REASON_RDPMC", 0x0F, -1)
    idc.AddConstEx(id, "EXIT_REASON_RDTSC", 0x10, -1)
    idc.AddConstEx(id, "EXIT_REASON_RSM", 0x11, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMCALL", 0x12, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMCLEAR", 0x13, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMLAUNCH", 0x14, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMPTRLD", 0x15, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMPTRST", 0x16, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMREAD", 0x17, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMRESUME", 0x18, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMWRITE", 0x19, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMXOFF", 0x1A, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMXON", 0x1B, -1)
    idc.AddConstEx(id, "EXIT_REASON_CR_ACCESS", 0x1C, -1)
    idc.AddConstEx(id, "EXIT_REASON_DR_ACCESS", 0x1D, -1)
    idc.AddConstEx(id, "EXIT_REASON_IO_INSTRUCTION", 0x1E, -1)
    idc.AddConstEx(id, "EXIT_REASON_MSR_READ", 0x1F, -1)
    idc.AddConstEx(id, "EXIT_REASON_MSR_WRITE", 0x20, -1)
    idc.AddConstEx(id, "EXIT_REASON_INVALID_GUEST_STATE", 0X21, -1)
    idc.AddConstEx(id, "EXIT_REASON_MSR_LOADING", 0x22, -1)

    idc.AddConstEx(id, "EXIT_REASON_MWAIT_INSTRUCTION", 0x24, -1)
    idc.AddConstEx(id, "EXIT_REASON_MONITOR_TRAP_FLAG", 0x25, -1)

    idc.AddConstEx(id, "EXIT_REASON_MONITOR_INSTRUCTION", 0x27, -1)
    idc.AddConstEx(id, "EXIT_REASON_PAUSE_INSTRUCTION", 0x28, -1)
    idc.AddConstEx(id, "EXIT_REASON_MCE_DURING_VMENTRY", 0x29, -1)

    idc.AddConstEx(id, "EXIT_REASON_TPR_BELOW_THRESHOLD", 0x2B, -1)
    idc.AddConstEx(id, "EXIT_REASON_APIC_ACCESS", 0x2C, -1)
    idc.AddConstEx(id, "EXIT_REASON_VIRTUALIZED_EOI", 0x2D, -1)
    idc.AddConstEx(id, "EXIT_REASON_ACCESS_GDTR_OR_IDTR", 0x2E, -1)
    idc.AddConstEx(id, "EXIT_REASON_ACCESS_LDTR_OR_TR", 0x2F, -1)
    idc.AddConstEx(id, "EXIT_REASON_EPT_VIOLATION", 0x30, -1)
    idc.AddConstEx(id, "EXIT_REASON_EPT_MISCONFIG", 0x31, -1)
    idc.AddConstEx(id, "EXIT_REASON_INVEPT", 0x32, -1)
    idc.AddConstEx(id, "EXIT_REASON_RDTSCP", 0x33, -1)
    idc.AddConstEx(id, "EXIT_REASON_VMX_PREEMPTION_TIMER_EXPIRED", 0x34, -1)
    idc.AddConstEx(id, "EXIT_REASON_INVVPID", 0x35, -1)
    idc.AddConstEx(id, "EXIT_REASON_WBINVD", 0x36, -1)
    idc.AddConstEx(id, "EXIT_REASON_XSETBV", 0x37, -1)
    idc.AddConstEx(id, "EXIT_REASON_APIC_WRITE", 0x38, -1)
    idc.AddConstEx(id, "EXIT_REASON_RDRAND", 0x39, -1)
    idc.AddConstEx(id, "EXIT_REASON_INVPCID", 0x3A, -1)
    idc.AddConstEx(id, "EXIT_REASON_EXECUTE_VMFUNC", 0x3B, -1)
    idc.AddConstEx(id, "EXIT_REASON_EXECUTE_ENCLS", 0x3C, -1)
    idc.AddConstEx(id, "EXIT_REASON_RDSEED", 0x3D, -1)
    idc.AddConstEx(id, "EXIT_REASON_PML_FULL", 0x3E, -1)
    idc.AddConstEx(id, "EXIT_REASON_XSAVES", 0x3F, -1)
    idc.AddConstEx(id, "EXIT_REASON_XRSTORS", 0x40, -1)
    idc.AddConstEx(id, "EXIT_REASON_PCOMMIT", 0x41, -1)

    # ------------------------------------------------------------------------------------- #
    # MSR
    # ------------------------------------------------------------------------------------- #
    id = idc.AddEnum(0, "IA32_MSR_LIST_ENUM", idaapi.hexflag())
    idc.AddConstEx(id, "HV_X64_MSR_GUEST_OS_ID", 0x40000000, -1)
    idc.AddConstEx(id, "HV_X64_MSR_HYPERCALL", 0x40000001, -1)
    idc.AddConstEx(id, "HV_X64_MSR_VP_INDEX", 0x40000002, -1)
    idc.AddConstEx(id, "HV_X64_MSR_RESET", 0x40000003, -1)
    idc.AddConstEx(id, "HV_X64_MSR_VP_RUNTIME", 0x40000010, -1)
    idc.AddConstEx(id, "HV_X64_MSR_TIME_REF_COUNT", 0x40000020, -1)
    idc.AddConstEx(id, "HV_X64_MSR_REFERENCE_TSC", 0x40000021, -1)
    idc.AddConstEx(id, "HV_X64_MSR_TSC_FREQUENCY", 0x40000022, -1)
    idc.AddConstEx(id, "HV_X64_MSR_APIC_FREQUENCY", 0x40000023, -1)
    idc.AddConstEx(id, "HV_X64_MSR_EOI", 0x40000070, -1)
    idc.AddConstEx(id, "HV_X64_MSR_ICR", 0x40000071, -1)
    idc.AddConstEx(id, "HV_X64_MSR_TPR", 0x40000072, -1)
    idc.AddConstEx(id, "HV_X64_MSR_VP_ASSIST_PAGE", 0x40000073, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SCONTROL", 0x40000080, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SVERSION", 0x40000081, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SIEFP", 0x40000082, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SIMP", 0x40000083, -1)
    idc.AddConstEx(id, "HV_X64_MSR_EOM", 0x40000084, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT0", 0x40000090, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT1", 0x40000091, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT2", 0x40000092, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT3", 0x40000093, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT4", 0x40000094, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT5", 0x40000095, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT6", 0x40000096, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT7", 0x40000097, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT8", 0x40000098, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT9", 0x40000099, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT10", 0x4000009a, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT11", 0x4000009b, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT12", 0x4000009c, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT13", 0x4000009d, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT14", 0x4000009e, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SINT15", 0x4000009f, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER0_CONFIG", 0x400000b0, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER0_COUNT", 0x400000b1, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER1_CONFIG", 0x400000b2, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER1_COUNT", 0x400000b3, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER2_CONFIG", 0x400000b4, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER2_COUNT", 0x400000b5, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER3_CONFIG", 0x400000b6, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STIMER3_COUNT", 0x400000b7, -1)
    idc.AddConstEx(id, "HV_X64_MSR_POWER_STATE_TRIGGER_C1", 0x400000c1, -1)
    idc.AddConstEx(id, "HV_X64_MSR_POWER_STATE_TRIGGER_C2", 0x400000c2, -1)
    idc.AddConstEx(id, "HV_X64_MSR_POWER_STATE_TRIGGER_C3", 0x400000c3, -1)
    idc.AddConstEx(id, "HV_X64_MSR_POWER_STATE_CONFIG_C1", 0x400000d1, -1)
    idc.AddConstEx(id, "HV_X64_MSR_POWER_STATE_CONFIG_C2", 0x400000d2, -1)
    idc.AddConstEx(id, "HV_X64_MSR_POWER_STATE_CONFIG_C3", 0x400000d3, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STATS_PARTITION_RETAIL_PAGE", 0x400000e0, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STATS_PARTITION_INTERNAL_PAGE", 0x400000e1, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STATS_VP_RETAIL_PAGE", 0x400000e2, -1)
    idc.AddConstEx(id, "HV_X64_MSR_STATS_VP_INTERNAL_PAGE", 0x400000e3, -1)
    idc.AddConstEx(id, "HV_X64_MSR_GUEST_IDLE", 0x400000f0, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SYNTH_DEBUG_CONTROL", 0x400000f1, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SYNTH_DEBUG_STATUS", 0x400000f2, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SYNTH_DEBUG_SEND_BUFFER", 0x400000f3, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SYNTH_DEBUG_RECEIVE_BUFFER", 0x400000f4, -1)
    idc.AddConstEx(id, "HV_X64_MSR_SYNTH_DEBUG_PENDING_BUFFER", 0x400000f5, -1)
    idc.AddConstEx(id, "HV_X64_MSR_CRASH_P0", 0x40000100, -1)
    idc.AddConstEx(id, "HV_X64_MSR_CRASH_P1", 0x40000101, -1)
    idc.AddConstEx(id, "HV_X64_MSR_CRASH_P2", 0x40000102, -1)
    idc.AddConstEx(id, "HV_X64_MSR_CRASH_P3", 0x40000103, -1)
    idc.AddConstEx(id, "HV_X64_MSR_CRASH_P4", 0x40000104, -1)
    idc.AddConstEx(id, "HV_X64_MSR_CRASH_CTL", 0x40000105, -1)
    idc.AddConstEx(id, "HV_X64_MSR_REENLIGHTENMENT_CONTROL", 0x40000106, -1)
    idc.AddConstEx(id, "HV_X64_MSR_TSC_EMULATION_CONTROL", 0x40000107, -1)
    idc.AddConstEx(id, "HV_X64_MSR_TSC_EMULATION_STATUS", 0x40000108, -1)
    idc.AddConstEx(id, "MSR_XEON_D_PPIN_CTL", 0x4e, -1)
    idc.AddConstEx(id, "MSR_XEON_D_PPIN", 0x4f, -1)
    idc.AddConstEx(id, "MSR_XEON_D_PLATFORM_INFO", 0xce, -1)
    idc.AddConstEx(id, "MSR_XEON_D_PKG_CST_CONFIG_CONTROL", 0xe2, -1)
    idc.AddConstEx(id, "IA32_MCG_CAP", 0x179, -1)
    idc.AddConstEx(id, "MSR_XEON_D_SMM_MCA_CAP", 0x17d, -1)
    idc.AddConstEx(id, "MSR_XEON_D_TEMPERATURE_TARGET", 0x1a2, -1)
    idc.AddConstEx(id, "MSR_XEON_D_TURBO_RATIO_LIMIT", 0x1ad, -1)
    idc.AddConstEx(id, "MSR_XEON_D_TURBO_RATIO_LIMIT1", 0x1ae, -1)
    idc.AddConstEx(id, "MSR_XEON_D_RAPL_POWER_UNIT", 0x606, -1)
    idc.AddConstEx(id, "MSR_XEON_D_DRAM_POWER_LIMIT", 0x618, -1)
    idc.AddConstEx(id, "MSR_XEON_D_DRAM_ENERGY_STATUS", 0x619, -1)
    idc.AddConstEx(id, "MSR_XEON_D_DRAM_PERF_STATUS", 0x61b, -1)
    idc.AddConstEx(id, "MSR_XEON_D_DRAM_POWER_INFO", 0x61c, -1)
    idc.AddConstEx(id, "MSR_XEON_D_MSRUNCORE_RATIO_LIMIT", 0x620, -1)
    idc.AddConstEx(id, "MSR_XEON_D_PP0_ENERGY_STATUS", 0x639, -1)
    idc.AddConstEx(id, "MSR_XEON_D_CORE_PERF_LIMIT_REASONS", 0x690, -1)
    idc.AddConstEx(id, "IA32_QM_EVTSEL", 0xc8d, -1)
    idc.AddConstEx(id, "IA32_PQR_ASSOC", 0xc8f, -1)
    idc.AddConstEx(id, "MSR_XEON_D_TURBO_RATIO_LIMIT3", 0x1ac, -1)
    idc.AddConstEx(id, "IA32_L3_QOS_CFG", 0xc81, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_0", 0xc90, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_1", 0xc91, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_2", 0xc92, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_3", 0xc93, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_4", 0xc94, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_5", 0xc95, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_6", 0xc96, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_7", 0xc97, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_8", 0xc98, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_9", 0xc99, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_10", 0xc9a, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_11", 0xc9b, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_12", 0xc9c, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_13", 0xc9d, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_14", 0xc9e, -1)
    idc.AddConstEx(id, "MSR_XEON_D_IA32_L3_QOS_MASK_15", 0xc9f, -1)
    idc.AddConstEx(id, "MSR_SPEC_CTRL", 0x48, -1)
    idc.AddConstEx(id, "MSR_PRED_CMD", 0x49, -1)
    idc.AddConstEx(id, "MSR_LBR_SELECT", 0x1c8, -1)
    idc.AddConstEx(id, "MSR_LBR_TOS", 0x1c9, -1)
    idc.AddConstEx(id, "MSR_LBR_NHM_FROM", 0x680, -1)
    idc.AddConstEx(id, "MSR_LBR_NHM_TO", 0x6c0, -1)
    idc.AddConstEx(id, "MSR_LBR_CORE_FROM", 0x40, -1)
    idc.AddConstEx(id, "MSR_LBR_CORE_TO", 0x60, -1)
    idc.AddConstEx(id, "MSR_IA32_LASTBRANCHFROMIP", 0x1db, -1)
    idc.AddConstEx(id, "MSR_IA32_LASTBRANCHTOIP", 0x1dc, -1)
    idc.AddConstEx(id, "MSR_IA32_LASTINTFROMIP", 0x1dd, -1)
    idc.AddConstEx(id, "MSR_IA32_LASTINTTOIP", 0x1de, -1)
    idc.AddConstEx(id, "MSR_PKG_C3_RESIDENCY", 0x3f8, -1)
    idc.AddConstEx(id, "MSR_PKG_C6_RESIDENCY", 0x3f9, -1)
    idc.AddConstEx(id, "MSR_PKG_C7_RESIDENCY", 0x3fa, -1)
    idc.AddConstEx(id, "MSR_CORE_C3_RESIDENCY", 0x3fc, -1)
    idc.AddConstEx(id, "MSR_CORE_C6_RESIDENCY", 0x3fd, -1)
    idc.AddConstEx(id, "MSR_CORE_C7_RESIDENCY", 0x3fe, -1)
    idc.AddConstEx(id, "MSR_KNL_CORE_C6_RESIDENCY", 0x3ff, -1)
    idc.AddConstEx(id, "MSR_PKG_C2_RESIDENCY", 0x60d, -1)
    idc.AddConstEx(id, "MSR_PKG_C8_RESIDENCY", 0x630, -1)
    idc.AddConstEx(id, "MSR_PKG_C9_RESIDENCY", 0x631, -1)
    idc.AddConstEx(id, "MSR_PKG_C10_RESIDENCY", 0x632, -1)
    idc.AddConstEx(id, "MSR_IDT_FCR1", 0x107, -1)
    idc.AddConstEx(id, "MSR_IDT_FCR2", 0x108, -1)
    idc.AddConstEx(id, "MSR_IDT_FCR3", 0x109, -1)
    idc.AddConstEx(id, "MSR_IDT_FCR4", 0x10a, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR0", 0x110, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR1", 0x111, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR2", 0x112, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR3", 0x113, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR4", 0x114, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR5", 0x115, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR6", 0x116, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR7", 0x117, -1)
    idc.AddConstEx(id, "MSR_IDT_MCR_CTRL", 0x120, -1)
    idc.AddConstEx(id, "MSR_FSB_FREQ", 0xcd, -1)
    idc.AddConstEx(id, "MSR_IA32_BBL_CR_CTL", 0x119, -1)
    idc.AddConstEx(id, "MSR_VM_CR", 0xc0010114, -1)
    idc.AddConstEx(id, "MSR_VM_IGNNE", 0xc0010115, -1)
    idc.AddConstEx(id, "MSR_VM_HSAVE_PA", 0xc0010117, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX64K_BASE", 0x0, -1)
    idc.AddConstEx(id, "IA32_P5_MC_TYPE", 0x1, -1)
    idc.AddConstEx(id, "IA32_MONITOR_FILTER_SIZE", 0x6, -1)
    idc.AddConstEx(id, "IA32_TIME_STAMP_COUNTER", 0x10, -1)
    idc.AddConstEx(id, "IA32_PLATFORM_ID", 0x17, -1)
    idc.AddConstEx(id, "IA32_APIC_BASE", 0x1b, -1)
    idc.AddConstEx(id, "IA32_FEATURE_CONTROL", 0x3a, -1)
    idc.AddConstEx(id, "IA32_TSC_ADJUST", 0x3b, -1)
    idc.AddConstEx(id, "IA32_BIOS_UPDT_TRIG", 0x79, -1)
    idc.AddConstEx(id, "IA32_BIOS_SIGN_ID", 0x8b, -1)
    idc.AddConstEx(id, "IA32_SGXLEPUBKEYHASH0", 0x8c, -1)
    idc.AddConstEx(id, "IA32_SGXLEPUBKEYHASH1", 0x8d, -1)
    idc.AddConstEx(id, "IA32_SGXLEPUBKEYHASH2", 0x8e, -1)
    idc.AddConstEx(id, "IA32_SGXLEPUBKEYHASH3", 0x8f, -1)
    idc.AddConstEx(id, "IA32_SMM_MONITOR_CTL", 0x9b, -1)
    idc.AddConstEx(id, "IA32_SMBASE", 0x9e, -1)
    idc.AddConstEx(id, "IA32_PMC0", 0xc1, -1)
    idc.AddConstEx(id, "IA32_PMC1", 0xc2, -1)
    idc.AddConstEx(id, "IA32_PMC2", 0xc3, -1)
    idc.AddConstEx(id, "IA32_PMC3", 0xc4, -1)
    idc.AddConstEx(id, "IA32_PMC4", 0xc5, -1)
    idc.AddConstEx(id, "IA32_PMC5", 0xc6, -1)
    idc.AddConstEx(id, "IA32_PMC6", 0xc7, -1)
    idc.AddConstEx(id, "IA32_PMC7", 0xc8, -1)
    idc.AddConstEx(id, "IA32_MPERF", 0xe7, -1)
    idc.AddConstEx(id, "IA32_APERF", 0xe8, -1)
    idc.AddConstEx(id, "IA32_MTRRCAP", 0xfe, -1)
    idc.AddConstEx(id, "IA32_SYSENTER_CS", 0x174, -1)
    idc.AddConstEx(id, "IA32_SYSENTER_ESP", 0x175, -1)
    idc.AddConstEx(id, "IA32_SYSENTER_EIP", 0x176, -1)
    idc.AddConstEx(id, "IA32_MCG_STATUS", 0x17a, -1)
    idc.AddConstEx(id, "IA32_MCG_CTL", 0x17b, -1)
    idc.AddConstEx(id, "IA32_PERFEVTSEL0", 0x186, -1)
    idc.AddConstEx(id, "IA32_PERFEVTSEL1", 0x187, -1)
    idc.AddConstEx(id, "IA32_PERFEVTSEL2", 0x188, -1)
    idc.AddConstEx(id, "IA32_PERFEVTSEL3", 0x189, -1)
    idc.AddConstEx(id, "IA32_PERF_STATUS", 0x198, -1)
    idc.AddConstEx(id, "IA32_PERF_CTL", 0x199, -1)
    idc.AddConstEx(id, "IA32_CLOCK_MODULATION", 0x19a, -1)
    idc.AddConstEx(id, "IA32_THERM_INTERRUPT", 0x19b, -1)
    idc.AddConstEx(id, "IA32_THERM_STATUS", 0x19c, -1)
    idc.AddConstEx(id, "IA32_MISC_ENABLE", 0x1a0, -1)
    idc.AddConstEx(id, "IA32_ENERGY_PERF_BIAS", 0x1b0, -1)
    idc.AddConstEx(id, "IA32_PACKAGE_THERM_STATUS", 0x1b1, -1)
    idc.AddConstEx(id, "IA32_PACKAGE_THERM_INTERRUPT", 0x1b2, -1)
    idc.AddConstEx(id, "IA32_DEBUGCTL", 0x1d9, -1)
    idc.AddConstEx(id, "IA32_SMRR_PHYSBASE", 0x1f2, -1)
    idc.AddConstEx(id, "IA32_SMRR_PHYSMASK", 0x1f3, -1)
    idc.AddConstEx(id, "IA32_PLATFORM_DCA_CAP", 0x1f8, -1)
    idc.AddConstEx(id, "IA32_CPU_DCA_CAP", 0x1f9, -1)
    idc.AddConstEx(id, "IA32_DCA_0_CAP", 0x1fa, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE0", 0x200, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE1", 0x202, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE2", 0x204, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE3", 0x206, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE4", 0x208, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE5", 0x20a, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE6", 0x20c, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE7", 0x20e, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE8", 0x210, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSBASE9", 0x212, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK0", 0x201, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK1", 0x203, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK2", 0x205, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK3", 0x207, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK4", 0x209, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK5", 0x20b, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK6", 0x20d, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK7", 0x20f, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK8", 0x211, -1)
    idc.AddConstEx(id, "IA32_MTRR_PHYSMASK9", 0x213, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX64K_SIZE", 0x10000, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX64K_00000", 0x250, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX16K_BASE", 0x80000, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX16K_SIZE", 0x4000, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX16K_80000", 0x258, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX16K_A0000", 0x259, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_BASE", 0xc0000, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_SIZE", 0x1000, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_C0000", 0x268, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_C8000", 0x269, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_D0000", 0x26a, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_D8000", 0x26b, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_E0000", 0x26c, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_E8000", 0x26d, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_F0000", 0x26e, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX4K_F8000", 0x26f, -1)
    idc.AddConstEx(id, "IA32_MTRR_FIX_COUNT", 0x58, -1)
    idc.AddConstEx(id, "IA32_MTRR_VARIABLE_COUNT", 0xff, -1)
    idc.AddConstEx(id, "IA32_MTRR_COUNT", 0x157, -1)
    idc.AddConstEx(id, "IA32_PAT", 0x277, -1)
    idc.AddConstEx(id, "IA32_MC0_CTL2", 0x280, -1)
    idc.AddConstEx(id, "IA32_MC1_CTL2", 0x281, -1)
    idc.AddConstEx(id, "IA32_MC2_CTL2", 0x282, -1)
    idc.AddConstEx(id, "IA32_MC3_CTL2", 0x283, -1)
    idc.AddConstEx(id, "IA32_MC4_CTL2", 0x284, -1)
    idc.AddConstEx(id, "IA32_MC5_CTL2", 0x285, -1)
    idc.AddConstEx(id, "IA32_MC6_CTL2", 0x286, -1)
    idc.AddConstEx(id, "IA32_MC7_CTL2", 0x287, -1)
    idc.AddConstEx(id, "IA32_MC8_CTL2", 0x288, -1)
    idc.AddConstEx(id, "IA32_MC9_CTL2", 0x289, -1)
    idc.AddConstEx(id, "IA32_MC10_CTL2", 0x28a, -1)
    idc.AddConstEx(id, "IA32_MC11_CTL2", 0x28b, -1)
    idc.AddConstEx(id, "IA32_MC12_CTL2", 0x28c, -1)
    idc.AddConstEx(id, "IA32_MC13_CTL2", 0x28d, -1)
    idc.AddConstEx(id, "IA32_MC14_CTL2", 0x28e, -1)
    idc.AddConstEx(id, "IA32_MC15_CTL2", 0x28f, -1)
    idc.AddConstEx(id, "IA32_MC16_CTL2", 0x290, -1)
    idc.AddConstEx(id, "IA32_MC17_CTL2", 0x291, -1)
    idc.AddConstEx(id, "IA32_MC18_CTL2", 0x292, -1)
    idc.AddConstEx(id, "IA32_MC19_CTL2", 0x293, -1)
    idc.AddConstEx(id, "IA32_MC20_CTL2", 0x294, -1)
    idc.AddConstEx(id, "IA32_MC21_CTL2", 0x295, -1)
    idc.AddConstEx(id, "IA32_MC22_CTL2", 0x296, -1)
    idc.AddConstEx(id, "IA32_MC23_CTL2", 0x297, -1)
    idc.AddConstEx(id, "IA32_MC24_CTL2", 0x298, -1)
    idc.AddConstEx(id, "IA32_MC25_CTL2", 0x299, -1)
    idc.AddConstEx(id, "IA32_MC26_CTL2", 0x29a, -1)
    idc.AddConstEx(id, "IA32_MC27_CTL2", 0x29b, -1)
    idc.AddConstEx(id, "IA32_MC28_CTL2", 0x29c, -1)
    idc.AddConstEx(id, "IA32_MC29_CTL2", 0x29d, -1)
    idc.AddConstEx(id, "IA32_MC30_CTL2", 0x29e, -1)
    idc.AddConstEx(id, "IA32_MC31_CTL2", 0x29f, -1)
    idc.AddConstEx(id, "IA32_MTRR_DEF_TYPE", 0x2ff, -1)
    idc.AddConstEx(id, "IA32_FIXED_CTR0", 0x309, -1)
    idc.AddConstEx(id, "IA32_FIXED_CTR1", 0x30a, -1)
    idc.AddConstEx(id, "IA32_FIXED_CTR2", 0x30b, -1)
    idc.AddConstEx(id, "IA32_PERF_CAPABILITIES", 0x345, -1)
    idc.AddConstEx(id, "IA32_FIXED_CTR_CTRL", 0x38d, -1)
    idc.AddConstEx(id, "IA32_PERF_GLOBAL_STATUS", 0x38e, -1)
    idc.AddConstEx(id, "IA32_PERF_GLOBAL_CTRL", 0x38f, -1)
    idc.AddConstEx(id, "IA32_PERF_GLOBAL_STATUS_RESET", 0x390, -1)
    idc.AddConstEx(id, "IA32_PERF_GLOBAL_STATUS_SET", 0x391, -1)
    idc.AddConstEx(id, "IA32_PERF_GLOBAL_INUSE", 0x392, -1)
    idc.AddConstEx(id, "IA32_PEBS_ENABLE", 0x3f1, -1)
    idc.AddConstEx(id, "IA32_MC0_CTL", 0x400, -1)
    idc.AddConstEx(id, "IA32_MC1_CTL", 0x404, -1)
    idc.AddConstEx(id, "IA32_MC2_CTL", 0x408, -1)
    idc.AddConstEx(id, "IA32_MC3_CTL", 0x40c, -1)
    idc.AddConstEx(id, "IA32_MC4_CTL", 0x410, -1)
    idc.AddConstEx(id, "IA32_MC5_CTL", 0x414, -1)
    idc.AddConstEx(id, "IA32_MC6_CTL", 0x418, -1)
    idc.AddConstEx(id, "IA32_MC7_CTL", 0x41c, -1)
    idc.AddConstEx(id, "IA32_MC8_CTL", 0x420, -1)
    idc.AddConstEx(id, "IA32_MC9_CTL", 0x424, -1)
    idc.AddConstEx(id, "IA32_MC10_CTL", 0x428, -1)
    idc.AddConstEx(id, "IA32_MC11_CTL", 0x42c, -1)
    idc.AddConstEx(id, "IA32_MC12_CTL", 0x430, -1)
    idc.AddConstEx(id, "IA32_MC13_CTL", 0x434, -1)
    idc.AddConstEx(id, "IA32_MC14_CTL", 0x438, -1)
    idc.AddConstEx(id, "IA32_MC15_CTL", 0x43c, -1)
    idc.AddConstEx(id, "IA32_MC16_CTL", 0x440, -1)
    idc.AddConstEx(id, "IA32_MC17_CTL", 0x444, -1)
    idc.AddConstEx(id, "IA32_MC18_CTL", 0x448, -1)
    idc.AddConstEx(id, "IA32_MC19_CTL", 0x44c, -1)
    idc.AddConstEx(id, "IA32_MC20_CTL", 0x450, -1)
    idc.AddConstEx(id, "IA32_MC21_CTL", 0x454, -1)
    idc.AddConstEx(id, "IA32_MC22_CTL", 0x458, -1)
    idc.AddConstEx(id, "IA32_MC23_CTL", 0x45c, -1)
    idc.AddConstEx(id, "IA32_MC24_CTL", 0x460, -1)
    idc.AddConstEx(id, "IA32_MC25_CTL", 0x464, -1)
    idc.AddConstEx(id, "IA32_MC26_CTL", 0x468, -1)
    idc.AddConstEx(id, "IA32_MC27_CTL", 0x46c, -1)
    idc.AddConstEx(id, "IA32_MC28_CTL", 0x470, -1)
    idc.AddConstEx(id, "IA32_MC0_STATUS", 0x401, -1)
    idc.AddConstEx(id, "IA32_MC1_STATUS", 0x405, -1)
    idc.AddConstEx(id, "IA32_MC2_STATUS", 0x409, -1)
    idc.AddConstEx(id, "IA32_MC3_STATUS", 0x40d, -1)
    idc.AddConstEx(id, "IA32_MC4_STATUS", 0x411, -1)
    idc.AddConstEx(id, "IA32_MC5_STATUS", 0x415, -1)
    idc.AddConstEx(id, "IA32_MC6_STATUS", 0x419, -1)
    idc.AddConstEx(id, "IA32_MC7_STATUS", 0x41d, -1)
    idc.AddConstEx(id, "IA32_MC8_STATUS", 0x421, -1)
    idc.AddConstEx(id, "IA32_MC9_STATUS", 0x425, -1)
    idc.AddConstEx(id, "IA32_MC10_STATUS", 0x429, -1)
    idc.AddConstEx(id, "IA32_MC11_STATUS", 0x42d, -1)
    idc.AddConstEx(id, "IA32_MC12_STATUS", 0x431, -1)
    idc.AddConstEx(id, "IA32_MC13_STATUS", 0x435, -1)
    idc.AddConstEx(id, "IA32_MC14_STATUS", 0x439, -1)
    idc.AddConstEx(id, "IA32_MC15_STATUS", 0x43d, -1)
    idc.AddConstEx(id, "IA32_MC16_STATUS", 0x441, -1)
    idc.AddConstEx(id, "IA32_MC17_STATUS", 0x445, -1)
    idc.AddConstEx(id, "IA32_MC18_STATUS", 0x449, -1)
    idc.AddConstEx(id, "IA32_MC19_STATUS", 0x44d, -1)
    idc.AddConstEx(id, "IA32_MC20_STATUS", 0x451, -1)
    idc.AddConstEx(id, "IA32_MC21_STATUS", 0x455, -1)
    idc.AddConstEx(id, "IA32_MC22_STATUS", 0x459, -1)
    idc.AddConstEx(id, "IA32_MC23_STATUS", 0x45d, -1)
    idc.AddConstEx(id, "IA32_MC24_STATUS", 0x461, -1)
    idc.AddConstEx(id, "IA32_MC25_STATUS", 0x465, -1)
    idc.AddConstEx(id, "IA32_MC26_STATUS", 0x469, -1)
    idc.AddConstEx(id, "IA32_MC27_STATUS", 0x46d, -1)
    idc.AddConstEx(id, "IA32_MC28_STATUS", 0x471, -1)
    idc.AddConstEx(id, "IA32_MC0_ADDR", 0x402, -1)
    idc.AddConstEx(id, "IA32_MC1_ADDR", 0x406, -1)
    idc.AddConstEx(id, "IA32_MC2_ADDR", 0x40a, -1)
    idc.AddConstEx(id, "IA32_MC3_ADDR", 0x40e, -1)
    idc.AddConstEx(id, "IA32_MC4_ADDR", 0x412, -1)
    idc.AddConstEx(id, "IA32_MC5_ADDR", 0x416, -1)
    idc.AddConstEx(id, "IA32_MC6_ADDR", 0x41a, -1)
    idc.AddConstEx(id, "IA32_MC7_ADDR", 0x41e, -1)
    idc.AddConstEx(id, "IA32_MC8_ADDR", 0x422, -1)
    idc.AddConstEx(id, "IA32_MC9_ADDR", 0x426, -1)
    idc.AddConstEx(id, "IA32_MC10_ADDR", 0x42a, -1)
    idc.AddConstEx(id, "IA32_MC11_ADDR", 0x42e, -1)
    idc.AddConstEx(id, "IA32_MC12_ADDR", 0x432, -1)
    idc.AddConstEx(id, "IA32_MC13_ADDR", 0x436, -1)
    idc.AddConstEx(id, "IA32_MC14_ADDR", 0x43a, -1)
    idc.AddConstEx(id, "IA32_MC15_ADDR", 0x43e, -1)
    idc.AddConstEx(id, "IA32_MC16_ADDR", 0x442, -1)
    idc.AddConstEx(id, "IA32_MC17_ADDR", 0x446, -1)
    idc.AddConstEx(id, "IA32_MC18_ADDR", 0x44a, -1)
    idc.AddConstEx(id, "IA32_MC19_ADDR", 0x44e, -1)
    idc.AddConstEx(id, "IA32_MC20_ADDR", 0x452, -1)
    idc.AddConstEx(id, "IA32_MC21_ADDR", 0x456, -1)
    idc.AddConstEx(id, "IA32_MC22_ADDR", 0x45a, -1)
    idc.AddConstEx(id, "IA32_MC23_ADDR", 0x45e, -1)
    idc.AddConstEx(id, "IA32_MC24_ADDR", 0x462, -1)
    idc.AddConstEx(id, "IA32_MC25_ADDR", 0x466, -1)
    idc.AddConstEx(id, "IA32_MC26_ADDR", 0x46a, -1)
    idc.AddConstEx(id, "IA32_MC27_ADDR", 0x46e, -1)
    idc.AddConstEx(id, "IA32_MC28_ADDR", 0x472, -1)
    idc.AddConstEx(id, "IA32_MC0_MISC", 0x403, -1)
    idc.AddConstEx(id, "IA32_MC1_MISC", 0x407, -1)
    idc.AddConstEx(id, "IA32_MC2_MISC", 0x40b, -1)
    idc.AddConstEx(id, "IA32_MC3_MISC", 0x40f, -1)
    idc.AddConstEx(id, "IA32_MC4_MISC", 0x413, -1)
    idc.AddConstEx(id, "IA32_MC5_MISC", 0x417, -1)
    idc.AddConstEx(id, "IA32_MC6_MISC", 0x41b, -1)
    idc.AddConstEx(id, "IA32_MC7_MISC", 0x41f, -1)
    idc.AddConstEx(id, "IA32_MC8_MISC", 0x423, -1)
    idc.AddConstEx(id, "IA32_MC9_MISC", 0x427, -1)
    idc.AddConstEx(id, "IA32_MC10_MISC", 0x42b, -1)
    idc.AddConstEx(id, "IA32_MC11_MISC", 0x42f, -1)
    idc.AddConstEx(id, "IA32_MC12_MISC", 0x433, -1)
    idc.AddConstEx(id, "IA32_MC13_MISC", 0x437, -1)
    idc.AddConstEx(id, "IA32_MC14_MISC", 0x43b, -1)
    idc.AddConstEx(id, "IA32_MC15_MISC", 0x43f, -1)
    idc.AddConstEx(id, "IA32_MC16_MISC", 0x443, -1)
    idc.AddConstEx(id, "IA32_MC17_MISC", 0x447, -1)
    idc.AddConstEx(id, "IA32_MC18_MISC", 0x44b, -1)
    idc.AddConstEx(id, "IA32_MC19_MISC", 0x44f, -1)
    idc.AddConstEx(id, "IA32_MC20_MISC", 0x453, -1)
    idc.AddConstEx(id, "IA32_MC21_MISC", 0x457, -1)
    idc.AddConstEx(id, "IA32_MC22_MISC", 0x45b, -1)
    idc.AddConstEx(id, "IA32_MC23_MISC", 0x45f, -1)
    idc.AddConstEx(id, "IA32_MC24_MISC", 0x463, -1)
    idc.AddConstEx(id, "IA32_MC25_MISC", 0x467, -1)
    idc.AddConstEx(id, "IA32_MC26_MISC", 0x46b, -1)
    idc.AddConstEx(id, "IA32_MC27_MISC", 0x46f, -1)
    idc.AddConstEx(id, "IA32_MC28_MISC", 0x473, -1)
    idc.AddConstEx(id, "IA32_VMX_BASIC", 0x480, -1)
    idc.AddConstEx(id, "IA32_VMX_PINBASED_CTLS", 0x481, -1)
    idc.AddConstEx(id, "IA32_VMX_PROCBASED_CTLS", 0x482, -1)
    idc.AddConstEx(id, "IA32_VMX_EXIT_CTLS", 0x483, -1)
    idc.AddConstEx(id, "IA32_VMX_ENTRY_CTLS", 0x484, -1)
    idc.AddConstEx(id, "IA32_VMX_MISC", 0x485, -1)
    idc.AddConstEx(id, "IA32_VMX_CR0_FIXED0", 0x486, -1)
    idc.AddConstEx(id, "IA32_VMX_CR0_FIXED1", 0x487, -1)
    idc.AddConstEx(id, "IA32_VMX_CR4_FIXED0", 0x488, -1)
    idc.AddConstEx(id, "IA32_VMX_CR4_FIXED1", 0x489, -1)
    idc.AddConstEx(id, "IA32_VMX_VMCS_ENUM", 0x48a, -1)
    idc.AddConstEx(id, "IA32_VMX_PROCBASED_CTLS2", 0x48b, -1)
    idc.AddConstEx(id, "IA32_VMX_EPT_VPID_CAP", 0x48c, -1)
    idc.AddConstEx(id, "IA32_VMX_TRUE_PINBASED_CTLS", 0x48d, -1)
    idc.AddConstEx(id, "IA32_VMX_TRUE_PROCBASED_CTLS", 0x48e, -1)
    idc.AddConstEx(id, "IA32_VMX_TRUE_EXIT_CTLS", 0x48f, -1)
    idc.AddConstEx(id, "IA32_VMX_TRUE_ENTRY_CTLS", 0x490, -1)
    idc.AddConstEx(id, "IA32_VMX_VMFUNC", 0x491, -1)
    idc.AddConstEx(id, "IA32_A_PMC0", 0x4c1, -1)
    idc.AddConstEx(id, "IA32_A_PMC1", 0x4c2, -1)
    idc.AddConstEx(id, "IA32_A_PMC2", 0x4c3, -1)
    idc.AddConstEx(id, "IA32_A_PMC3", 0x4c4, -1)
    idc.AddConstEx(id, "IA32_A_PMC4", 0x4c5, -1)
    idc.AddConstEx(id, "IA32_A_PMC5", 0x4c6, -1)
    idc.AddConstEx(id, "IA32_A_PMC6", 0x4c7, -1)
    idc.AddConstEx(id, "IA32_A_PMC7", 0x4c8, -1)
    idc.AddConstEx(id, "IA32_MCG_EXT_CTL", 0x4d0, -1)
    idc.AddConstEx(id, "IA32_SGX_SVN_STATUS", 0x500, -1)
    idc.AddConstEx(id, "IA32_RTIT_OUTPUT_BASE", 0x560, -1)
    idc.AddConstEx(id, "IA32_RTIT_OUTPUT_MASK_PTRS", 0x561, -1)
    idc.AddConstEx(id, "IA32_RTIT_CTL", 0x570, -1)
    idc.AddConstEx(id, "IA32_RTIT_STATUS", 0x571, -1)
    idc.AddConstEx(id, "IA32_RTIT_CR3_MATCH", 0x572, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR0_A", 0x580, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR1_A", 0x582, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR2_A", 0x584, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR3_A", 0x586, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR0_B", 0x581, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR1_B", 0x583, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR2_B", 0x585, -1)
    idc.AddConstEx(id, "IA32_RTIT_ADDR3_B", 0x587, -1)
    idc.AddConstEx(id, "IA32_DS_AREA", 0x600, -1)
    idc.AddConstEx(id, "IA32_TSC_DEADLINE", 0x6e0, -1)
    idc.AddConstEx(id, "IA32_PM_ENABLE", 0x770, -1)
    idc.AddConstEx(id, "IA32_HWP_CAPABILITIES", 0x771, -1)
    idc.AddConstEx(id, "IA32_HWP_REQUEST_PKG", 0x772, -1)
    idc.AddConstEx(id, "IA32_HWP_INTERRUPT", 0x773, -1)
    idc.AddConstEx(id, "IA32_HWP_REQUEST", 0x774, -1)
    idc.AddConstEx(id, "IA32_HWP_STATUS", 0x777, -1)
    idc.AddConstEx(id, "IA32_X2APIC_APICID", 0x802, -1)
    idc.AddConstEx(id, "IA32_X2APIC_VERSION", 0x803, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TPR", 0x808, -1)
    idc.AddConstEx(id, "IA32_X2APIC_PPR", 0x80a, -1)
    idc.AddConstEx(id, "IA32_X2APIC_EOI", 0x80b, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LDR", 0x80d, -1)
    idc.AddConstEx(id, "IA32_X2APIC_SIVR", 0x80f, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR0", 0x810, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR1", 0x811, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR2", 0x812, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR3", 0x813, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR4", 0x814, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR5", 0x815, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR6", 0x816, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ISR7", 0x817, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR0", 0x818, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR1", 0x819, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR2", 0x81a, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR3", 0x81b, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR4", 0x81c, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR5", 0x81d, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR6", 0x81e, -1)
    idc.AddConstEx(id, "IA32_X2APIC_TMR7", 0x81f, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR0", 0x820, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR1", 0x821, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR2", 0x822, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR3", 0x823, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR4", 0x824, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR5", 0x825, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR6", 0x826, -1)
    idc.AddConstEx(id, "IA32_X2APIC_IRR7", 0x827, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ESR", 0x828, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LVT_CMCI", 0x82f, -1)
    idc.AddConstEx(id, "IA32_X2APIC_ICR", 0x830, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LVT_TIMER", 0x832, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LVT_THERMAL", 0x833, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LVT_PMI", 0x834, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LVT_LINT0", 0x835, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LVT_LINT1", 0x836, -1)
    idc.AddConstEx(id, "IA32_X2APIC_LVT_ERROR", 0x837, -1)
    idc.AddConstEx(id, "IA32_X2APIC_INIT_COUNT", 0x838, -1)
    idc.AddConstEx(id, "IA32_X2APIC_CUR_COUNT", 0x839, -1)
    idc.AddConstEx(id, "IA32_X2APIC_DIV_CONF", 0x83e, -1)
    idc.AddConstEx(id, "IA32_X2APIC_SELF_IPI", 0x83f, -1)
    idc.AddConstEx(id, "IA32_DEBUG_INTERFACE", 0xc80, -1)
    idc.AddConstEx(id, "IA32_L2_QOS_CFG", 0xc82, -1)
    idc.AddConstEx(id, "IA32_QM_CTR", 0xc8e, -1)
    idc.AddConstEx(id, "IA32_BNDCFGS", 0xd90, -1)
    idc.AddConstEx(id, "IA32_XSS", 0xda0, -1)
    idc.AddConstEx(id, "IA32_PKG_HDC_CTL", 0xdb0, -1)
    idc.AddConstEx(id, "IA32_PM_CTL1", 0xdb1, -1)
    idc.AddConstEx(id, "IA32_THREAD_STALL", 0xdb2, -1)
    idc.AddConstEx(id, "IA32_EFER", 0xc0000080, -1)
    idc.AddConstEx(id, "IA32_STAR", 0xc0000081, -1)
    idc.AddConstEx(id, "IA32_LSTAR", 0xc0000082, -1)
    idc.AddConstEx(id, "IA32_CSTAR", 0xc0000083, -1)
    idc.AddConstEx(id, "IA32_FMASK", 0xc0000084, -1)
    idc.AddConstEx(id, "IA32_FS_BASE", 0xc0000100, -1)
    idc.AddConstEx(id, "IA32_GS_BASE", 0xc0000101, -1)
    idc.AddConstEx(id, "IA32_KERNEL_GS_BASE", 0xc0000102, -1)
    idc.AddConstEx(id, "IA32_TSC_AUX", 0xc0000103, -1)
    # ------------------------------------------------------------------------------------- #
    # VMCS 
    # ------------------------------------------------------------------------------------- #
    id = idc.AddEnum(0, "IA32_VMCS_LIST_ENUM", idaapi.hexflag())
    idc.AddConstEx(id, "VMCS_CTRL_VPID", 0x0, -1)
    idc.AddConstEx(id, "VMCS_CTRL_POSTED_INTR_NOTIFY_VECTOR", 0x2, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EPTP_INDEX", 0x4, -1)
    idc.AddConstEx(id, "VMCS_GUEST_ES_SEL", 0x800, -1)
    idc.AddConstEx(id, "VMCS_GUEST_CS_SEL", 0x802, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SS_SEL", 0x804, -1)
    idc.AddConstEx(id, "VMCS_GUEST_DS_SEL", 0x806, -1)
    idc.AddConstEx(id, "VMCS_GUEST_FS_SEL", 0x808, -1)
    idc.AddConstEx(id, "VMCS_GUEST_GS_SEL", 0x80a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_LDTR_SEL", 0x80c, -1)
    idc.AddConstEx(id, "VMCS_GUEST_TR_SEL", 0x80e, -1)
    idc.AddConstEx(id, "VMCS_GUEST_INTR_STATUS", 0x810, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PML_INDEX", 0x812, -1)
    idc.AddConstEx(id, "VMCS_HOST_ES_SEL", 0xc00, -1)
    idc.AddConstEx(id, "VMCS_HOST_CS_SEL", 0xc02, -1)
    idc.AddConstEx(id, "VMCS_HOST_SS_SEL", 0xc04, -1)
    idc.AddConstEx(id, "VMCS_HOST_DS_SEL", 0xc06, -1)
    idc.AddConstEx(id, "VMCS_HOST_FS_SEL", 0xc08, -1)
    idc.AddConstEx(id, "VMCS_HOST_GS_SEL", 0xc0a, -1)
    idc.AddConstEx(id, "VMCS_HOST_TR_SEL", 0xc0c, -1)
    idc.AddConstEx(id, "VMCS_CTRL_IO_BITMAP_A", 0x2000, -1)
    idc.AddConstEx(id, "VMCS_CTRL_IO_BITMAP_B", 0x2002, -1)
    idc.AddConstEx(id, "VMCS_CTRL_MSR_BITMAP", 0x2004, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VMEXIT_MSR_STORE", 0x2006, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VMEXIT_MSR_LOAD", 0x2008, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VMENTRY_MSR_LOAD", 0x200a, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EXEC_VMCS_PTR", 0x200c, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PML_ADDR", 0x200e, -1)
    idc.AddConstEx(id, "VMCS_CTRL_TSC_OFFSET", 0x2010, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VAPIC_PAGEADDR", 0x2012, -1)
    idc.AddConstEx(id, "VMCS_CTRL_APIC_ACCESSADDR", 0x2014, -1)
    idc.AddConstEx(id, "VMCS_CTRL_POSTED_INTR_DESC", 0x2016, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VMFUNC_CTRLS", 0x2018, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EPTP", 0x201a, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EOI_BITMAP_0", 0x201c, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EOI_BITMAP_1", 0x201e, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EOI_BITMAP_2", 0x2020, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EOI_BITMAP_3", 0x2022, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EPTP_LIST", 0x2024, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VMREAD_BITMAP", 0x2026, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VMWRITE_BITMAP", 0x2028, -1)
    idc.AddConstEx(id, "VMCS_CTRL_VIRTXCPT_INFO_ADDR", 0x202a, -1)
    idc.AddConstEx(id, "VMCS_CTRL_XSS_EXITING_BITMAP", 0x202c, -1)
    idc.AddConstEx(id, "VMCS_CTRL_ENCLS_EXITING_BITMAP", 0x202e, -1)
    idc.AddConstEx(id, "VMCS_CTRL_TSC_MULTIPLIER", 0x2032, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PHYS_ADDR", 0x2400, -1)
    idc.AddConstEx(id, "VMCS_GUEST_VMCS_LINK_PTR", 0x2800, -1)
    idc.AddConstEx(id, "VMCS_GUEST_DEBUGCTL", 0x2802, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PAT", 0x2804, -1)
    idc.AddConstEx(id, "VMCS_GUEST_EFER", 0x2806, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PERF_GLOBAL_CTRL", 0x2808, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PDPTE0", 0x280a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PDPTE1", 0x280c, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PDPTE2", 0x280e, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PDPTE3", 0x2810, -1)
    idc.AddConstEx(id, "VMCS_HOST_PAT", 0x2c00, -1)
    idc.AddConstEx(id, "VMCS_HOST_EFER", 0x2c02, -1)
    idc.AddConstEx(id, "VMCS_HOST_PERF_GLOBAL_CTRL", 0x2c04, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PIN_EXEC", 0x4000, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PROC_EXEC", 0x4002, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EXCEPTION_BITMAP", 0x4004, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PAGEFAULT_ERROR_MASK", 0x4006, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PAGEFAULT_ERROR_MATCH", 0x4008, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR3_TARGET_COUNT", 0x400a, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EXIT", 0x400c, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EXIT_MSR_STORE_COUNT", 0x400e, -1)
    idc.AddConstEx(id, "VMCS_CTRL_EXIT_MSR_LOAD_COUNT", 0x4010, -1)
    idc.AddConstEx(id, "VMCS_CTRL_ENTRY", 0x4012, -1)
    idc.AddConstEx(id, "VMCS_CTRL_ENTRY_MSR_LOAD_COUNT", 0x4014, -1)
    idc.AddConstEx(id, "VMCS_CTRL_ENTRY_INTERRUPTION_INFO", 0x4016, -1)
    idc.AddConstEx(id, "VMCS_CTRL_ENTRY_EXCEPTION_ERRCODE", 0x4018, -1)
    idc.AddConstEx(id, "VMCS_CTRL_ENTRY_INSTR_LENGTH", 0x401a, -1)
    idc.AddConstEx(id, "VMCS_CTRL_TPR_THRESHOLD", 0x401c, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PROC_EXEC2", 0x401e, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PLE_GAP", 0x4020, -1)
    idc.AddConstEx(id, "VMCS_CTRL_PLE_WINDOW", 0x4022, -1)
    idc.AddConstEx(id, "VMCS_VM_INSTR_ERROR", 0x4400, -1)
    idc.AddConstEx(id, "VMCS_EXIT_REASON", 0x4402, -1)
    idc.AddConstEx(id, "VMCS_EXIT_INTERRUPTION_INFO", 0x4404, -1)
    idc.AddConstEx(id, "VMCS_EXIT_INTERRUPTION_ERROR_CODE", 0x4406, -1)
    idc.AddConstEx(id, "VMCS_IDT_VECTORING_INFO", 0x4408, -1)
    idc.AddConstEx(id, "VMCS_IDT_VECTORING_ERROR_CODE", 0x440a, -1)
    idc.AddConstEx(id, "VMCS_EXIT_INSTR_LENGTH", 0x440c, -1)
    idc.AddConstEx(id, "VMCS_EXIT_INSTR_INFO", 0x440e, -1)
    idc.AddConstEx(id, "VMCS_GUEST_ES_LIMIT", 0x4800, -1)
    idc.AddConstEx(id, "VMCS_GUEST_CS_LIMIT", 0x4802, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SS_LIMIT", 0x4804, -1)
    idc.AddConstEx(id, "VMCS_GUEST_DS_LIMIT", 0x4806, -1)
    idc.AddConstEx(id, "VMCS_GUEST_FS_LIMIT", 0x4808, -1)
    idc.AddConstEx(id, "VMCS_GUEST_GS_LIMIT", 0x480a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_LDTR_LIMIT", 0x480c, -1)
    idc.AddConstEx(id, "VMCS_GUEST_TR_LIMIT", 0x480e, -1)
    idc.AddConstEx(id, "VMCS_GUEST_GDTR_LIMIT", 0x4810, -1)
    idc.AddConstEx(id, "VMCS_GUEST_IDTR_LIMIT", 0x4812, -1)
    idc.AddConstEx(id, "VMCS_GUEST_ES_ACCESS_RIGHTS", 0x4814, -1)
    idc.AddConstEx(id, "VMCS_GUEST_CS_ACCESS_RIGHTS", 0x4816, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SS_ACCESS_RIGHTS", 0x4818, -1)
    idc.AddConstEx(id, "VMCS_GUEST_DS_ACCESS_RIGHTS", 0x481a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_FS_ACCESS_RIGHTS", 0x481c, -1)
    idc.AddConstEx(id, "VMCS_GUEST_GS_ACCESS_RIGHTS", 0x481e, -1)
    idc.AddConstEx(id, "VMCS_GUEST_LDTR_ACCESS_RIGHTS", 0x4820, -1)
    idc.AddConstEx(id, "VMCS_GUEST_TR_ACCESS_RIGHTS", 0x4822, -1)
    idc.AddConstEx(id, "VMCS_GUEST_INTERRUPTIBILITY_STATE", 0x4824, -1)
    idc.AddConstEx(id, "VMCS_GUEST_ACTIVITY_STATE", 0x4826, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SMBASE", 0x4828, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SYSENTER_CS", 0x482a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PREEMPT_TIMER_VALUE", 0x482e, -1)
    idc.AddConstEx(id, "VMCS_SYSENTER_CS", 0x4c00, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR0_MASK", 0x6000, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR4_MASK", 0x6002, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR0_READ_SHADOW", 0x6004, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR4_READ_SHADOW", 0x6006, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR3_TARGET_VAL0", 0x6008, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR3_TARGET_VAL1", 0x600a, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR3_TARGET_VAL2", 0x600c, -1)
    idc.AddConstEx(id, "VMCS_CTRL_CR3_TARGET_VAL3", 0x600e, -1)
    idc.AddConstEx(id, "VMCS_EXIT_QUALIFICATION", 0x6400, -1)
    idc.AddConstEx(id, "VMCS_IO_RCX", 0x6402, -1)
    idc.AddConstEx(id, "VMCS_IO_RSX", 0x6404, -1)
    idc.AddConstEx(id, "VMCS_IO_RDI", 0x6406, -1)
    idc.AddConstEx(id, "VMCS_IO_RIP", 0x6408, -1)
    idc.AddConstEx(id, "VMCS_EXIT_GUEST_LINEAR_ADDR", 0x640a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_CR0", 0x6800, -1)
    idc.AddConstEx(id, "VMCS_GUEST_CR3", 0x6802, -1)
    idc.AddConstEx(id, "VMCS_GUEST_CR4", 0x6804, -1)
    idc.AddConstEx(id, "VMCS_GUEST_ES_BASE", 0x6806, -1)
    idc.AddConstEx(id, "VMCS_GUEST_CS_BASE", 0x6808, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SS_BASE", 0x680a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_DS_BASE", 0x680c, -1)
    idc.AddConstEx(id, "VMCS_GUEST_FS_BASE", 0x680e, -1)
    idc.AddConstEx(id, "VMCS_GUEST_GS_BASE", 0x6810, -1)
    idc.AddConstEx(id, "VMCS_GUEST_LDTR_BASE", 0x6812, -1)
    idc.AddConstEx(id, "VMCS_GUEST_TR_BASE", 0x6814, -1)
    idc.AddConstEx(id, "VMCS_GUEST_GDTR_BASE", 0x6816, -1)
    idc.AddConstEx(id, "VMCS_GUEST_IDTR_BASE", 0x6818, -1)
    idc.AddConstEx(id, "VMCS_GUEST_DR7", 0x681a, -1)
    idc.AddConstEx(id, "VMCS_GUEST_RSP", 0x681c, -1)
    idc.AddConstEx(id, "VMCS_GUEST_RIP", 0x681e, -1)
    idc.AddConstEx(id, "VMCS_GUEST_RFLAGS", 0x6820, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PENDING_DEBUG_EXCEPTIONS", 0x6822, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SYSENTER_ESP", 0x6824, -1)
    idc.AddConstEx(id, "VMCS_GUEST_SYSENTER_EIP", 0x6826, -1)
    idc.AddConstEx(id, "VMCS_HOST_CR0", 0x6c00, -1)
    idc.AddConstEx(id, "VMCS_HOST_CR3", 0x6c02, -1)
    idc.AddConstEx(id, "VMCS_HOST_CR4", 0x6c04, -1)
    idc.AddConstEx(id, "VMCS_HOST_FS_BASE", 0x6c06, -1)
    idc.AddConstEx(id, "VMCS_HOST_GS_BASE", 0x6c08, -1)
    idc.AddConstEx(id, "VMCS_HOST_TR_BASE", 0x6c0a, -1)
    idc.AddConstEx(id, "VMCS_HOST_GDTR_BASE", 0x6c0c, -1)
    idc.AddConstEx(id, "VMCS_HOST_IDTR_BASE", 0x6c0e, -1)
    idc.AddConstEx(id, "VMCS_HOST_SYSENTER_ESP", 0x6c10, -1)
    idc.AddConstEx(id, "VMCS_HOST_SYSENTER_EIP", 0x6c12, -1)
    idc.AddConstEx(id, "VMCS_HOST_RSP", 0x6c14, -1)
    idc.AddConstEx(id, "VMCS_HOST_RIP", 0x6c16, -1)
    idc.AddConstEx(id, "VMCS_GUEST_PDPTR3_HIGH", 0x2811, -1)
    idc.AddConstEx(id, "VMCS_GUEST_BNDCFGS", 0x2812, -1)
    idc.AddConstEx(id, "VMCS_GUEST_BNDCFGS_HIGH", 0x2813, -1)
    idc.AddConstEx(id, "VMCS_GUEST_RTIT_CTL", 0x2814, -1)
    idc.AddConstEx(id, "VMCS_GUEST_RTIT_CTL_HIGH", 0x2815, -1)

# ------------------------------------------------------------------------------------- #
class KnowledgeDb():
    def __init__(self):
        self.ia32_vmcs_db = {
            0x00000000 : 'VMCS_CTRL_VPID',
            0x00000002 : 'VMCS_CTRL_POSTED_INTR_NOTIFY_VECTOR',
            0x00000004 : 'VMCS_CTRL_EPTP_INDEX',
            0x00000800 : 'VMCS_GUEST_ES_SEL',
            0x00000802 : 'VMCS_GUEST_CS_SEL',
            0x00000804 : 'VMCS_GUEST_SS_SEL',
            0x00000806 : 'VMCS_GUEST_DS_SEL',
            0x00000808 : 'VMCS_GUEST_FS_SEL',
            0x0000080A : 'VMCS_GUEST_GS_SEL',
            0x0000080C : 'VMCS_GUEST_LDTR_SEL',
            0x0000080E : 'VMCS_GUEST_TR_SEL',
            0x00000810 : 'VMCS_GUEST_INTR_STATUS',
            0x00000812 : 'VMCS_GUEST_PML_INDEX',
            0x00000C00 : 'VMCS_HOST_ES_SEL',
            0x00000C02 : 'VMCS_HOST_CS_SEL',
            0x00000C04 : 'VMCS_HOST_SS_SEL',
            0x00000C06 : 'VMCS_HOST_DS_SEL',
            0x00000C08 : 'VMCS_HOST_FS_SEL',
            0x00000C0A : 'VMCS_HOST_GS_SEL',
            0x00000C0C : 'VMCS_HOST_TR_SEL',
            0x00002000 : 'VMCS_CTRL_IO_BITMAP_A',
            0x00002002 : 'VMCS_CTRL_IO_BITMAP_B',
            0x00002004 : 'VMCS_CTRL_MSR_BITMAP',
            0x00002006 : 'VMCS_CTRL_VMEXIT_MSR_STORE',
            0x00002008 : 'VMCS_CTRL_VMEXIT_MSR_LOAD',
            0x0000200A : 'VMCS_CTRL_VMENTRY_MSR_LOAD',
            0x0000200C : 'VMCS_CTRL_EXEC_VMCS_PTR',
            0x0000200E : 'VMCS_CTRL_PML_ADDR',
            0x00002010 : 'VMCS_CTRL_TSC_OFFSET',
            0x00002012 : 'VMCS_CTRL_VAPIC_PAGEADDR',
            0x00002014 : 'VMCS_CTRL_APIC_ACCESSADDR',
            0x00002016 : 'VMCS_CTRL_POSTED_INTR_DESC',
            0x00002018 : 'VMCS_CTRL_VMFUNC_CTRLS',
            0x0000201A : 'VMCS_CTRL_EPTP',
            0x0000201C : 'VMCS_CTRL_EOI_BITMAP_0',
            0x0000201E : 'VMCS_CTRL_EOI_BITMAP_1',
            0x00002020 : 'VMCS_CTRL_EOI_BITMAP_2',
            0x00002022 : 'VMCS_CTRL_EOI_BITMAP_3',
            0x00002024 : 'VMCS_CTRL_EPTP_LIST',
            0x00002026 : 'VMCS_CTRL_VMREAD_BITMAP',
            0x00002028 : 'VMCS_CTRL_VMWRITE_BITMAP',
            0x0000202A : 'VMCS_CTRL_VIRTXCPT_INFO_ADDR',
            0x0000202C : 'VMCS_CTRL_XSS_EXITING_BITMAP',
            0x0000202E : 'VMCS_CTRL_ENCLS_EXITING_BITMAP',
            0x00002032 : 'VMCS_CTRL_TSC_MULTIPLIER',
            0x00002400 : 'VMCS_GUEST_PHYS_ADDR',
            0x00002800 : 'VMCS_GUEST_VMCS_LINK_PTR',
            0x00002802 : 'VMCS_GUEST_DEBUGCTL',
            0x00002804 : 'VMCS_GUEST_PAT',
            0x00002806 : 'VMCS_GUEST_EFER',
            0x00002808 : 'VMCS_GUEST_PERF_GLOBAL_CTRL',
            0x0000280A : 'VMCS_GUEST_PDPTE0',
            0x0000280C : 'VMCS_GUEST_PDPTE1',
            0x0000280E : 'VMCS_GUEST_PDPTE2',
            0x00002810 : 'VMCS_GUEST_PDPTE3',
            0x00002C00 : 'VMCS_HOST_PAT',
            0x00002C02 : 'VMCS_HOST_EFER',
            0x00002C04 : 'VMCS_HOST_PERF_GLOBAL_CTRL',
            0x00004000 : 'VMCS_CTRL_PIN_EXEC',
            0x00004002 : 'VMCS_CTRL_PROC_EXEC',
            0x00004004 : 'VMCS_CTRL_EXCEPTION_BITMAP',
            0x00004006 : 'VMCS_CTRL_PAGEFAULT_ERROR_MASK',
            0x00004008 : 'VMCS_CTRL_PAGEFAULT_ERROR_MATCH',
            0x0000400A : 'VMCS_CTRL_CR3_TARGET_COUNT',
            0x0000400C : 'VMCS_CTRL_EXIT',
            0x0000400E : 'VMCS_CTRL_EXIT_MSR_STORE_COUNT',
            0x00004010 : 'VMCS_CTRL_EXIT_MSR_LOAD_COUNT',
            0x00004012 : 'VMCS_CTRL_ENTRY',
            0x00004014 : 'VMCS_CTRL_ENTRY_MSR_LOAD_COUNT',
            0x00004016 : 'VMCS_CTRL_ENTRY_INTERRUPTION_INFO',
            0x00004018 : 'VMCS_CTRL_ENTRY_EXCEPTION_ERRCODE',
            0x0000401A : 'VMCS_CTRL_ENTRY_INSTR_LENGTH',
            0x0000401C : 'VMCS_CTRL_TPR_THRESHOLD',
            0x0000401E : 'VMCS_CTRL_PROC_EXEC2',
            0x00004020 : 'VMCS_CTRL_PLE_GAP',
            0x00004022 : 'VMCS_CTRL_PLE_WINDOW',
            0x00004400 : 'VMCS_VM_INSTR_ERROR',
            0x00004402 : 'VMCS_EXIT_REASON',
            0x00004404 : 'VMCS_EXIT_INTERRUPTION_INFO',
            0x00004406 : 'VMCS_EXIT_INTERRUPTION_ERROR_CODE',
            0x00004408 : 'VMCS_IDT_VECTORING_INFO',
            0x0000440A : 'VMCS_IDT_VECTORING_ERROR_CODE',
            0x0000440C : 'VMCS_EXIT_INSTR_LENGTH',
            0x0000440E : 'VMCS_EXIT_INSTR_INFO',
            0x00004800 : 'VMCS_GUEST_ES_LIMIT',
            0x00004802 : 'VMCS_GUEST_CS_LIMIT',
            0x00004804 : 'VMCS_GUEST_SS_LIMIT',
            0x00004806 : 'VMCS_GUEST_DS_LIMIT',
            0x00004808 : 'VMCS_GUEST_FS_LIMIT',
            0x0000480A : 'VMCS_GUEST_GS_LIMIT',
            0x0000480C : 'VMCS_GUEST_LDTR_LIMIT',
            0x0000480E : 'VMCS_GUEST_TR_LIMIT',
            0x00004810 : 'VMCS_GUEST_GDTR_LIMIT',
            0x00004812 : 'VMCS_GUEST_IDTR_LIMIT',
            0x00004814 : 'VMCS_GUEST_ES_ACCESS_RIGHTS',
            0x00004816 : 'VMCS_GUEST_CS_ACCESS_RIGHTS',
            0x00004818 : 'VMCS_GUEST_SS_ACCESS_RIGHTS',
            0x0000481A : 'VMCS_GUEST_DS_ACCESS_RIGHTS',
            0x0000481C : 'VMCS_GUEST_FS_ACCESS_RIGHTS',
            0x0000481E : 'VMCS_GUEST_GS_ACCESS_RIGHTS',
            0x00004820 : 'VMCS_GUEST_LDTR_ACCESS_RIGHTS',
            0x00004822 : 'VMCS_GUEST_TR_ACCESS_RIGHTS',
            0x00004824 : 'VMCS_GUEST_INTERRUPTIBILITY_STATE',
            0x00004826 : 'VMCS_GUEST_ACTIVITY_STATE',
            0x00004828 : 'VMCS_GUEST_SMBASE',
            0x0000482A : 'VMCS_GUEST_SYSENTER_CS',
            0x0000482E : 'VMCS_GUEST_PREEMPT_TIMER_VALUE',
            0x00004C00 : 'VMCS_SYSENTER_CS',
            0x00006000 : 'VMCS_CTRL_CR0_MASK',
            0x00006002 : 'VMCS_CTRL_CR4_MASK',
            0x00006004 : 'VMCS_CTRL_CR0_READ_SHADOW',
            0x00006006 : 'VMCS_CTRL_CR4_READ_SHADOW',
            0x00006008 : 'VMCS_CTRL_CR3_TARGET_VAL0',
            0x0000600A : 'VMCS_CTRL_CR3_TARGET_VAL1',
            0x0000600C : 'VMCS_CTRL_CR3_TARGET_VAL2',
            0x0000600E : 'VMCS_CTRL_CR3_TARGET_VAL3',
            0x00006400 : 'VMCS_EXIT_QUALIFICATION',
            0x00006402 : 'VMCS_IO_RCX',
            0x00006404 : 'VMCS_IO_RSX',
            0x00006406 : 'VMCS_IO_RDI',
            0x00006408 : 'VMCS_IO_RIP',
            0x0000640A : 'VMCS_EXIT_GUEST_LINEAR_ADDR',
            0x00006800 : 'VMCS_GUEST_CR0',
            0x00006802 : 'VMCS_GUEST_CR3',
            0x00006804 : 'VMCS_GUEST_CR4',
            0x00006806 : 'VMCS_GUEST_ES_BASE',
            0x00006808 : 'VMCS_GUEST_CS_BASE',
            0x0000680A : 'VMCS_GUEST_SS_BASE',
            0x0000680C : 'VMCS_GUEST_DS_BASE',
            0x0000680E : 'VMCS_GUEST_FS_BASE',
            0x00006810 : 'VMCS_GUEST_GS_BASE',
            0x00006812 : 'VMCS_GUEST_LDTR_BASE',
            0x00006814 : 'VMCS_GUEST_TR_BASE',
            0x00006816 : 'VMCS_GUEST_GDTR_BASE',
            0x00006818 : 'VMCS_GUEST_IDTR_BASE',
            0x0000681A : 'VMCS_GUEST_DR7',
            0x0000681C : 'VMCS_GUEST_RSP',
            0x0000681E : 'VMCS_GUEST_RIP',
            0x00006820 : 'VMCS_GUEST_RFLAGS',
            0x00006822 : 'VMCS_GUEST_PENDING_DEBUG_EXCEPTIONS',
            0x00006824 : 'VMCS_GUEST_SYSENTER_ESP',
            0x00006826 : 'VMCS_GUEST_SYSENTER_EIP',
            0x00006C00 : 'VMCS_HOST_CR0',
            0x00006C02 : 'VMCS_HOST_CR3',
            0x00006C04 : 'VMCS_HOST_CR4',
            0x00006C06 : 'VMCS_HOST_FS_BASE',
            0x00006C08 : 'VMCS_HOST_GS_BASE',
            0x00006C0A : 'VMCS_HOST_TR_BASE',
            0x00006C0C : 'VMCS_HOST_GDTR_BASE',
            0x00006C0E : 'VMCS_HOST_IDTR_BASE',
            0x00006C10 : 'VMCS_HOST_SYSENTER_ESP',
            0x00006C12 : 'VMCS_HOST_SYSENTER_EIP',
            0x00006C14 : 'VMCS_HOST_RSP',
            0x00006C16 : 'VMCS_HOST_RIP',

                    # new https://lore.kernel.org/patchwork/patch/1002950/.
            0x00002811 : 'VMCS_GUEST_PDPTR3_HIGH',
            0x00002812 : 'VMCS_GUEST_BNDCFGS',
            0x00002813 : 'VMCS_GUEST_BNDCFGS_HIGH',
            0x00002814 : 'VMCS_GUEST_RTIT_CTL',
            0x00002815 : 'VMCS_GUEST_RTIT_CTL_HIGH'
            }

        self.ia32_msr_db = {

            #
            # Hyper-V MSRs
            #
            0x40000000 : 'HV_X64_MSR_GUEST_OS_ID',
            0x40000001 : 'HV_X64_MSR_HYPERCALL',
            0x40000002 : 'HV_X64_MSR_VP_INDEX',
            0x40000003 : 'HV_X64_MSR_RESET',
            0x40000010 : 'HV_X64_MSR_VP_RUNTIME',
            0x40000020 : 'HV_X64_MSR_TIME_REF_COUNT',
            0x40000021 : 'HV_X64_MSR_REFERENCE_TSC',
            0x40000022 : 'HV_X64_MSR_TSC_FREQUENCY',
            0x40000023 : 'HV_X64_MSR_APIC_FREQUENCY',
            0x40000070 : 'HV_X64_MSR_EOI',
            0x40000071 : 'HV_X64_MSR_ICR',
            0x40000072 : 'HV_X64_MSR_TPR',
            0x40000073 : 'HV_X64_MSR_VP_ASSIST_PAGE',
            0x40000080 : 'HV_X64_MSR_SCONTROL',
            0x40000081 : 'HV_X64_MSR_SVERSION',
            0x40000082 : 'HV_X64_MSR_SIEFP',
            0x40000083 : 'HV_X64_MSR_SIMP',
            0x40000084 : 'HV_X64_MSR_EOM',
            0x40000090 : 'HV_X64_MSR_SINT0',
            0x40000091 : 'HV_X64_MSR_SINT1',
            0x40000092 : 'HV_X64_MSR_SINT2',
            0x40000093 : 'HV_X64_MSR_SINT3',
            0x40000094 : 'HV_X64_MSR_SINT4',
            0x40000095 : 'HV_X64_MSR_SINT5',
            0x40000096 : 'HV_X64_MSR_SINT6',
            0x40000097 : 'HV_X64_MSR_SINT7',
            0x40000098 : 'HV_X64_MSR_SINT8',
            0x40000099 : 'HV_X64_MSR_SINT9',
            0x4000009A : 'HV_X64_MSR_SINT10',
            0x4000009B : 'HV_X64_MSR_SINT11',
            0x4000009C : 'HV_X64_MSR_SINT12',
            0x4000009D : 'HV_X64_MSR_SINT13',
            0x4000009E : 'HV_X64_MSR_SINT14',
            0x4000009F : 'HV_X64_MSR_SINT15',
            0x400000B0 : 'HV_X64_MSR_STIMER0_CONFIG',
            0x400000B1 : 'HV_X64_MSR_STIMER0_COUNT',
            0x400000B2 : 'HV_X64_MSR_STIMER1_CONFIG',
            0x400000B3 : 'HV_X64_MSR_STIMER1_COUNT',
            0x400000B4 : 'HV_X64_MSR_STIMER2_CONFIG',
            0x400000B5 : 'HV_X64_MSR_STIMER2_COUNT',
            0x400000B6 : 'HV_X64_MSR_STIMER3_CONFIG',
            0x400000B7 : 'HV_X64_MSR_STIMER3_COUNT',
            0x400000C1 : 'HV_X64_MSR_POWER_STATE_TRIGGER_C1',
            0x400000C2 : 'HV_X64_MSR_POWER_STATE_TRIGGER_C2',
            0x400000C3 : 'HV_X64_MSR_POWER_STATE_TRIGGER_C3',
            0x400000D1 : 'HV_X64_MSR_POWER_STATE_CONFIG_C1',
            0x400000D2 : 'HV_X64_MSR_POWER_STATE_CONFIG_C2',
            0x400000D3 : 'HV_X64_MSR_POWER_STATE_CONFIG_C3',
            0x400000E0 : 'HV_X64_MSR_STATS_PARTITION_RETAIL_PAGE',
            0x400000E1 : 'HV_X64_MSR_STATS_PARTITION_INTERNAL_PAGE',
            0x400000E2 : 'HV_X64_MSR_STATS_VP_RETAIL_PAGE',
            0x400000E3 : 'HV_X64_MSR_STATS_VP_INTERNAL_PAGE',
            0x400000F0 : 'HV_X64_MSR_GUEST_IDLE',
            0x400000F1 : 'HV_X64_MSR_SYNTH_DEBUG_CONTROL',
            0x400000F2 : 'HV_X64_MSR_SYNTH_DEBUG_STATUS',
            0x400000F3 : 'HV_X64_MSR_SYNTH_DEBUG_SEND_BUFFER',
            0x400000F4 : 'HV_X64_MSR_SYNTH_DEBUG_RECEIVE_BUFFER',
            0x400000F5 : 'HV_X64_MSR_SYNTH_DEBUG_PENDING_BUFFER',
            0x40000100 : 'HV_X64_MSR_CRASH_P0',
            0x40000101 : 'HV_X64_MSR_CRASH_P1',
            0x40000102 : 'HV_X64_MSR_CRASH_P2',
            0x40000103 : 'HV_X64_MSR_CRASH_P3',
            0x40000104 : 'HV_X64_MSR_CRASH_P4',
            0x40000105 : 'HV_X64_MSR_CRASH_CTL',

            0x40000106 : 'HV_X64_MSR_REENLIGHTENMENT_CONTROL',
            0x40000107 : 'HV_X64_MSR_TSC_EMULATION_CONTROL',
            0x40000108 : 'HV_X64_MSR_TSC_EMULATION_STATUS',


            #
            # XEON MSRs
            #
            0x0000004E : 'MSR_XEON_D_PPIN_CTL',
            0x0000004F : 'MSR_XEON_D_PPIN',
            0x000000CE : 'MSR_XEON_D_PLATFORM_INFO',
            0x000000E2 : 'MSR_XEON_D_PKG_CST_CONFIG_CONTROL',
            0x00000179 : 'MSR_XEON_D_IA32_MCG_CAP',
            0x0000017D : 'MSR_XEON_D_SMM_MCA_CAP',
            0x000001A2 : 'MSR_XEON_D_TEMPERATURE_TARGET',
            0x000001AD : 'MSR_XEON_D_TURBO_RATIO_LIMIT',
            0x000001AE : 'MSR_XEON_D_TURBO_RATIO_LIMIT1',
            0x00000606 : 'MSR_XEON_D_RAPL_POWER_UNIT',
            0x00000618 : 'MSR_XEON_D_DRAM_POWER_LIMIT',
            0x00000619 : 'MSR_XEON_D_DRAM_ENERGY_STATUS',
            0x0000061B : 'MSR_XEON_D_DRAM_PERF_STATUS',
            0x0000061C : 'MSR_XEON_D_DRAM_POWER_INFO',
            0x00000620 : 'MSR_XEON_D_MSRUNCORE_RATIO_LIMIT',
            0x00000639 : 'MSR_XEON_D_PP0_ENERGY_STATUS',
            0x00000690 : 'MSR_XEON_D_CORE_PERF_LIMIT_REASONS',
            0x00000C8D : 'MSR_XEON_D_IA32_QM_EVTSEL',
            0x00000C8F : 'MSR_XEON_D_IA32_PQR_ASSOC',
            0x000001AC : 'MSR_XEON_D_TURBO_RATIO_LIMIT3',
            0x00000C81 : 'MSR_XEON_D_IA32_L3_QOS_CFG',
            0x00000C90 : 'MSR_XEON_D_IA32_L3_QOS_MASK_0',
            0x00000C91 : 'MSR_XEON_D_IA32_L3_QOS_MASK_1',
            0x00000C92 : 'MSR_XEON_D_IA32_L3_QOS_MASK_2',
            0x00000C93 : 'MSR_XEON_D_IA32_L3_QOS_MASK_3',
            0x00000C94 : 'MSR_XEON_D_IA32_L3_QOS_MASK_4',
            0x00000C95 : 'MSR_XEON_D_IA32_L3_QOS_MASK_5',
            0x00000C96 : 'MSR_XEON_D_IA32_L3_QOS_MASK_6',
            0x00000C97 : 'MSR_XEON_D_IA32_L3_QOS_MASK_7',
            0x00000C98 : 'MSR_XEON_D_IA32_L3_QOS_MASK_8',
            0x00000C99 : 'MSR_XEON_D_IA32_L3_QOS_MASK_9',
            0x00000C9A : 'MSR_XEON_D_IA32_L3_QOS_MASK_10',
            0x00000C9B : 'MSR_XEON_D_IA32_L3_QOS_MASK_11',
            0x00000C9C : 'MSR_XEON_D_IA32_L3_QOS_MASK_12',
            0x00000C9D : 'MSR_XEON_D_IA32_L3_QOS_MASK_13',
            0x00000C9E : 'MSR_XEON_D_IA32_L3_QOS_MASK_14',
            0x00000C9F : 'MSR_XEON_D_IA32_L3_QOS_MASK_15',



            # 
            # QQQQQQQQQQ
            #

            0x00000048 : 'MSR_SPEC_CTRL',
            0x00000049 : 'MSR_PRED_CMD',

            0x000001C8 : 'MSR_LBR_SELECT',
            0x000001C9 : 'MSR_LBR_TOS',
            0x00000680 : 'MSR_LBR_NHM_FROM',
            0x000006C0 : 'MSR_LBR_NHM_TO',
            0x00000040 : 'MSR_LBR_CORE_FROM',
            0x00000060 : 'MSR_LBR_CORE_TO',


            0x000001DB : 'MSR_IA32_LASTBRANCHFROMIP',
            0x000001DC : 'MSR_IA32_LASTBRANCHTOIP',
            0x000001DD : 'MSR_IA32_LASTINTFROMIP',
            0x000001DE : 'MSR_IA32_LASTINTTOIP',

            # https://github.com/collectd/collectd/blob/master/src/msr-index.h
            #
            # C-state Residency Counters
            #
            0x000003F8 : 'MSR_PKG_C3_RESIDENCY',
            0x000003F9 : 'MSR_PKG_C6_RESIDENCY',
            0x000003FA : 'MSR_ATOM_PKG_C6_RESIDENCY',
            0x000003FA : 'MSR_PKG_C7_RESIDENCY',
            0x000003FC : 'MSR_CORE_C3_RESIDENCY',
            0x000003FD : 'MSR_CORE_C6_RESIDENCY',
            0x000003FE : 'MSR_CORE_C7_RESIDENCY',
            0x000003FF : 'MSR_KNL_CORE_C6_RESIDENCY',
            0x0000060D : 'MSR_PKG_C2_RESIDENCY',
            0x00000630 : 'MSR_PKG_C8_RESIDENCY',
            0x00000631 : 'MSR_PKG_C9_RESIDENCY',
            0x00000632 : 'MSR_PKG_C10_RESIDENCY',


            # https://sites.uclouvain.be/SystInfo/usr/include/asm/msr-index.h.html

            #
            # Centaur-Hauls/IDT defined MSRs.
            #
            0x00000107 : 'MSR_IDT_FCR1',
            0x00000108 : 'MSR_IDT_FCR2',
            0x00000109 : 'MSR_IDT_FCR3',
            0x0000010A : 'MSR_IDT_FCR4',
            0x00000110 : 'MSR_IDT_MCR0',
            0x00000111 : 'MSR_IDT_MCR1',
            0x00000112 : 'MSR_IDT_MCR2',
            0x00000113 : 'MSR_IDT_MCR3',
            0x00000114 : 'MSR_IDT_MCR4',
            0x00000115 : 'MSR_IDT_MCR5',
            0x00000116 : 'MSR_IDT_MCR6',
            0x00000117 : 'MSR_IDT_MCR7',
            0x00000120 : 'MSR_IDT_MCR_CTRL',


            0x000000CD : 'MSR_FSB_FREQ',
            0x00000119 : 'MSR_IA32_BBL_CR_CTL',

            # 
            # QQQQQQQQQQ
            #

            # AMD-V MSRs
            0xC0010114 : 'MSR_VM_CR',
            0xC0010115 : 'MSR_VM_IGNNE',
            0xC0010117 : 'MSR_VM_HSAVE_PA',



            ###########################################################
            # start https://github.com/wbenny/ia32-doc

            # @defgroup intel_manual \
            #           Intel Manual
            # BEGIN
            # @defgroup model_specific_registers \
            #           Model Specific Registers
            # BEGIN
            # @defgroup ia32_p5_mc \
            #           IA32_P5_MC_(x)
            # BEGIN
            0x00000000 : 'IA32_P5_MC_ADDR',
            0x00000001 : 'IA32_P5_MC_TYPE',
            # END

            0x00000006 : 'IA32_MONITOR_FILTER_SIZE',
            0x00000010 : 'IA32_TIME_STAMP_COUNTER',
            0x00000017 : 'IA32_PLATFORM_ID',
            0x0000001B : 'IA32_APIC_BASE',
            0x0000003A : 'IA32_FEATURE_CONTROL',
            0x0000003B : 'IA32_TSC_ADJUST',
            0x00000079 : 'IA32_BIOS_UPDT_TRIG',
            0x0000008B : 'IA32_BIOS_SIGN_ID',
            # @defgroup ia32_sgxlepubkeyhash \
            #           IA32_SGXLEPUBKEYHASH[(64*n+63):(64*n)]
            # BEGIN
            0x0000008C : 'IA32_SGXLEPUBKEYHASH0',
            0x0000008D : 'IA32_SGXLEPUBKEYHASH1',
            0x0000008E : 'IA32_SGXLEPUBKEYHASH2',
            0x0000008F : 'IA32_SGXLEPUBKEYHASH3',
            # END

            0x0000009B : 'IA32_SMM_MONITOR_CTL',
            0x0000009E : 'IA32_SMBASE',
            # @defgroup ia32_pmc \
            #           IA32_PMC(n)
            # BEGIN
            0x000000C1 : 'IA32_PMC0',
            0x000000C2 : 'IA32_PMC1',
            0x000000C3 : 'IA32_PMC2',
            0x000000C4 : 'IA32_PMC3',
            0x000000C5 : 'IA32_PMC4',
            0x000000C6 : 'IA32_PMC5',
            0x000000C7 : 'IA32_PMC6',
            0x000000C8 : 'IA32_PMC7',
            # END

            0x000000E7 : 'IA32_MPERF',
            0x000000E8 : 'IA32_APERF',
            0x000000FE : 'IA32_MTRRCAP',
            0x00000174 : 'IA32_SYSENTER_CS',
            0x00000175 : 'IA32_SYSENTER_ESP',
            0x00000176 : 'IA32_SYSENTER_EIP',
            0x00000179 : 'IA32_MCG_CAP',
            0x0000017A : 'IA32_MCG_STATUS',
            0x0000017B : 'IA32_MCG_CTL',
            # @defgroup ia32_perfevtsel \
            #           IA32_PERFEVTSEL(n)
            # BEGIN
            0x00000186 : 'IA32_PERFEVTSEL0',
            0x00000187 : 'IA32_PERFEVTSEL1',
            0x00000188 : 'IA32_PERFEVTSEL2',
            0x00000189 : 'IA32_PERFEVTSEL3',
            # END

            0x00000198 : 'IA32_PERF_STATUS',
            0x00000199 : 'IA32_PERF_CTL',
            0x0000019A : 'IA32_CLOCK_MODULATION',
            0x0000019B : 'IA32_THERM_INTERRUPT',
            0x0000019C : 'IA32_THERM_STATUS',
            0x000001A0 : 'IA32_MISC_ENABLE',
            0x000001B0 : 'IA32_ENERGY_PERF_BIAS',
            0x000001B1 : 'IA32_PACKAGE_THERM_STATUS',
            0x000001B2 : 'IA32_PACKAGE_THERM_INTERRUPT',
            0x000001D9 : 'IA32_DEBUGCTL',
            0x000001F2 : 'IA32_SMRR_PHYSBASE',
            0x000001F3 : 'IA32_SMRR_PHYSMASK',
            0x000001F8 : 'IA32_PLATFORM_DCA_CAP',
            0x000001F9 : 'IA32_CPU_DCA_CAP',
            0x000001FA : 'IA32_DCA_0_CAP',
            # @defgroup ia32_mtrr_physbase \
            #           IA32_MTRR_PHYSBASE(n)
            # BEGIN
            0x00000200 : 'IA32_MTRR_PHYSBASE0',
            0x00000202 : 'IA32_MTRR_PHYSBASE1',
            0x00000204 : 'IA32_MTRR_PHYSBASE2',
            0x00000206 : 'IA32_MTRR_PHYSBASE3',
            0x00000208 : 'IA32_MTRR_PHYSBASE4',
            0x0000020A : 'IA32_MTRR_PHYSBASE5',
            0x0000020C : 'IA32_MTRR_PHYSBASE6',
            0x0000020E : 'IA32_MTRR_PHYSBASE7',
            0x00000210 : 'IA32_MTRR_PHYSBASE8',
            0x00000212 : 'IA32_MTRR_PHYSBASE9',
            # END

            # @defgroup ia32_mtrr_physmask \
            #           IA32_MTRR_PHYSMASK(n)
            # BEGIN
            0x00000201 : 'IA32_MTRR_PHYSMASK0',
            0x00000203 : 'IA32_MTRR_PHYSMASK1',
            0x00000205 : 'IA32_MTRR_PHYSMASK2',
            0x00000207 : 'IA32_MTRR_PHYSMASK3',
            0x00000209 : 'IA32_MTRR_PHYSMASK4',
            0x0000020B : 'IA32_MTRR_PHYSMASK5',
            0x0000020D : 'IA32_MTRR_PHYSMASK6',
            0x0000020F : 'IA32_MTRR_PHYSMASK7',
            0x00000211 : 'IA32_MTRR_PHYSMASK8',
            0x00000213 : 'IA32_MTRR_PHYSMASK9',
            # END

            # @defgroup ia32_mtrr_fix \
            #           IA32_MTRR_FIX(x)
            # BEGIN
            # @defgroup ia32_mtrr_fix64k \
            #           IA32_MTRR_FIX64K(x)
            # BEGIN
            0x00000000 : 'IA32_MTRR_FIX64K_BASE',
            0x00010000 : 'IA32_MTRR_FIX64K_SIZE',
            0x00000250 : 'IA32_MTRR_FIX64K_00000',
            # END

            # @defgroup ia32_mtrr_fix16k \
            #           IA32_MTRR_FIX16K(x)
            # BEGIN
            0x00080000 : 'IA32_MTRR_FIX16K_BASE',
            0x00004000 : 'IA32_MTRR_FIX16K_SIZE',
            0x00000258 : 'IA32_MTRR_FIX16K_80000',
            0x00000259 : 'IA32_MTRR_FIX16K_A0000',
            # END

            # @defgroup ia32_mtrr_fix4k \
            #           IA32_MTRR_FIX4K(x)
            # BEGIN
            0x000C0000 : 'IA32_MTRR_FIX4K_BASE',
            0x00001000 : 'IA32_MTRR_FIX4K_SIZE',
            0x00000268 : 'IA32_MTRR_FIX4K_C0000',
            0x00000269 : 'IA32_MTRR_FIX4K_C8000',
            0x0000026A : 'IA32_MTRR_FIX4K_D0000',
            0x0000026B : 'IA32_MTRR_FIX4K_D8000',
            0x0000026C : 'IA32_MTRR_FIX4K_E0000',
            0x0000026D : 'IA32_MTRR_FIX4K_E8000',
            0x0000026E : 'IA32_MTRR_FIX4K_F0000',
            0x0000026F : 'IA32_MTRR_FIX4K_F8000',
            # END


            # """
            # >>> hex(((1 + 2 + 8) * 8))
            # '0x58'
            # """
            # '((1 + 2 + 8) * 8) : 'IA32_MTRR_FIX_COUNT',

            0x00000058 : 'IA32_MTRR_FIX_COUNT',
            0x000000FF : 'IA32_MTRR_VARIABLE_COUNT',
            # """>>> hex(0x00000058 + 0x000000FF)
            # '0x157'
            # """
            # '(IA32_MTRR_FIX_COUNT + IA32_MTRR_VARIABLE_COUNT) : 'IA32_MTRR_COUNT',
            0x00000157 : 'IA32_MTRR_COUNT',

            # END

            0x00000277 : 'IA32_PAT',
            # @defgroup ia32_mc_ctl2 \
            #           IA32_MC(i)_CTL2
            # BEGIN
            0x00000280 : 'IA32_MC0_CTL2',
            0x00000281 : 'IA32_MC1_CTL2',
            0x00000282 : 'IA32_MC2_CTL2',
            0x00000283 : 'IA32_MC3_CTL2',
            0x00000284 : 'IA32_MC4_CTL2',
            0x00000285 : 'IA32_MC5_CTL2',
            0x00000286 : 'IA32_MC6_CTL2',
            0x00000287 : 'IA32_MC7_CTL2',
            0x00000288 : 'IA32_MC8_CTL2',
            0x00000289 : 'IA32_MC9_CTL2',
            0x0000028A : 'IA32_MC10_CTL2',
            0x0000028B : 'IA32_MC11_CTL2',
            0x0000028C : 'IA32_MC12_CTL2',
            0x0000028D : 'IA32_MC13_CTL2',
            0x0000028E : 'IA32_MC14_CTL2',
            0x0000028F : 'IA32_MC15_CTL2',
            0x00000290 : 'IA32_MC16_CTL2',
            0x00000291 : 'IA32_MC17_CTL2',
            0x00000292 : 'IA32_MC18_CTL2',
            0x00000293 : 'IA32_MC19_CTL2',
            0x00000294 : 'IA32_MC20_CTL2',
            0x00000295 : 'IA32_MC21_CTL2',
            0x00000296 : 'IA32_MC22_CTL2',
            0x00000297 : 'IA32_MC23_CTL2',
            0x00000298 : 'IA32_MC24_CTL2',
            0x00000299 : 'IA32_MC25_CTL2',
            0x0000029A : 'IA32_MC26_CTL2',
            0x0000029B : 'IA32_MC27_CTL2',
            0x0000029C : 'IA32_MC28_CTL2',
            0x0000029D : 'IA32_MC29_CTL2',
            0x0000029E : 'IA32_MC30_CTL2',
            0x0000029F : 'IA32_MC31_CTL2',
            # END

            0x000002FF : 'IA32_MTRR_DEF_TYPE',
            # @defgroup ia32_fixed_ctr \
            #           IA32_FIXED_CTR(n)
            # BEGIN
            0x00000309 : 'IA32_FIXED_CTR0',
            0x0000030A : 'IA32_FIXED_CTR1',
            0x0000030B : 'IA32_FIXED_CTR2',
            # END

            0x00000345 : 'IA32_PERF_CAPABILITIES',
            0x0000038D : 'IA32_FIXED_CTR_CTRL',
            0x0000038E : 'IA32_PERF_GLOBAL_STATUS',
            0x0000038F : 'IA32_PERF_GLOBAL_CTRL',
            0x00000390 : 'IA32_PERF_GLOBAL_STATUS_RESET',
            0x00000391 : 'IA32_PERF_GLOBAL_STATUS_SET',
            0x00000392 : 'IA32_PERF_GLOBAL_INUSE',
            0x000003F1 : 'IA32_PEBS_ENABLE',
            # @defgroup ia32_mc_ctl \
            #           IA32_MC(i)_CTL
            # BEGIN
            0x00000400 : 'IA32_MC0_CTL',
            0x00000404 : 'IA32_MC1_CTL',
            0x00000408 : 'IA32_MC2_CTL',
            0x0000040C : 'IA32_MC3_CTL',
            0x00000410 : 'IA32_MC4_CTL',
            0x00000414 : 'IA32_MC5_CTL',
            0x00000418 : 'IA32_MC6_CTL',
            0x0000041C : 'IA32_MC7_CTL',
            0x00000420 : 'IA32_MC8_CTL',
            0x00000424 : 'IA32_MC9_CTL',
            0x00000428 : 'IA32_MC10_CTL',
            0x0000042C : 'IA32_MC11_CTL',
            0x00000430 : 'IA32_MC12_CTL',
            0x00000434 : 'IA32_MC13_CTL',
            0x00000438 : 'IA32_MC14_CTL',
            0x0000043C : 'IA32_MC15_CTL',
            0x00000440 : 'IA32_MC16_CTL',
            0x00000444 : 'IA32_MC17_CTL',
            0x00000448 : 'IA32_MC18_CTL',
            0x0000044C : 'IA32_MC19_CTL',
            0x00000450 : 'IA32_MC20_CTL',
            0x00000454 : 'IA32_MC21_CTL',
            0x00000458 : 'IA32_MC22_CTL',
            0x0000045C : 'IA32_MC23_CTL',
            0x00000460 : 'IA32_MC24_CTL',
            0x00000464 : 'IA32_MC25_CTL',
            0x00000468 : 'IA32_MC26_CTL',
            0x0000046C : 'IA32_MC27_CTL',
            0x00000470 : 'IA32_MC28_CTL',
            # END

            # @defgroup ia32_mc_status \
            #           IA32_MC(i)_STATUS
            # BEGIN
            0x00000401 : 'IA32_MC0_STATUS',
            0x00000405 : 'IA32_MC1_STATUS',
            0x00000409 : 'IA32_MC2_STATUS',
            0x0000040D : 'IA32_MC3_STATUS',
            0x00000411 : 'IA32_MC4_STATUS',
            0x00000415 : 'IA32_MC5_STATUS',
            0x00000419 : 'IA32_MC6_STATUS',
            0x0000041D : 'IA32_MC7_STATUS',
            0x00000421 : 'IA32_MC8_STATUS',
            0x00000425 : 'IA32_MC9_STATUS',
            0x00000429 : 'IA32_MC10_STATUS',
            0x0000042D : 'IA32_MC11_STATUS',
            0x00000431 : 'IA32_MC12_STATUS',
            0x00000435 : 'IA32_MC13_STATUS',
            0x00000439 : 'IA32_MC14_STATUS',
            0x0000043D : 'IA32_MC15_STATUS',
            0x00000441 : 'IA32_MC16_STATUS',
            0x00000445 : 'IA32_MC17_STATUS',
            0x00000449 : 'IA32_MC18_STATUS',
            0x0000044D : 'IA32_MC19_STATUS',
            0x00000451 : 'IA32_MC20_STATUS',
            0x00000455 : 'IA32_MC21_STATUS',
            0x00000459 : 'IA32_MC22_STATUS',
            0x0000045D : 'IA32_MC23_STATUS',
            0x00000461 : 'IA32_MC24_STATUS',
            0x00000465 : 'IA32_MC25_STATUS',
            0x00000469 : 'IA32_MC26_STATUS',
            0x0000046D : 'IA32_MC27_STATUS',
            0x00000471 : 'IA32_MC28_STATUS',
            # END

            # @defgroup ia32_mc_addr \
            #           IA32_MC(i)_ADDR
            # BEGIN
            0x00000402 : 'IA32_MC0_ADDR',
            0x00000406 : 'IA32_MC1_ADDR',
            0x0000040A : 'IA32_MC2_ADDR',
            0x0000040E : 'IA32_MC3_ADDR',
            0x00000412 : 'IA32_MC4_ADDR',
            0x00000416 : 'IA32_MC5_ADDR',
            0x0000041A : 'IA32_MC6_ADDR',
            0x0000041E : 'IA32_MC7_ADDR',
            0x00000422 : 'IA32_MC8_ADDR',
            0x00000426 : 'IA32_MC9_ADDR',
            0x0000042A : 'IA32_MC10_ADDR',
            0x0000042E : 'IA32_MC11_ADDR',
            0x00000432 : 'IA32_MC12_ADDR',
            0x00000436 : 'IA32_MC13_ADDR',
            0x0000043A : 'IA32_MC14_ADDR',
            0x0000043E : 'IA32_MC15_ADDR',
            0x00000442 : 'IA32_MC16_ADDR',
            0x00000446 : 'IA32_MC17_ADDR',
            0x0000044A : 'IA32_MC18_ADDR',
            0x0000044E : 'IA32_MC19_ADDR',
            0x00000452 : 'IA32_MC20_ADDR',
            0x00000456 : 'IA32_MC21_ADDR',
            0x0000045A : 'IA32_MC22_ADDR',
            0x0000045E : 'IA32_MC23_ADDR',
            0x00000462 : 'IA32_MC24_ADDR',
            0x00000466 : 'IA32_MC25_ADDR',
            0x0000046A : 'IA32_MC26_ADDR',
            0x0000046E : 'IA32_MC27_ADDR',
            0x00000472 : 'IA32_MC28_ADDR',
            # END

            # @defgroup ia32_mc_misc \
            #           IA32_MC(i)_MISC
            # BEGIN
            0x00000403 : 'IA32_MC0_MISC',
            0x00000407 : 'IA32_MC1_MISC',
            0x0000040B : 'IA32_MC2_MISC',
            0x0000040F : 'IA32_MC3_MISC',
            0x00000413 : 'IA32_MC4_MISC',
            0x00000417 : 'IA32_MC5_MISC',
            0x0000041B : 'IA32_MC6_MISC',
            0x0000041F : 'IA32_MC7_MISC',
            0x00000423 : 'IA32_MC8_MISC',
            0x00000427 : 'IA32_MC9_MISC',
            0x0000042B : 'IA32_MC10_MISC',
            0x0000042F : 'IA32_MC11_MISC',
            0x00000433 : 'IA32_MC12_MISC',
            0x00000437 : 'IA32_MC13_MISC',
            0x0000043B : 'IA32_MC14_MISC',
            0x0000043F : 'IA32_MC15_MISC',
            0x00000443 : 'IA32_MC16_MISC',
            0x00000447 : 'IA32_MC17_MISC',
            0x0000044B : 'IA32_MC18_MISC',
            0x0000044F : 'IA32_MC19_MISC',
            0x00000453 : 'IA32_MC20_MISC',
            0x00000457 : 'IA32_MC21_MISC',
            0x0000045B : 'IA32_MC22_MISC',
            0x0000045F : 'IA32_MC23_MISC',
            0x00000463 : 'IA32_MC24_MISC',
            0x00000467 : 'IA32_MC25_MISC',
            0x0000046B : 'IA32_MC26_MISC',
            0x0000046F : 'IA32_MC27_MISC',
            0x00000473 : 'IA32_MC28_MISC',
            # END

            0x00000480 : 'IA32_VMX_BASIC',
            0x00000481 : 'IA32_VMX_PINBASED_CTLS',
            0x00000482 : 'IA32_VMX_PROCBASED_CTLS',
            0x00000483 : 'IA32_VMX_EXIT_CTLS',
            0x00000484 : 'IA32_VMX_ENTRY_CTLS',
            0x00000485 : 'IA32_VMX_MISC',
            0x00000486 : 'IA32_VMX_CR0_FIXED0',
            0x00000487 : 'IA32_VMX_CR0_FIXED1',
            0x00000488 : 'IA32_VMX_CR4_FIXED0',
            0x00000489 : 'IA32_VMX_CR4_FIXED1',
            0x0000048A : 'IA32_VMX_VMCS_ENUM',
            0x0000048B : 'IA32_VMX_PROCBASED_CTLS2',
            0x0000048C : 'IA32_VMX_EPT_VPID_CAP',
            # @defgroup ia32_vmx_true_ctls \
            #           IA32_VMX_TRUE_(x)_CTLS
            # BEGIN
            0x0000048D : 'IA32_VMX_TRUE_PINBASED_CTLS',
            0x0000048E : 'IA32_VMX_TRUE_PROCBASED_CTLS',
            0x0000048F : 'IA32_VMX_TRUE_EXIT_CTLS',
            0x00000490 : 'IA32_VMX_TRUE_ENTRY_CTLS',
            # END

            0x00000491 : 'IA32_VMX_VMFUNC',
            # @defgroup ia32_a_pmc \
            #           IA32_A_PMC(n)
            # BEGIN
            0x000004C1 : 'IA32_A_PMC0',
            0x000004C2 : 'IA32_A_PMC1',
            0x000004C3 : 'IA32_A_PMC2',
            0x000004C4 : 'IA32_A_PMC3',
            0x000004C5 : 'IA32_A_PMC4',
            0x000004C6 : 'IA32_A_PMC5',
            0x000004C7 : 'IA32_A_PMC6',
            0x000004C8 : 'IA32_A_PMC7',
            # END

            0x000004D0 : 'IA32_MCG_EXT_CTL',
            0x00000500 : 'IA32_SGX_SVN_STATUS',
            0x00000560 : 'IA32_RTIT_OUTPUT_BASE',
            0x00000561 : 'IA32_RTIT_OUTPUT_MASK_PTRS',
            0x00000570 : 'IA32_RTIT_CTL',
            0x00000571 : 'IA32_RTIT_STATUS',
            0x00000572 : 'IA32_RTIT_CR3_MATCH',
            # @defgroup ia32_rtit_addr \
            #           IA32_RTIT_ADDR(x)
            # BEGIN
            # @defgroup ia32_rtit_addr_a \
            #           IA32_RTIT_ADDR(n)_A
            # BEGIN
            0x00000580 : 'IA32_RTIT_ADDR0_A',
            0x00000582 : 'IA32_RTIT_ADDR1_A',
            0x00000584 : 'IA32_RTIT_ADDR2_A',
            0x00000586 : 'IA32_RTIT_ADDR3_A',
            # END

            # @defgroup ia32_rtit_addr_b \
            #           IA32_RTIT_ADDR(n)_B
            # BEGIN
            0x00000581 : 'IA32_RTIT_ADDR0_B',
            0x00000583 : 'IA32_RTIT_ADDR1_B',
            0x00000585 : 'IA32_RTIT_ADDR2_B',
            0x00000587 : 'IA32_RTIT_ADDR3_B',
            # END

            # END

            0x00000600 : 'IA32_DS_AREA',
            0x000006E0 : 'IA32_TSC_DEADLINE',
            0x00000770 : 'IA32_PM_ENABLE',
            0x00000771 : 'IA32_HWP_CAPABILITIES',
            0x00000772 : 'IA32_HWP_REQUEST_PKG',
            0x00000773 : 'IA32_HWP_INTERRUPT',
            0x00000774 : 'IA32_HWP_REQUEST',
            0x00000777 : 'IA32_HWP_STATUS',
            0x00000802 : 'IA32_X2APIC_APICID',
            0x00000803 : 'IA32_X2APIC_VERSION',
            0x00000808 : 'IA32_X2APIC_TPR',
            0x0000080A : 'IA32_X2APIC_PPR',
            0x0000080B : 'IA32_X2APIC_EOI',
            0x0000080D : 'IA32_X2APIC_LDR',
            0x0000080F : 'IA32_X2APIC_SIVR',
            # @defgroup ia32_x2apic_isr \
            #           IA32_X2APIC_ISR(n)
            # BEGIN
            0x00000810 : 'IA32_X2APIC_ISR0',
            0x00000811 : 'IA32_X2APIC_ISR1',
            0x00000812 : 'IA32_X2APIC_ISR2',
            0x00000813 : 'IA32_X2APIC_ISR3',
            0x00000814 : 'IA32_X2APIC_ISR4',
            0x00000815 : 'IA32_X2APIC_ISR5',
            0x00000816 : 'IA32_X2APIC_ISR6',
            0x00000817 : 'IA32_X2APIC_ISR7',
            # END

            # @defgroup ia32_x2apic_tmr \
            #           IA32_X2APIC_TMR(n)
            # BEGIN
            0x00000818 : 'IA32_X2APIC_TMR0',
            0x00000819 : 'IA32_X2APIC_TMR1',
            0x0000081A : 'IA32_X2APIC_TMR2',
            0x0000081B : 'IA32_X2APIC_TMR3',
            0x0000081C : 'IA32_X2APIC_TMR4',
            0x0000081D : 'IA32_X2APIC_TMR5',
            0x0000081E : 'IA32_X2APIC_TMR6',
            0x0000081F : 'IA32_X2APIC_TMR7',
            # END

            # @defgroup ia32_x2apic_irr \
            #           IA32_X2APIC_IRR(n)
            # BEGIN
            0x00000820 : 'IA32_X2APIC_IRR0',
            0x00000821 : 'IA32_X2APIC_IRR1',
            0x00000822 : 'IA32_X2APIC_IRR2',
            0x00000823 : 'IA32_X2APIC_IRR3',
            0x00000824 : 'IA32_X2APIC_IRR4',
            0x00000825 : 'IA32_X2APIC_IRR5',
            0x00000826 : 'IA32_X2APIC_IRR6',
            0x00000827 : 'IA32_X2APIC_IRR7',
            # END

            0x00000828 : 'IA32_X2APIC_ESR',
            0x0000082F : 'IA32_X2APIC_LVT_CMCI',
            0x00000830 : 'IA32_X2APIC_ICR',
            0x00000832 : 'IA32_X2APIC_LVT_TIMER',
            0x00000833 : 'IA32_X2APIC_LVT_THERMAL',
            0x00000834 : 'IA32_X2APIC_LVT_PMI',
            0x00000835 : 'IA32_X2APIC_LVT_LINT0',
            0x00000836 : 'IA32_X2APIC_LVT_LINT1',
            0x00000837 : 'IA32_X2APIC_LVT_ERROR',
            0x00000838 : 'IA32_X2APIC_INIT_COUNT',
            0x00000839 : 'IA32_X2APIC_CUR_COUNT',
            0x0000083E : 'IA32_X2APIC_DIV_CONF',
            0x0000083F : 'IA32_X2APIC_SELF_IPI',
            0x00000C80 : 'IA32_DEBUG_INTERFACE',
            0x00000C81 : 'IA32_L3_QOS_CFG',
            0x00000C82 : 'IA32_L2_QOS_CFG',
            0x00000C8D : 'IA32_QM_EVTSEL',
            0x00000C8E : 'IA32_QM_CTR',
            0x00000C8F : 'IA32_PQR_ASSOC',
            0x00000D90 : 'IA32_BNDCFGS',
            0x00000DA0 : 'IA32_XSS',
            0x00000DB0 : 'IA32_PKG_HDC_CTL',
            0x00000DB1 : 'IA32_PM_CTL1',
            0x00000DB2 : 'IA32_THREAD_STALL',
            0xC0000080 : 'IA32_EFER',
            0xC0000081 : 'IA32_STAR',
            0xC0000082 : 'IA32_LSTAR',
            0xC0000083 : 'IA32_CSTAR',
            0xC0000084 : 'IA32_FMASK',
            0xC0000100 : 'IA32_FS_BASE',
            0xC0000101 : 'IA32_GS_BASE',
            0xC0000102 : 'IA32_KERNEL_GS_BASE',
            0xC0000103 : 'IA32_TSC_AUX',
            # END

            # end https://github.com/wbenny/ia32-doc
            #######################################################
        }

    def GetIa32VmcsDb(self):
        return self.ia32_vmcs_db

    def GetIa32MsrDb(self):
        return self.ia32_msr_db

# ------------------------------------------------------------------------------------- #
def PrettyHex(v):
    return '0x{0:08X}'.format(v)

def GetJumpAddr(call_addr, func_ea):
    func_addr = idc.LocByName(idc.GetFunctionName(func_ea))
    ret = int(call_addr - func_addr)
    return hex(ret).replace("0x", "")

def GetFunctionNamePlusRva(inst_ea):
        jump_addr = GetJumpAddr(inst_ea, inst_ea)
        return idc.GetFunctionName(inst_ea)+ '+' + jump_addr


def AddressToHex(ea):
    return hex(ea).rstrip('L')

def GetDelimeter(input_str,max_len = 30, delim_char = " "):
    if input_str == None:
        return delim_char * (max_len - len('NoneNone'))
    else:
        return delim_char * (max_len - len(input_str))

def GetValueWithDelimeter(input_str,max_len = 30, delim_char = " "):
    if input_str == None:
        return GetDelimeter(input_str,max_len,delim_char)
    return input_str+GetDelimeter(input_str,max_len,delim_char)

# ------------------------------------------------------------------------------------- #
class Config (object):
    PLUGIN_TITLTE = "IA32 VMX Helper"
    COLUMN_NAMES  = ["Address", "Mnemonic",  "Value_Hex", "Value_Name", "Function_Name"]
    COLUMN_SIZE   = [20,     10,         18, 40 , 30]
    PLUGIN_COLUMNS = [list(c) for c in
                      zip(COLUMN_NAMES, COLUMN_SIZE)]

    ROW_ITEM = collections.namedtuple("RowItem",  COLUMN_NAMES)

class PluginPage(idaapi.Choose2):
    
    def __init__(self, title, columns, items, icon, embedded=False):

        idaapi.Choose2.__init__(self, title, columns, embedded=embedded)
        self.items = items
        self.icon = icon

    def GetItems(self):
        return self.items

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

    def OnClose(self):
        pass

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

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

    def OnSelectLine(self, index):
        try:
            row = self.items[index][0]
            addr = int(row,16)
            idc.Jump(addr)
        except Exception, e:
            idc.Warning('Exception : Selected item {}  = {}'.format(index,str(e))) 
            pass





class ReportEntry(object):
    
    def __init__(self, addr,mnemonic,value_hex,value_name,func_name):
        self.addr = addr
        self.mnemonic = mnemonic
        self.value_hex = value_hex
        self.value_name = value_name
        self.func_name = func_name

    def get_row(self):

        addr = "0x%x" % (self.addr)
        return list(Config.ROW_ITEM(addr, self.mnemonic, self.value_hex,self.value_name,self.func_name))

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

def GetEnumItemByConst(const_value,constants):
    if constants == None:
        # idc.Warning('constants not loaded')
        return None
    for item_name,item_constant, enum_id in constants:
        if item_constant == const_value:
            return [item_name,enum_id]
    return None

class SymbolicConstant(object):
    def __init__(self,enum_name):
        self.enum_name = enum_name
        self.result = []
        self.count = 0

    def LoadEnums(self):
        enum_id = idc.GetEnum(self.enum_name)
        enum_count = idc.GetEnumSize(enum_id)
        
        if enum_id == idaapi.BADADDR:
            return None

        item_constant = idc.GetFirstConst(enum_id, -1)
        item_name     = idc.GetConstName(idc.GetConstEx(enum_id, item_constant, 0, -1))

        self.result.append([item_name,item_constant, enum_id])

        for i in range(enum_count):
            item_constant = idc.GetNextConst(enum_id, item_constant, -1)
            item_name = idc.GetConstName(idc.GetConstEx(enum_id, item_constant, 0, -1))

            self.result.append([item_name,item_constant, enum_id])

        return self.result
# ------------------------------------------------------------------------------------- #
OPERAND_INDEX_0 = 0
OPERAND_INDEX_1 = 1

class Instruction_t(object):
    def __init__(self,inst_ea):
        self.inst_ea = inst_ea
        self.mnemonic = idc.GetMnem(self.inst_ea)
        self.operand_0 = idc.GetOpnd(self.inst_ea, 0)
        self.operand_1 = idc.GetOpnd(self.inst_ea, 1)

        self.operand_value_0 = idc.GetOperandValue(self.inst_ea,0)
        self.operand_value_1 = idc.GetOperandValue(self.inst_ea,1)

    def GetOperandValue_32bit(self,operand_index):
        return idc.GetOperandValue(self.inst_ea,operand_index) & 0xffffffff

    def GetOperandType(self,operand_index):
        return idc.GetOpType(self.inst_ea,operand_index)
    
    @property
    def NextInstruction(self):
        return Instruction_t(idc.NextHead(self.inst_ea))

    @property
    def PrevInstruction(self):
        return Instruction_t(idc.PrevHead(self.inst_ea))
    
    @property
    def FunctionName(self):
        return idc.GetFunctionName(self.inst_ea)

    def IsCode(self):
        return idaapi.isCode(idaapi.getFlags(self.inst_ea))

    def IsCallInstruction(self):
        return self.mnemonic == 'call'

    def GetRvaFromBaseOfFunction(self):
        func_ea = idc.LocByName(self.FunctionName)
        return int(self.inst_ea - func_ea)

    def IsXoredThisReg(self,reg):
        return self.mnemonic == 'xor' and self.operand_0 == self.operand_1 == reg;

    def IsXoredWithItself(self):
        return self.mnemonic == 'xor' and self.operand_0 == self.operand_1;

# ------------------------------------------------------------------------------------- #
# @ IA32 MSR ANALYZER BEGIN ==>
# ------------------------------------------------------------------------------------- #
class Ia32MsrAnalyzePass(object):
    def __init__(self,func_ea,symbolic_constants,enable_symbolic_constants = False,enable_comment = False, \
        enable_colorize = False , color = 0x00441E):

        self.func_ea = func_ea
        self.inst_ea = None

        self.symbolic_constants = symbolic_constants

        self.ia32_msr_db = g_knowledge_db.GetIa32MsrDb()

        self.enable_symbolic_constants = enable_symbolic_constants
        self.enable_comment = enable_comment
        self.enable_colorize = enable_colorize
        self.color = color

        self.msr_code_imm_ea = None

        self.gui_report_rows = []

    def IsVmxInstruction(self,inst_ea):
        mnemonic = idc.GetMnem(inst_ea)
        return mnemonic =='rdmsr' or mnemonic=='wrmsr'

    def SetSymbolcsConstantOfMsrCode(self,inst_ea,msr_code):
        ret = GetEnumItemByConst(msr_code,self.symbolic_constants)
        if ret == None: 
            return
        item_name, enum_id = ret
        if enum_id  and self.ia32_msr_db.get(msr_code)== item_name:
            idc.OpEnumEx(inst_ea, 1,enum_id,0 )

    def GetMsrCodeFromOperand(self,inst_ea):
        curr_inst = Instruction_t(inst_ea)

        while True:
            curr_inst = curr_inst.PrevInstruction
            if(curr_inst.mnemonic == 'mov' and curr_inst.operand_0[-2:] == 'cx'):
                if (curr_inst.GetOperandType(OPERAND_INDEX_1) == idc.o_reg and cfg_enable_smart_mode == True):
                    reg_name = curr_inst.operand_1
                    rev_inst = Instruction_t(curr_inst.inst_ea)

                    if(reg_name[-2:] == 'ax' and rev_inst.PrevInstruction.IsCallInstruction()):
                        """
                        rcx value came from a function call

                        mov     r10d, 0A0013h
                        mov     ecx, r10d
                        call    sub_FFFFF800002E7C78
                        mov     ecx, eax
                        rdmsr                   ; rdmsr(Not imm value) 
                        """
                        return None

                    """
                    mov rax, 0xC0000082 ; IA32_LSTAR
                    mov rcx,rax
                    """
                    for i in range(cfg_smart_mode_max_back_step):
                        rev_inst = rev_inst.PrevInstruction 
                        if rev_inst.mnemonic == 'mov' and rev_inst.operand_0 == reg_name and \
                            rev_inst.GetOperandType(OPERAND_INDEX_1) == idc.o_imm:

                            self.msr_code_imm_ea = rev_inst.inst_ea
                            return rev_inst.GetOperandValue_32bit(OPERAND_INDEX_1)

                # mov rcx, 0xC0000082 ; IA32_LSTAR
                elif(curr_inst.GetOperandType(OPERAND_INDEX_1) == idc.o_imm):
                    self.msr_code_imm_ea = curr_inst.inst_ea
                    return curr_inst.GetOperandValue_32bit(OPERAND_INDEX_1)
                    
            # Not imm value :-(
            elif (curr_inst.mnemonic == 'lea' and  curr_inst.operand_0 in ['rcx','ecx']):
                return None
            # xor rcx,rcx 
            elif (curr_inst.operand_0 in ['rcx','ecx'] and curr_inst.IsXoredThisReg(curr_inst.operand_0)):
                self.msr_code_imm_ea = curr_inst.inst_ea
                return 0
            else:
                continue

    def GetMsrName(self,msr_code):
        name = self.ia32_msr_db.get(msr_code)
        if (name == None):
            name = PrettyHex(msr_code)
        return name

    def GetGuiReportRows(self):
        return self.gui_report_rows

    def RunPass(self):
        for self.inst_ea in idautils.FuncItems(self.func_ea):
            mnemonic = idc.GetMnem(self.inst_ea)
            if (self.IsVmxInstruction(self.inst_ea) == False):
                continue

            msr_name = None
            msr_hex  = None
            msr_code = self.GetMsrCodeFromOperand(self.inst_ea)

            if(msr_code == None):
                msr_name = 'Not imm value'
                msr_hex  = 'Not imm value'
            else:
                msr_name = self.GetMsrName(msr_code)
                msr_hex  = PrettyHex(msr_code)

            if(self.enable_comment):
                comment_message = '{}({})'.format(mnemonic,msr_name)
                idc.MakeComm(self.inst_ea, comment_message)
            
            if(self.enable_colorize):
                idc.SetColor(self.inst_ea,idc.CIC_ITEM, self.color)

            if(self.enable_symbolic_constants and msr_code != None):
                self.SetSymbolcsConstantOfMsrCode(self.msr_code_imm_ea,msr_code)

            self.gui_report_rows.append(ReportEntry(
                                        self.inst_ea,
                                        mnemonic,
                                        msr_hex,
                                        msr_name,
                                        GetFunctionNamePlusRva(self.inst_ea))
                                        )

# ------------------------------------------------------------------------------------- #
# @ IA32 MSR ANALYZER END <==
# ------------------------------------------------------------------------------------- #

# ------------------------------------------------------------------------------------- #
# @IA32 VMCS ANALYZER BEGIN ==>
# ------------------------------------------------------------------------------------- #

class Ia32VmcsAnalyzePass(object):
    def __init__(self,func_ea,symbolic_constants,enable_symbolic_constants = False, enable_comment = False, \
        enable_colorize = False , color = 0x00441E):

        self.func_ea = func_ea
        self.inst_ea = None

        self.symbolic_constants = symbolic_constants

        self.ia32_vmcs_db = g_knowledge_db.GetIa32VmcsDb()

        self.enable_symbolic_constants = enable_symbolic_constants
        self.enable_comment = enable_comment
        self.enable_colorize = enable_colorize
        self.color = color

        self.vmcs_code_imm_ea = None

        self.gui_report_rows = []

    def IsVmxInstruction(self,inst_ea):
        mnemonic = idc.GetMnem(inst_ea)
        return mnemonic =='vmread' or mnemonic=='vmwrite'

    def SetSymbolcsConstantOfVmcsCode(self,inst_ea,vmcs_code):
        ret = GetEnumItemByConst(vmcs_code,self.symbolic_constants)
        if ret == None: 
            return
        item_name, enum_id = ret
        if enum_id  and self.ia32_vmcs_db.get(vmcs_code)== item_name:
            idc.OpEnumEx(inst_ea, 1,enum_id,0 )

    def GetValueRegister(self,inst_ea):
        """
        vmread  buffer, vmcs_code
        vmwrite vmcs_code, buffer
        """
        inst_t = Instruction_t(inst_ea)
        if( inst_t.mnemonic == 'vmread' ):#and inst_t.GetOperandType(0) == idc.o_reg ):
            return inst_t.operand_0
        elif( inst_t.mnemonic == 'vmwrite'):# and inst_t.GetOperandType(1) == idc.o_reg ):
            return inst_t.operand_1
        else:
            return None

    def GetCodeRegister(self,inst_ea):
        """
        vmread  buffer, vmcs_code
        vmwrite vmcs_code, buffer
        """
        inst_t = Instruction_t(inst_ea)
        if( inst_t.mnemonic == 'vmread' and inst_t.GetOperandType(1) == idc.o_reg ):
            return inst_t.operand_1
        elif( inst_t.mnemonic == 'vmwrite' and inst_t.GetOperandType(0) == idc.o_reg ):
            return inst_t.operand_0
        else:
            return None

    def GetVmcsCodeFromOperand(self,inst_ea):
        curr_inst = Instruction_t(inst_ea)

        code_reg = self.GetCodeRegister(inst_ea)
        if( code_reg == None ):
            return None

        while True:
            curr_inst = curr_inst.PrevInstruction

            # mov  ['cx', 'ecx', 'rcx'] [-2:] ==> 'cx'  , ###
            if( curr_inst.mnemonic == 'mov' and curr_inst.operand_0[-2:] == code_reg[-2:]):
                if(curr_inst.GetOperandType(OPERAND_INDEX_1) == idc.o_imm):
                    self.vmcs_code_imm_ea = curr_inst.inst_ea
                    return curr_inst.GetOperandValue_32bit(OPERAND_INDEX_1)

                # if mov ecx, rcx keep looking for mov rax, imm
                elif(curr_inst.GetOperandType(OPERAND_INDEX_1) == idc.o_reg and cfg_enable_smart_mode == True):
                    reg_name = curr_inst.operand_1
                    rev_inst = Instruction_t(curr_inst.inst_ea)

                    if(reg_name[-2:] == 'ax' and rev_inst.PrevInstruction.IsCallInstruction()):
                        """
                        register value came from a function call
                        """
                        return None
                    
                    # while True:
                    for i in range(cfg_smart_mode_max_back_step):
                        rev_inst = rev_inst.PrevInstruction
                        if rev_inst.mnemonic == 'mov' and rev_inst.operand_0 == reg_name and \
                            rev_inst.GetOperandType(OPERAND_INDEX_1) == idc.o_imm:

                            self.vmcs_code_imm_ea = rev_inst.inst_ea
                            return rev_inst.GetOperandValue_32bit(OPERAND_INDEX_1)

            # Not imm value :-(
            elif (curr_inst.mnemonic == 'lea' and curr_inst.operand_0[-2:] == code_reg[-2:]):
                return None
            # xor ecx,ecx means vmcs code is 0x0 VMCS_CTRL_VPID
            elif(curr_inst.operand_0[-2:] == code_reg[-2:] and curr_inst.IsXoredThisReg(code_reg)):
                self.vmcs_code_imm_ea = curr_inst.inst_ea
                return 0

    def GetVmcsName(self,vmcs_code):
        name = self.ia32_vmcs_db.get(vmcs_code)
        if (name == None):
            name = PrettyHex(vmcs_code)
        return name

    def GetGuiReportRows(self):
        return self.gui_report_rows

    def RunPass(self):
        for self.inst_ea in idautils.FuncItems(self.func_ea):
            mnemonic = idc.GetMnem(self.inst_ea)
            if (self.IsVmxInstruction(self.inst_ea) == False):
                continue

            vmcs_name = None
            vmcs_hex  = None
            vmcs_code = self.GetVmcsCodeFromOperand(self.inst_ea)

            if(vmcs_code == None):
                vmcs_name = 'Not imm value'
                vmcs_hex  = 'Not imm value'
            else:
                vmcs_name = self.GetVmcsName(vmcs_code)
                vmcs_hex  = PrettyHex(vmcs_code)

            if(self.enable_comment):
                comment_message = '{}({}, {})'.format(mnemonic,vmcs_name,self.GetValueRegister(self.inst_ea))
                idc.MakeComm(self.inst_ea, comment_message)
            
            if(self.enable_colorize):
                idc.SetColor(self.inst_ea,idc.CIC_ITEM, self.color)

            if(self.enable_symbolic_constants and vmcs_code != None):
                self.SetSymbolcsConstantOfVmcsCode(self.vmcs_code_imm_ea,vmcs_code)

            self.gui_report_rows.append(ReportEntry(
                                        self.inst_ea,
                                        mnemonic,
                                        vmcs_hex,
                                        vmcs_name,
                                        GetFunctionNamePlusRva(self.inst_ea))
                                        )

# ------------------------------------------------------------------------------------- #
# @IA32 VMCS ANALYZER END <==
# ------------------------------------------------------------------------------------- #

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


class Ia32VmxHelper(idaapi.plugin_t):
    def __init__(self):
        self.enable_comment  = 0
        self.enable_colorize = 0
        self.enable_symbolic_constants = 0
        self.gui_report_rows = []

        self.ia32_msr_symbolic_constants = SymbolicConstant('IA32_MSR_LIST_ENUM').LoadEnums()
        self.ia32_vmcs_symbolic_constants = SymbolicConstant('IA32_VMCS_LIST_ENUM').LoadEnums()

        global g_knowledge_db
        g_knowledge_db = KnowledgeDb()

        flags = idaapi.PLUGIN_UNL
        comment = ''
        help = ''
        wanted_name = ''
        # wanted_hotkey = 'Ctrl-Alt-Q'

    def AddRowsToGuiTable(self,rows):
        if self._plugin_page is None:
            self._plugin_page = PluginPage(
                    Config.PLUGIN_TITLTE, Config.PLUGIN_COLUMNS,rows, self.icon_id)
        else:
            self._plugin_page.SetItems(rows)

        self._plugin_page.Show()

    def init(self):
        self._plugin_page = None
        self.icon_id = None
        return idaapi.PLUGIN_OK

    def term(self):
        pass

    def run(self,_=0):
        only_show = idc.AskYN(1,'Only show the result')
        if(only_show == 0):
            add_enums = idc.AskYN(0,'Do you want to add symbolic constants enums?, **You Only need to do this for once not more!**')
            if(add_enums == 1):
                AddSymbolicConstantsEnumsToIda()
                if(self.ia32_msr_symbolic_constants == None):
                    self.ia32_msr_symbolic_constants = SymbolicConstant('IA32_MSR_LIST_ENUM').LoadEnums()
                if(self.ia32_vmcs_symbolic_constants == None):
                    self.ia32_vmcs_symbolic_constants = SymbolicConstant('IA32_VMCS_LIST_ENUM').LoadEnums()

            self.enable_comment  = idc.AskYN(0,'Do you want to add comment on all decoded values?')
            self.enable_colorize = idc.AskYN(0,'Do you want to colorize all decoded values?')
            self.enable_symbolic_constants = idc.AskYN(0,'Do you want to apply symbolic constants on all decoded values?')

        for func_ea in idautils.Functions():
            isa32_msr_pass = Ia32MsrAnalyzePass(func_ea,
                                                self.ia32_msr_symbolic_constants,
                                                self.enable_symbolic_constants,
                                                self.enable_comment,
                                                self.enable_colorize,
                                                0x00441E)

            isa32_vmcs_pass = Ia32VmcsAnalyzePass(func_ea,
                                                self.ia32_vmcs_symbolic_constants,
                                                self.enable_symbolic_constants,
                                                self.enable_comment,
                                                self.enable_colorize,
                                                0x00441E)
            
            isa32_msr_pass.RunPass()
            isa32_vmcs_pass.RunPass()
            

            for i in isa32_msr_pass.GetGuiReportRows():
                 self.gui_report_rows.append(i.get_row())

            for i in isa32_vmcs_pass.GetGuiReportRows():
                 self.gui_report_rows.append(i.get_row())

        self.AddRowsToGuiTable(self.gui_report_rows)
# ------------------------------------------------------------------------------------- #
def PLUGIN_ENTRY():
    return Ia32VmxHelper()

def main():
    f = PLUGIN_ENTRY()
    f.init()
    f.run()
    f.term()

if __name__ == '__main__':
    main()


repository: idapython/pyhexraysdeob
file: pyhexraysdeob.py
description: A port of Rolf Rolles' https://github.com/RolfRolles/HexRaysDeob to Python
wanted_name: Python Microcode explorer (pyhexraysdeob)
comment: Show microcode
created_at: 2019-10-15 13:12:47+00:00
pushed_at: 2022-05-04 09:18:08+00:00
url:
https://github.com/idapython/pyhexraysdeob/blob/dd3588b99228a83e67ceaa32ca2fc513af0cc424/pyhexraysdeob.py

content:
import os

import ida_idaapi
import ida_loader
import ida_kernwin
import ida_hexrays

my_dirname, _ = os.path.split(__file__)

setattr(ida_hexrays, "MMAT_DEOB_MAP", getattr(ida_hexrays, "MMAT_LOCOPT"))

class pyhexraysdeob_t(ida_idaapi.plugin_t):
    flags = 0
    comment = "Show microcode"
    help = ""
    wanted_name = "Python Microcode explorer (pyhexraysdeob)"
    wanted_hotkey = ""

    def __init__(self):
        self.black_list = []
        self.white_list = []
        self.activated = False

    def toggle_activated(self):
        if not self.activated:
            # Install our block and instruction optimization classes.
            import pattern_deobfuscate
            self.oco = pattern_deobfuscate.obf_compiler_optimizer_t()
            self.oco.install()
            import unflattener
            self.cfu = unflattener.cf_unflattener_t(self)
            self.cfu.install()
        else:
            # Uninstall our block and instruction optimization classes.
            self.oco.remove()
            self.oco = None
            self.cfu.remove()
            self.cfu = None
        self.activated = not self.activated
        print("%s is now %sactivated" % (self.wanted_name, "" if self.activated else "de-"))

    def init(self):
        if not ida_hexrays.init_hexrays_plugin():
            print("pyhexraysdeob: no decompiler, skipping")
            return ida_idaapi.PLUGIN_SKIP
        print("Hex-rays version %s has been detected, %s ready to use" % (
            ida_hexrays.get_hexrays_version(),
            self.wanted_name))

        import sys
        modules_path = os.path.join(my_dirname, "pyhexraysdeob_modules")
        if not modules_path in sys.path:
            sys.path.append(modules_path)

        return ida_idaapi.PLUGIN_OK

    def run(self, arg):
        if arg == 0:
            self.toggle_activated()
        elif arg == 0xbeef:
            self.flags |= ida_loader.PLUGIN_UNL
        elif arg == 2:
            fix_calls_to_alloca_probe() # unimp
        elif arg == 3:
            show_microcode_explorer() # unimp
        return True

    def term(self):
        if self.activated:
            self.toggle_activated()

def PLUGIN_ENTRY():
    return pyhexraysdeob_t()


repository: hdbreaker/RicardoNarvaja_ExploitingWithIDAPRO_Desde0
file: 35-INTRODUCCION AL REVERSING CON IDA PRO DESDE CERO PARTE 35/idasploiter.py
description: Nuevo Curso Completo de Exploiting con IDA Pro pro Ricardo Narvaja
wanted_name: IDA Sploiter
comment: Exploit development and vulnerability research toolkit.
created_at: 2018-05-27 20:23:39+00:00
pushed_at: 2018-05-27 20:31:59+00:00
url:
https://github.com/hdbreaker/RicardoNarvaja_ExploitingWithIDAPRO_Desde0/blob/c7bbc6def93b6c45200e88af584b6c3e34c5fc54/35-INTRODUCCION%20AL%20REVERSING%20CON%20IDA%20PRO%20DESDE%20CERO%20PARTE%2035/idasploiter.py

content:
#!/usr/bin/env python
#
# IDA Sploiter is an exploit development and vulnerability research environment
# implemented as a plugin for Hex-Ray's IDA Pro disassembler.

IDASPLOITER_VERSION = "1.0"

# Copyright (C) 2014 Peter Kacherginsky
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met: 
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer. 
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
# 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 OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# Performance profiling
import cProfile
import pstats

# IDA libraries
import idaapi
import idautils
import idc
from idaapi import Form, Choose2, plugin_t

# Python libraries
import os
import binascii
import string
import textwrap
import copy
import csv
import itertools

from struct import pack, unpack
from ctypes import *

###############################################################################
# Data Tables and Structures
###############################################################################

###############################################################################
# 00-FF Single byte to unicode transforms
# NOTE: Some unicode characters can have two bytes
#       http://www.phenoelit.org/stuff/Phenoelit20c3.pdf

# Ascii_to_Unicode_transforms
ASCII_TO_UNICODE_GENERAL = [
    "\x00\x00", "\x00\x01", "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05",
    "\x00\x06", "\x00\x07", "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b", 
    "\x00\x0c", "\x00\x0d", "\x00\x0e", "\x00\x0f", "\x00\x10", "\x00\x11", 
    "\x00\x12", "\x00\x13", "\x00\x14", "\x00\x15", "\x00\x16", "\x00\x17", 
    "\x00\x18", "\x00\x19", "\x00\x1a", "\x00\x1b", "\x00\x1c", "\x00\x1d", 
    "\x00\x1e", "\x00\x1f", "\x00\x20", "\x00\x21", "\x00\x22", "\x00\x23", 
    "\x00\x24", "\x00\x25", "\x00\x26", "\x00\x27", "\x00\x28", "\x00\x29", 
    "\x00\x2a", "\x00\x2b", "\x00\x2c", "\x00\x2d", "\x00\x2e", "\x00\x2f", 
    "\x00\x30", "\x00\x31", "\x00\x32", "\x00\x33", "\x00\x34", "\x00\x35", 
    "\x00\x36", "\x00\x37", "\x00\x38", "\x00\x39", "\x00\x3a", "\x00\x3b", 
    "\x00\x3c", "\x00\x3d", "\x00\x3e", "\x00\x3f", "\x00\x40", "\x00\x41", 
    "\x00\x42", "\x00\x43", "\x00\x44", "\x00\x45", "\x00\x46", "\x00\x47", 
    "\x00\x48", "\x00\x49", "\x00\x4a", "\x00\x4b", "\x00\x4c", "\x00\x4d", 
    "\x00\x4e", "\x00\x4f", "\x00\x50", "\x00\x51", "\x00\x52", "\x00\x53", 
    "\x00\x54", "\x00\x55", "\x00\x56", "\x00\x57", "\x00\x58", "\x00\x59", 
    "\x00\x5a", "\x00\x5b", "\x00\x5c", "\x00\x5d", "\x00\x5e", "\x00\x5f", 
    "\x00\x60", "\x00\x61", "\x00\x62", "\x00\x63", "\x00\x64", "\x00\x65", 
    "\x00\x66", "\x00\x67", "\x00\x68", "\x00\x69", "\x00\x6a", "\x00\x6b", 
    "\x00\x6c", "\x00\x6d", "\x00\x6e", "\x00\x6f", "\x00\x70", "\x00\x71", 
    "\x00\x72", "\x00\x73", "\x00\x74", "\x00\x75", "\x00\x76", "\x00\x77", 
    "\x00\x78", "\x00\x79", "\x00\x7a", "\x00\x7b", "\x00\x7c", "\x00\x7d", 
    "\x00\x7e", "\x00\x7f"
]

ASCII_TO_UNICODE_ANSI = ASCII_TO_UNICODE_GENERAL + [
    "\x20\xac", "\x00\x81", "\x20\x1a", "\x01\x92", "\x20\x1e", "\x20\x26",
    "\x20\x20", "\x20\x21", "\x02\xc6", "\x20\x30", "\x01\x60", "\x20\x39",
    "\x01\x52", "\x00\x8d", "\x01\x7d", "\x00\x8f", "\x90\x00", "\x20\x18",
    "\x20\x19", "\x20\x1c", "\x20\x1d", "\x20\x22", "\x20\x13", "\x20\x14",
    "\x02\xdc", "\x21\x22", "\x01\x61", "\x3a\x20", "\x01\x53", "\x00\x9d",
    "\x01\x7e", "\x01\x78", "\x00\xa0", "\x00\xa1", "\x00\xa2", "\x00\xa3",
    "\x00\xa4", "\x00\xa5", "\x00\xa6", "\x00\xa7", "\x00\xa8", "\x00\xa9",
    "\x00\xaa", "\x00\xab", "\x00\xac", "\x00\xad", "\x00\xae", "\x00\xaf",
    "\x00\xb0", "\x00\xb1", "\x00\xb2", "\x00\xb3", "\x00\xb4", "\x00\xb5",
    "\x00\xb6", "\x00\xb7", "\x00\xb8", "\x00\xb9", "\x00\xba", "\x00\xbb",
    "\x00\xbc", "\x00\xbd", "\x00\xbe", "\x00\xbf", "\x00\xc0", "\x00\xc1",
    "\x00\xc2", "\x00\xc3", "\x00\xc4", "\x00\xc5", "\x00\xc6", "\x00\xc7",
    "\x00\xc8", "\x00\xc9", "\x00\xca", "\x00\xcb", "\x00\xcc", "\x00\xcd",
    "\x00\xce", "\x00\xcf", "\x00\xd0", "\x00\xd1", "\x00\xd2", "\x00\xd3",
    "\x00\xd4", "\x00\xd5", "\x00\xd6", "\x00\xd7", "\x00\xd8", "\x00\xd9",
    "\x00\xda", "\x00\xdb", "\x00\xdc", "\x00\xdd", "\x00\xde", "\x00\xdf",
    "\x00\xe0", "\x00\xe1", "\x00\xe2", "\x00\xe3", "\x00\xe4", "\x00\xe5",
    "\x00\xe6", "\x00\xe7", "\x00\xe8", "\x00\xe9", "\x00\xea", "\x00\xeb",
    "\x00\xec", "\x00\xed", "\x00\xee", "\x00\xef", "\x00\xf0", "\x00\xf1",
    "\x00\xf2", "\x00\xf3", "\x00\xf4", "\x00\xf5", "\x00\xf6", "\x00\xf7",
    "\x00\xf8", "\x00\xf9", "\x00\xfa", "\x00\xfb", "\x00\xfc", "\x00\xfd",
    "\x00\xfe", "\x00\xff"
]

ASCII_TO_UNICODE_OEM = ASCII_TO_UNICODE_GENERAL + [
    "\x00\xc7", "\x00\xfc", "\x00\xe9", "\x00\xe2", "\x00\xe4", "\x00\xe0",
    "\x00\xe5", "\x00\xe7", "\x00\xea", "\x00\xeb", "\x00\xe8", "\x00\xef",
    "\x00\xee", "\x00\xec", "\x00\xc4", "\x00\xc5", "\x00\xc9", "\x00\xe6",
    "\x00\xc6", "\x00\xf4", "\x00\xf6", "\x00\xf2", "\x00\xfb", "\x00\xf9",
    "\x00\xff", "\x00\xd6", "\x00\xdc", "\x00\xf8", "\x00\xa3", "\x00\xd8",
    "\x00\xd7", "\x01\x92", "\x00\xe1", "\x00\xed", "\x00\xf3", "\x00\xfa",
    "\x00\xf1", "\x00\xd1", "\x00\xaa", "\x00\xba", "\x00\xbf", "\x00\xae",
    "\x00\xac", "\x00\xbd", "\x00\xbc", "\x00\xa1", "\x00\xab", "\x00\xbb",
    "\x25\x91", "\x25\x92", "\x25\x93", "\x25\x02", "\x25\x24", "\x00\xc1",
    "\x00\xc2", "\x00\xc0", "\x00\xa9", "\x25\x63", "\x25\x51", "\x25\x57",
    "\x25\x5d", "\x00\xa2", "\x00\xa5", "\x25\x10", "\x25\x14", "\x25\x34",
    "\x25\x2c", "\x25\x1c", "\x25\x00", "\x25\x3c", "\x00\xe3", "\x00\xc3",
    "\x25\x5a", "\x25\x54", "\x25\x69", "\x25\x66", "\x25\x60", "\x25\x50",
    "\x25\x6c", "\x00\xa4", "\x00\xf0", "\x00\xd0", "\x00\xca", "\x00\xcb",
    "\x00\xc8", "\x01\x31", "\x00\xcd", "\x00\xce", "\x00\xcf", "\x25\x18",
    "\x25\x0c", "\x25\x88", "\x25\x84", "\x00\xa6", "\x00\xcc", "\x25\x80",
    "\x00\xd3", "\x00\xdf", "\x00\xd4", "\x00\xd2", "\x00\xf5", "\x00\xd5",
    "\x00\xb5", "\x00\xfe", "\x00\xde", "\x00\xda", "\x00\xdb", "\x00\xd9",
    "\x00\xfd", "\x00\xdd", "\x00\xaf", "\x00\xb4", "\x00\xad", "\x00\xb1",
    "\x20\x17", "\x00\xbe", "\x00\xb6", "\x00\xa7", "\x00\xf7", "\x00\xb8",
    "\x00\xb0", "\x00\xa8", "\x00\xb7", "\x00\xb9", "\x00\xb3", "\x00\xb2",
    "\x25\xa0", "\x00\xa0"
]

ASCII_TO_UNICODE_UTF7 = ASCII_TO_UNICODE_GENERAL + [
    "\xff\x80", "\xff\x81", "\xff\x82", "\xff\x83", "\xff\x84", "\xff\x85",
    "\xff\x86", "\xff\x87", "\xff\x88", "\xff\x89", "\xff\x8a", "\xff\x8b",
    "\xff\x8c", "\xff\x8d", "\xff\x8e", "\xff\x8f", "\xff\x90", "\xff\x91",
    "\xff\x92", "\xff\x93", "\xff\x94", "\xff\x95", "\xff\x96", "\xff\x97",
    "\xff\x98", "\xff\x99", "\xff\x9a", "\xff\x9b", "\xff\x9c", "\xff\x9d",
    "\xff\x9e", "\xff\x9f", "\xff\xa0", "\xff\xa1", "\xff\xa2", "\xff\xa3",
    "\xff\xa4", "\xff\xa5", "\xff\xa6", "\xff\xa7", "\xff\xa8", "\xff\xa9",
    "\xff\xaa", "\xff\xab", "\xff\xac", "\xff\xad", "\xff\xae", "\xff\xaf",
    "\xff\xb0", "\xff\xb1", "\xff\xb2", "\xff\xb3", "\xff\xb4", "\xff\xb5",
    "\xff\xb6", "\xff\xb7", "\xff\xb8", "\xff\xb9", "\xff\xba", "\xff\xbb",
    "\xff\xbc", "\xff\xbd", "\xff\xbe", "\xff\xbf", "\xff\xc0", "\xff\xc1",
    "\xff\xc2", "\xff\xc3", "\xff\xc4", "\xff\xc5", "\xff\xc6", "\xff\xc7",
    "\xff\xc8", "\xff\xc9", "\xff\xca", "\xff\xcb", "\xff\xcc", "\xff\xcd",
    "\xff\xce", "\xff\xcf", "\xff\xd0", "\xff\xd1", "\xff\xd2", "\xff\xd3",
    "\xff\xd4", "\xff\xd5", "\xff\xd6", "\xff\xd7", "\xff\xd8", "\xff\xd9",
    "\xff\xda", "\xff\xdb", "\xff\xdc", "\xff\xdd", "\xff\xde", "\xff\xdf",
    "\xff\xe0", "\xff\xe1", "\xff\xe2", "\xff\xe3", "\xff\xe4", "\xff\xe5",
    "\xff\xe6", "\xff\xe7", "\xff\xe8", "\xff\xe9", "\xff\xea", "\xff\xeb",
    "\xff\xec", "\xff\xed", "\xff\xee", "\xff\xef", "\xff\xf0", "\xff\xf1",
    "\xff\xf2", "\xff\xf3", "\xff\xf4", "\xff\xf5", "\xff\xf6", "\xff\xf7",
    "\xff\xf8", "\xff\xf9", "\xff\xfa", "\xff\xfb", "\xff\xfc", "\xff\xfd",
    "\xff\xfe", "\xff\xff"
]

ASCII_TO_UNICODE_UTF8 = ASCII_TO_UNICODE_GENERAL + [
    "\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", "\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\x01",
    "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05", "\x00\x06", "\x00\x07",
    "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b", "\x00\x0c", "\x00\x0d",
    "\x00\x0e", "\x00\x0f", "\x00\x10", "\x00\x11", "\x00\x12", "\x00\x13",
    "\x00\x14", "\x00\x15", "\x00\x16", "\x00\x17", "\x00\x18", "\x00\x19",
    "\x00\x1a", "\x00\x1b", "\x00\x1c", "\x00\x1d", "\x00\x1e", "\x00\x1f",
    "\x00\x00", "\x00\x01", "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05",
    "\x00\x06", "\x00\x07", "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b",
    "\x00\x0c", "\x00\x0d", "\x00\x0e", "\x00\x0f", "\x00\x00", "\x00\x01",
    "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05", "\x00\x06", "\x00\x07",
    "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b", "\x00\x0c", "\x00\x0d",
    "\x00\x0e", "\x00\x0f"
]

ASCII_TO_UNICODE = [ASCII_TO_UNICODE_ANSI,
                    ASCII_TO_UNICODE_OEM, 
                    ASCII_TO_UNICODE_UTF7, 
                    ASCII_TO_UNICODE_UTF8]

###############################################################################
# Microsoft Portable Executable and Common Object File Format Specification
# Revision 8.3 - February 6, 2013

class IMAGE_DOS_HEADER(Structure):
    _fields_ = [
                    ("signature" , c_char * 2),
                    ("lastsize"  , c_short),
                    ("nblocks"   , c_short),
                    ("nreloc"    , c_short),
                    ("hdrsize"   , c_short),
                    ("minalloc"  , c_short),
                    ("maxalloc"  , c_short),
                    ("ss"        , c_short),
                    ("sp"        , c_short),
                    ("checksum"  , c_short),
                    ("ip"        , c_short),
                    ("cs"        , c_short),
                    ("relocpos"  , c_short),
                    ("noverlay"  , c_short),
                    ("reserved1" , c_short * 4),
                    ("oem_id"    , c_short),
                    ("reserved2" , c_short * 10),
                    ("e_lfanew"  , c_long)
                ]

class IMAGE_FILE_HEADER(Structure):
    _fields_ = [
                    ("Machine"              , c_short),
                    ("NumberOfSections"     , c_short),
                    ("TimeDateStamp"        , c_long),
                    ("PointerToSymbolTable" , c_long),
                    ("NumberOfSymbols"      , c_long),
                    ("SizeOfOptionalHeader" , c_short),
                    ("Characteristics"      , c_short)
                ]

class IMAGE_DATA_DIRECTORY(Structure):
    _fields_ = [
                    ("VirtualAddress", c_long),
                    ("Size", c_long)
                ]

class IMAGE_OPTIONAL_HEADER(Structure):
    _fields_ = [
                    ("signature"               , c_short),
                    ("MajorLinkerVersion"      , c_byte),
                    ("MinorLinkerVersion"      , c_byte),
                    ("SizeOfCode"              , c_long),
                    ("SizeOfInitializedData"   , c_long),
                    ("SizeOfUninitializedData" , c_long),
                    ("AddressOfEntryPoint"     , c_long),
                    ("BaseOfCode"              , c_long),
                    ("BaseOfData"              , c_long),
                    ("ImageBase"               , c_long),
                    ("SectionAlignment"        , c_long),
                    ("FileAlignment"           , c_long),
                    ("MajorOSVersion"          , c_short),
                    ("MinorOSVersion"          , c_short),
                    ("MajorImageVersion"       , c_short),
                    ("MinorImageVersion"       , c_short),
                    ("MajorSubsystemVersion"   , c_short),
                    ("MinorSubsystemVersion"   , c_short),
                    ("Reserved"                , c_long),
                    ("SizeOfImage"             , c_long),
                    ("SizeOfHeaders"           , c_long),
                    ("Checksum"                , c_long),
                    ("Subsystem"               , c_short),
                    ("DLLCharacteristics"      , c_short),
                    ("SizeOfStackReserve"      , c_long),
                    ("SizeOfStackCommit"       , c_long),
                    ("SizeOfHeapReserve"       , c_long),
                    ("SizeOfHeapCommit"        , c_long),
                    ("LoaderFlags"             , c_long),
                    ("NumberOfRvaAndSizes"     , c_long),
                    ("DataDirectory"           , IMAGE_DATA_DIRECTORY * 16)
                ]

class IMAGE_OPTIONAL_HEADER64(Structure):
    _fields_ = [
                    ("signature"               , c_short),
                    ("MajorLinkerVersion"      , c_byte),
                    ("MinorLinkerVersion"      , c_byte),
                    ("SizeOfCode"              , c_long),
                    ("SizeOfInitializedData"   , c_long),
                    ("SizeOfUninitializedData" , c_long),
                    ("AddressOfEntryPoint"     , c_long),
                    ("BaseOfCode"              , c_long),
                    ("ImageBase"               , c_longlong),
                    ("SectionAlignment"        , c_long),
                    ("FileAlignment"           , c_long),
                    ("MajorOSVersion"          , c_short),
                    ("MinorOSVersion"          , c_short),
                    ("MajorImageVersion"       , c_short),
                    ("MinorImageVersion"       , c_short),
                    ("MajorSubsystemVersion"   , c_short),
                    ("MinorSubsystemVersion"   , c_short),
                    ("Reserved"                , c_long),
                    ("SizeOfImage"             , c_long),
                    ("SizeOfHeaders"           , c_long),
                    ("Checksum"                , c_long),
                    ("Subsystem"               , c_short),
                    ("DLLCharacteristics"      , c_short),
                    ("SizeOfStackReserve"      , c_longlong),
                    ("SizeOfStackCommit"       , c_longlong),
                    ("SizeOfHeapReserve"       , c_longlong),
                    ("SizeOfHeapCommit"        , c_longlong),
                    ("LoaderFlags"             , c_long),
                    ("NumberOfRvaAndSizes"     , c_long),
                    ("DataDirectory"           , IMAGE_DATA_DIRECTORY * 16)
                ]

class IMAGE_LOAD_CONFIG_DIRECTORY(Structure):
    _fields_ = [
                    ("Size"                          , c_long),
                    ("TimeDateStamp"                 , c_long),
                    ("MajorVersion"                  , c_short),
                    ("MinorVersion"                  , c_short),
                    ("GlobalFlagsClear"              , c_long),
                    ("GlobalFlagsSet"                , c_long),
                    ("CriticalSectionDefaultTimeout" , c_long),
                    ("DeCommitFreeBlockThreshold"    , c_long),
                    ("DeCommitTotalFreeThreshold"    , c_long),
                    ("LockPrefixTable"               , c_long),
                    ("MaximumAllocationSize"         , c_long),
                    ("VirtualMemoryThreshold"        , c_long),
                    ("ProcessHeapFlags"              , c_long),
                    ("ProcessAffinityMask"           , c_long),
                    ("CSDVersion"                    , c_short),
                    ("Reserved1"                     , c_short),
                    ("EditList"                      , c_long),
                    ("SecurityCookie"                , c_long),
                    ("SEHandlerTable"                , c_long),
                    ("SEHandlerCount"                , c_long)
                ]

class IMAGE_LOAD_CONFIG_DIRECTORY64(Structure):
    _fields_ = [
                    ("Size"                          , c_long),
                    ("TimeDateStamp"                 , c_long),
                    ("MajorVersion"                  , c_short),
                    ("MinorVersion"                  , c_short),
                    ("GlobalFlagsClear"              , c_long),
                    ("GlobalFlagsSet"                , c_long),
                    ("CriticalSectionDefaultTimeout" , c_long),
                    ("DeCommitFreeBlockThreshold"    , c_longlong),
                    ("DeCommitTotalFreeThreshold"    , c_longlong),
                    ("LockPrefixTable"               , c_longlong),
                    ("MaximumAllocationSize"         , c_longlong),
                    ("VirtualMemoryThreshold"        , c_longlong),
                    ("ProcessHeapFlags"              , c_long),
                    ("ProcessAffinityMask"           , c_long),
                    ("CSDVersion"                    , c_short),
                    ("Reserved1"                     , c_short),
                    ("EditList"                      , c_longlong),
                    ("SecurityCookie"                , c_longlong),
                    ("SEHandlerTable"                , c_longlong),
                    ("SEHandlerCount"                , c_longlong),
                ]

###############################################################################
# Module Class - Manages module characteristics

class Module():

    def __init__(self, name, size, base, rebase_to):

            self.addr = base

            # BUG: IDA's API does not always zero out the SWIG buffer.
            # BUG: ntdll does not have path information.
            if "\x00" in name:
                name,junk = name.split("\x00",1)

            self.file = os.path.basename(name)
            self.path = os.path.dirname(name)

            self.size = size

            # Parse the module as a PE file
            pe = PE(self.addr)

            if pe:
                self.NXCompat = pe.isNXCompat()
                self.ASLR     = pe.isDynamicBase()
                self.SafeSEH  = pe.isSafeSEH()
                self.GS       = pe.isGS()

###############################################################################
# PE Module Parser

class PE():

    # PE flags
    IMAGE_FILE_MACHINE_I386 = 0x14c

    MAGIC_PE32      = 0x10b
    MAGIC_PE32_PLUS = 0x20b

    IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE = 0x0040
    IMAGE_DLL_CHARACTERISTICS_NX_COMPAT    = 0x0100
    IMAGE_DLLCHARACTERISTICS_NO_SEH        = 0x0400

    def __init__(self, base):

        # PE headers
        self.dos_header            = None
        self.file_header           = None
        self.optional_header       = None
        self.load_config_directory = None
        self.arch64                = False

        # Verify DOS header magic
        magic_dos = idaapi.dbg_read_memory( base, 0x2)
        if magic_dos != "MZ":
            print "[idasploiter] Invalid DOS header magic."
            return None

        # Parse the DOS header
        self.dos_header = IMAGE_DOS_HEADER.from_buffer_copy(
            idaapi.dbg_read_memory(base, 0x40) )

        # Get offset to PE header
        base_pe = base + self.dos_header.e_lfanew

        # Verify PE header magic
        magic_pe = idaapi.dbg_read_memory( base_pe, 0x4 )
        if magic_pe != "PE\x00\x00":
            print "[idasploiter] Invalid PE header magic: %x" % unpack("I",magic_pe)[0]
            return None

        # Parse the FILE header
        # NOTE: IMAGE_FILE_HEADER size is 0x14
        self.file_header = IMAGE_FILE_HEADER.from_buffer_copy( 
            idaapi.dbg_read_memory( base_pe + 0x4, sizeof(IMAGE_FILE_HEADER) ) )

        # Parse the OPTIONAL header
        base_pe_optional =  base_pe + 0x4 + 0x14
        magic_pe_opt = idaapi.dbg_read_memory ( base_pe_optional, 0x2)
        magic_pe_opt = unpack("H", magic_pe_opt)[0]

        # PE32 Optional Header
        if magic_pe_opt == self.MAGIC_PE32:
            self.arch64 = False
            self.optional_header = IMAGE_OPTIONAL_HEADER.from_buffer_copy(
                idaapi.dbg_read_memory( base_pe_optional, self.file_header.SizeOfOptionalHeader ))

        # PE32+ Optional Header
        elif magic_pe_opt == self.MAGIC_PE32_PLUS:
            self.arch64 = True
            self.optional_header = IMAGE_OPTIONAL_HEADER64.from_buffer_copy(
                idaapi.dbg_read_memory( base_pe_optional, self.file_header.SizeOfOptionalHeader ))

        # Invalid PE Header
        else:
            print "[idasploiter] Invalid IMAGE_OPTIONAL_HEADER magic: %x" % unpack("H",magic_pe_opt)[0]
            return None

        # Load Configuration Table
        load_config_directory_rva = self.optional_header.DataDirectory[10].VirtualAddress

        # Parse Load Configuration Table if present
        if load_config_directory_rva:

            # Read LOAD CONFIG DIRECTORY size before parsing it
            load_config_directory_size = idaapi.dbg_read_memory( base + load_config_directory_rva, 0x4)
            load_config_directory_size = unpack("I",load_config_directory_size)[0]

            # Parse LOAD CONFIG DIRECTORY based on PE32 format
            if self.optional_header.signature == self.MAGIC_PE32:
                self.load_config_directory = IMAGE_LOAD_CONFIG_DIRECTORY.from_buffer_copy( idaapi.dbg_read_memory( base + load_config_directory_rva, load_config_directory_size ))

            # Parse LOAD CONFIG DIRECTORY based on PE32+ format
            elif self.optional_header.signature == self.MAGIC_PE32_PLUS:
                self.load_config_directory = IMAGE_LOAD_CONFIG_DIRECTORY64.from_buffer_copy( idaapi.dbg_read_memory( base + load_config_directory_rva, load_config_directory_size ))

            # Invalid PE header
            else:
                print "[idasploiter] Invalid IMAGE_OPTIONAL_HEADER magic: %x" % unpack("H",magic)[0]
                return None

    def isDynamicBase(self):
        if self.optional_header:
            if self.optional_header.DLLCharacteristics & self.IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE:
                return "Yes"
            else:
                return "No"
        else:
            return "Unk"

    def isNXCompat(self):
        if self.optional_header:
            if self.optional_header.DLLCharacteristics & self.IMAGE_DLL_CHARACTERISTICS_NX_COMPAT:
                return "Yes"
            else:
                return "No"
        else:
            return "Unk"

    def isSafeSEH(self):

        if self.file_header and self.optional_header:

            # SafeSEH is only applicable to 32-bit Windows.
            if self.file_header.Machine != self.IMAGE_FILE_MACHINE_I386:
                return "N/A"

            # NO_SEH flag indicates that the module does not use SEH and no handlers can be called in this image.
            elif self.optional_header.DLLCharacteristics & self.IMAGE_DLLCHARACTERISTICS_NO_SEH:
                return "No SEH"

            # SafeSEH is enabled when there is SEHandlerTable entry in the LOAD_CONFIG_DIRECTORY
            elif self.load_config_directory and self.load_config_directory.SEHandlerTable != 0:
                return "Yes"

            # Otherwise SafeSEH is not used
            else:
                return "No"

        else:
            return "Unk"

    def isGS(self):

        if self.file_header and self.optional_header:

            # NOTE: PE32+ does not necessarily provide a pointer to the Security Cookie in the
            #       LoadConfigDirectory. The address appears to be still present in the binary.

            # TODO: Need to do additional research into reliably detecting /GS in PE32+ files.

            # PE32+ binaries with missing LoadConfigDirectory may still used security cookies.
            if self.arch64 and not self.load_config_directory:
                return "Unk"

            if self.load_config_directory and self.load_config_directory.SecurityCookie != 0:

                # Attempt to read the security cookie
                # NOTE: The extra exception handler was added to account for cookie being stored in memory locations
                #       that are not accessible by the debugger (e.g. manually loading win32k.sys, security cookie
                #       address will point to kernel space which we can't read)
                try:
                    # Confirm the Security Cookie is not zero
                    if self.arch64:
                        cookie = idaapi.dbg_read_memory(self.load_config_directory.SecurityCookie, 8)
                        cookie = struct.unpack("<Q", cookie)[0]
                    else:
                        cookie = idaapi.dbg_read_memory(self.load_config_directory.SecurityCookie, 4)
                        cookie = struct.unpack("<I", cookie)[0]

                    if cookie != 0:
                        return "Yes"
                    else:
                        return "Null"

                except Exception, e:
                    return "Inv"

            else:
                return "No"

        else:
            return "Unk"

###############################################################################
# Pointer Class - Manager writable pointer characteristics

class Ptr():

    def __init__(self, module, ptr_ea, ptr_offset, ptr_charset, call_ea, insn_disas):
        self.module = module

        self.ptr_ea = ptr_ea
        self.ptr_offset = ptr_offset
        self.ptr_charset = ptr_charset

        self.call_ea = call_ea
        self.insn_disas = insn_disas

        self.name = idaapi.get_name(self.call_ea + self.ptr_offset, self.ptr_ea) or ""
        self.call_name = idaapi.get_func_name(self.call_ea) or ""

        self.p2p_ea = None
        self.p2p_offset = None
        self.p2p_charset = None

###############################################################################
# Function Pointer Search Engine

class FuncPtr():

    def __init__(self,sploiter):

        self.sploiter   = sploiter
        self.ptr_calls  = list()
        self.ptrs       = list()

    # NOTE: Pointer Offsets are interpreted as follows:
    # Example 1: Consider we have loaded fptr with a user controlled value
    #
    #              mov [fptr + offset], value
    #
    #            The fptr value we need is (fptr - offset) to overwrite
    #            the actual fptr location. Do this early so that applied address
    #            filters consider fptr bytes before offset adjustment.
    #
    #            Specify this positive offset in the Function Pointer form
    #            so that it would be applied to all discovered function pointers.
    #
    # Example 2: Consider we have a pointer to a pointer write condition:
    #
    #              mov fptr, [p2p + offset2]
    #              mov [fptr + offset1], value
    #
    #            The value (p2p - offset2) must point to the fptr location,
    #            but since we have another offset the target of (p2p - offset2)
    #            must be actually (fptr - offset1). Also do this early so that
    #            applied address filters consider p2p before offset adjustment.
    #
    #            Specify this positive p2p offset2 so it gets applied to the listing
    #            of pointers-to-pointers. All pointer-to-pointer(s) will also
    #            be calculated relative to the ptr - offset1
    #
    # NOTE: Only search for pointers to pointers within the same module. This
    #       was done purely for performance considerations since any readable
    #       pointer to pointer may be used.


    def search_pointers(self):
        
        # HACK: A separate flag is used to track user canceling the search,
        #       because multiple calls to idaapi.wasBreak() do not properly
        #       detect cancellations.
        breakFlag = False

        # Show wait dialog
        idaapi.show_wait_box("Searching writable function pointers...")

        for m in self.modules:
 
            ###################################################################
            # Locate all of the CALL and JMP instructions in the current module
            # which use an immediate operand.

            # List of call/jmp pointer calls in a given module
            ptr_calls = list()

            # Iterate over segments in the module
            # BUG: Iterating over all loaded segments is more stable than looking up by address
            for n in xrange(idaapi.get_segm_qty()):
                seg = idaapi.getnseg(n)

                # Segment in a selected modules
                if seg and seg.startEA >= m.addr and seg.endEA <= (m.addr + m.size):

                    # Locate executable segments
                    # NOTE: Each module may have multiple executable segments
                    # TODO: Search for "MOV REG, PTR # CALL REG"
                    if seg.perm & idaapi.SEGPERM_EXEC:

                        # Search all instances of CALL /2 imm32/64 - FF 15
                        # TODO: Alternative pointer calls using SIB: FF 14 E5 11 22 33 44 - call dword/qword ptr [0x44332211]
                        #                                            FF 14 65 11 22 33 44
                        #                                            FF 14 25 11 22 33 44
                        call_ea = seg.startEA
                        while True:
                            call_ea = idaapi.find_binary(call_ea + 1, seg.endEA, "FF 15", 16, idaapi.SEARCH_DOWN)
                            if call_ea == idaapi.BADADDR: break
                            ptr_calls.append( call_ea )

                        # Search all instances of JMP /2 imm32/64 - FF 25
                        # TODO: Alternative pointer calls using SIB: FF 24 E5 11 22 33 44 - jmp dword/qword ptr [0x44332211]
                        #                                            FF 24 65 11 22 33 44
                        #                                            FF 24 25 11 22 33 44
                        call_ea = seg.startEA
                        while True:
                            call_ea = idaapi.find_binary(call_ea + 1, seg.endEA, "FF 25", 16, idaapi.SEARCH_DOWN)
                            if call_ea == idaapi.BADADDR: break
                            ptr_calls.append( call_ea )

            ###################################################################
            # Extract all of the function pointers and make sure they are
            # are writable.

            # List of writable function pointer objects in a given module
            ptrs = list()

            for call_ea in ptr_calls:

                # Decode CALL/JMP instruction
                # NOTE: May result in invalid disassembly of split instructions
                insn_size = idaapi.decode_insn(call_ea)

                if insn_size:

                    insn = idaapi.cmd
                    insn_op1 = insn.Operands[0].type

                    # Verify first operand is a direct memory reference
                    if insn.Operands[0].type == idaapi.o_mem:

                        # Get operand address
                        ptr_ea = insn.Operands[0].addr

                        # Apply pointer offset
                        ptr_ea -= self.ptrOffset

                        # Locate segment where the pointer is located
                        ptr_seg = idaapi.getseg(ptr_ea)

                        # Make sure a valid segment writeable segment was found
                        if ptr_seg and ptr_seg.perm & idaapi.SEGPERM_WRITE:

                            # Get pointer charset
                            ptr_charset = self.sploiter.get_ptr_charset(ptr_ea)

                            # Filter the pointer
                            if not self.filterP2P:
                                if ptr_charset == None:                                    continue
                                if self.ptrNonull     and not "nonull"     in ptr_charset: continue
                                if self.ptrUnicode    and not "unicode"    in ptr_charset: continue
                                if self.ptrAscii      and not "ascii"      in ptr_charset: continue
                                if self.ptrAsciiPrint and not "asciiprint" in ptr_charset: continue
                                if self.ptrAlphaNum   and not "alphanum"   in ptr_charset: continue
                                if self.ptrNum        and not "numeric"    in ptr_charset: continue
                                if self.ptrAlpha      and not "alpha"      in ptr_charset: continue

                            # Increment the fptr counter

                            # Get pointer disassembly
                            insn_disas = idc.GetDisasmEx(call_ea, idaapi.GENDSM_FORCE_CODE)

                            # Add pointer to the list
                            ptr = Ptr(m.file, ptr_ea, self.ptrOffset, ptr_charset, call_ea, insn_disas)
                            ptrs.append(ptr)

            ###################################################################
            # Cache Pointers to Pointers

            ptr_ea_prefix_cache = dict()

            if self.searchP2P:

                # CACHE: Running repeated searches over the entire memory space is 
                #        very expensive. Let's cache all of the addresses containing
                #        bytes corresponding to discovered function pointers in a
                #        single search and simply reference this cache for each
                #        function pointer. Specifically running idaapi.find_binary()
                #        is much more expensive than idaapi.dbg_read_memory().
                #
                # NOTE:  For performance considerations, the cache works on a per 
                #        module basis, but could be expanded for the entire memory
                #        space.
                #
                # prefix_offset - how many bytes of discovered function
                #        pointers to cache.
                #                 
                #        Example: For function pointers 0x00401234, 0x00404321, 0x000405678
                #        we are going to use prefix_offset 2, so we will cache all of the
                #        values located at addresses 0x0040XXXX

                if self.sploiter.addr64:
                    pack_format   = "<Q"
                    addr_bytes    = 8
                    prefix_offset = 6
                else:
                    pack_format   = "<I"
                    addr_bytes    = 4
                    prefix_offset = 2

                # Set of unique N-byte address prefixes to search in memory
                ea_prefix_set = set()

                for ptr in ptrs:
                    ptr_ea = ptr.ptr_ea

                    ptr_bytes = struct.pack( pack_format, ptr_ea )
                    ptr_bytes = ptr_bytes[-prefix_offset:]

                    ea_prefix_set.add( ptr_bytes)

                # Search the module for all bytes corresponding to the prefix
                # and use them as candidates for pointers-to-pointers

                for ea_prefix in ea_prefix_set:

                    # NOTE: Make sure you search using 44 33 22 11 format and not 11223344
                    ea_prefix_str = " ".join(  ["%02x" % ord(b) for b in ea_prefix]  )

                    # Initialize search parameters for a given module
                    ea = m.addr
                    maxea = m.addr + m.size

                    while True:
                        ea = idaapi.find_binary(ea + 1, maxea, ea_prefix_str, 16, idaapi.SEARCH_DOWN)
                        if ea == idaapi.BADADDR: break

                        p2p_ea = ea - (addr_bytes - prefix_offset)

                        dbg_mem = idaapi.dbg_read_memory(p2p_ea, addr_bytes)
                        ptr_ea_prefix = unpack(pack_format, dbg_mem)[0]

                        if ptr_ea_prefix in ptr_ea_prefix_cache:
                            ptr_ea_prefix_cache[ptr_ea_prefix].add(p2p_ea)
                        else:
                            ptr_ea_prefix_cache[ptr_ea_prefix] = set([p2p_ea,])

                        # Detect search cancellation, but allow the loop below
                        # to run to create already cached/found function pointers

                        # Canceled
                        if breakFlag or idaapi.wasBreak():
                            breakFlag = True
                            break

                    # Canceled
                    if breakFlag or idaapi.wasBreak():
                        breakFlag = True
                        break


            ###################################################################
            # Locate Pointer to Pointers

            for ptr in ptrs:

                ptr_ea = ptr.ptr_ea

                # Locate pointers-to-pointers for a given function pointer in the cache
                if self.searchP2P and ptr_ea in ptr_ea_prefix_cache:

                    for p2p_ea in ptr_ea_prefix_cache[ptr_ea]:

                        # Apply pointer-to-pointer offset
                        p2p_ea -= self.p2pOffset

                        p2p_charset = self.sploiter.get_ptr_charset(p2p_ea)

                        # Filter the pointer
                        if self.filterP2P:
                            if p2p_charset == None:                                    continue
                            if self.ptrNonull     and not "nonull"     in p2p_charset: continue
                            if self.ptrUnicode    and not "unicode"    in p2p_charset: continue
                            if self.ptrAscii      and not "ascii"      in p2p_charset: continue
                            if self.ptrAsciiPrint and not "asciiprint" in p2p_charset: continue
                            if self.ptrAlphaNum   and not "alphanum"   in p2p_charset: continue
                            if self.ptrNum        and not "numeric"    in p2p_charset: continue
                            if self.ptrAlpha      and not "alpha"      in p2p_charset: continue

                        # Copy existing pointer object to modify it for the particular p
                        p2p = copy.copy(ptr)
                        p2p.p2p_ea = p2p_ea
                        p2p.p2p_offset = self.p2pOffset
                        p2p.p2p_charset = p2p_charset

                        # Apppend p2p specific pointer object to the global list
                        self.ptrs.append(p2p)

                        # Exceeded maximum number of pointers
                        if self.maxPtrs and len(self.ptrs) >= self.maxPtrs:
                            breakFlag = True
                            print "[idasploiter] Maximum number of pointers exceeded."
                            break

                # Simply append pointer object to the global list
                else:
                    self.ptrs.append(ptr)

                    # Exceeded maximum number of pointers
                    if self.maxPtrs and len(self.ptrs) >= self.maxPtrs:
                        breakFlag = True
                        print "[idasploiter] Maximum number of pointers exceeded."
                        break

                if breakFlag or idaapi.wasBreak():
                    breakFlag = True
                    break

            # Canceled
            # NOTE: Only works when started from GUI not script.
            if breakFlag or idaapi.wasBreak():
                breakFlag = True
                print "[idasploiter] Canceled."
                break

        print "[idasploiter] Found %d total pointers." % len(self.ptrs)
        idaapi.hide_wait_box()

###############################################################################
# Gadget Class - manages ROP gadget characteristics

class Gadget():

    def __init__(self, instructions, pivot, operations, chg_registers, use_registers):

        self.address = 0x0

        self.module = ""

        self.instructions = instructions
        self.size = len(instructions)

        self.pivot = pivot
        self.operations = operations

        self.chg_registers = chg_registers
        self.use_registers = use_registers

        self.ptr_charset = []

###############################################################################
# ROP Search Engine

class Rop():

    def __init__(self, sploiter):
        
        self.maxRopOffset = 40 # Maximum offset from the return instruction to look for gadgets. default: 40
        self.maxRopSize   = 6  # Maximum number of instructions to look for gadgets. default: 6
        self.maxRetnImm   = 64 # Maximum imm16 value in retn. default: 64
        self.maxJopImm    = 255 # Maximum jop [reg + IMM] value. default: 64
        self.maxRops      = 0  # Maximum number of ROP chains to find. default: 0 (unlimited)

        self.debug        = False

        self.regnames     = idaapi.ph_get_regnames()

        self.sploiter     = sploiter
        self.retns        = list()
        self.gadgets      = list()

        # Decoded instruction cache
        self.insn_cache = dict()

        # Extra bytes to read to ensure correct decoding of
        # RETN, RETN imm16, CALL /2, and JMP /4 instructions.
        self.dbg_read_extra = 6 # FF + ModR/M + SIB + disp32 

        self.insn_arithmetic_ops = ["inc","dec","neg", "add","sub","mul","imul","div","idiv","adc","sbb","lea"]
        self.insn_bit_ops = ["not","and","or","xor","shr","shl","sar","sal","shld","shrd","ror","rcr","rcl"]

    def get_o_reg_name(self, insn, n):

        reg_num = insn.Operands[n].reg
        reg_name = self.regnames[reg_num]

        # NOTE: IDA's x86/x86-64 regname array contains only register root names
        #       (e.g ax,cx,dx,etc.). However we can still figure out exact register
        #       size by looking at the operand 'dtyp' property.
        if reg_num < 8:

            # 32-bit register
            if insn.Operands[n].dtyp == idaapi.dt_dword:
                reg_name = 'e'+reg_name

            # 64-bit register
            elif insn.Operands[n].dtyp == idaapi.dt_qword:
                reg_name = 'r'+reg_name

            # 16-bit register otherwise

        return reg_name

    def search_retns(self):

        for m in self.modules:

            # Iterate over segments in the module
            # BUG: Iterating over all loaded segments is more stable than looking up by address
            for n in xrange(idaapi.get_segm_qty()):
                seg = idaapi.getnseg(n)

                # Locate executable segments in a selected modules
                # NOTE: Each module may have multiple executable segments
                if seg and \
                   seg.startEA >= m.addr and seg.endEA <= (m.addr + m.size) and \
                   seg.perm & idaapi.SEGPERM_EXEC:

                    #######################################################
                    # Search for ROP gadgets
                    if self.searchRop:

                        #Search all instances of RETN 
                        ea = seg.startEA
                        while True:
                            ea = idaapi.find_binary(ea + 1, seg.endEA, "C3", 16, idaapi.SEARCH_DOWN)
                            if ea == idaapi.BADADDR: break
                            self.retns.append( (ea, m.file))

                        # Search all instances of RETN imm16
                        ea = seg.startEA
                        while True:
                            ea = idaapi.find_binary(ea + 1, seg.endEA, "C2", 16, idaapi.SEARCH_DOWN)
                            if ea == idaapi.BADADDR: break

                            # Read imm16 value and filter large values
                            retn_imm16 = idaapi.dbg_read_memory ( ea + 1, 0x2)
                            retn_imm16 = unpack("<H", retn_imm16)[0]

                            if retn_imm16 <= self.maxRetnImm:

                                self.retns.append( (ea, m.file))

                    #######################################################
                    # Search for JOP gadgets
                    if self.searchJop:

                        # Search all instances of JMP reg (FF /4) and CALL reg (FF /2)
                        ea = seg.startEA
                        while True:
                            ea = idaapi.find_binary(ea + 1, seg.endEA, "FF", 16, idaapi.SEARCH_DOWN)
                            if ea == idaapi.BADADDR: break

                            # Read possible ModR/M, SIB, and IMM8/IMM32 bytes
                            jop = idaapi.dbg_read_memory ( ea + 1, 0x6)

                            ###################################################
                            # JMP/CALL reg
                            if jop[0] in ["\xe0","\xe1","\xe2","\xe3","\xe4","\xe5","\xe6","\xe7",
                                          "\xd0","\xd1","\xd2","\xd3","\xd4","\xd5","\xd6","\xd7"]:
                                self.retns.append( (ea, m.file))

                            ###################################################
                            # JMP/CALL [reg] no SIB
                            # NOTE: Do not include pure [disp] instruction.

                            # JMP/CALL [reg] no *SP,*BP
                            elif jop[0] in ["\x20","\x21","\x22","\x23","\x26","\x27", 
                                            "\x10","\x11","\x12","\x13","\x16","\x17"]:
                                self.retns.append( (ea, m.file))

                            # JMP/CALL [reg + imm8] no *SP
                            elif jop[0] in ["\x60","\x61","\x62","\x63","\x65","\x66","\x67",
                                            "\x50","\x51","\x52","\x53","\x55","\x56","\x57"]:
                                jop_imm8 = jop[1]
                                jop_imm8 = unpack("b", jop_imm8)[0] # signed

                                if jop_imm8 <= self.maxJopImm:
                                    self.retns.append( (ea, m.file))


                            # JMP/CALL [reg + imm32] no *SP
                            elif jop[0] in ["\xa0","\xa1","\xa2","\xa3","\xa5","\xa6","\xa7",
                                            "\x90","\x91","\x92","\x93","\x95","\x96","\x97"]:
                                jop_imm32 = jop[1:5]
                                jop_imm32 = unpack("<i", jop_imm32)[0] # signed

                                if jop_imm32 <= self.maxJopImm:
                                    self.retns.append( (ea, m.file))

                            ###################################################
                            # JMP/CALL [reg] with SIB
                            # NOTE: Do no include pure [disp] instructions in SIB ([*] - none)
                            elif (jop[0] in ["\x24","\x64","\xa4"] and not jop[1] in ["\x25","\x65","\xad","\xe5"]) or \
                                 (jop[0] in ["\x14","\x54","\x94"] and not jop[1] in ["\x25","\x65","\xad","\xe5"]):

                                 # Check for displacement
                                if jop[0] in ["\x64","\x54"]:
                                    jop_imm8 = jop[2]
                                    jop_imm8 = unpack("b", jop_imm8)[0] # signed

                                    if jop_imm8 <= self.maxJopImm:
                                        self.retns.append( (ea, m.file))

                                elif jop[0] in ["\xa4","\x94"]:
                                    jop_imm32 = jop[2:6]
                                    jop_imm32 = unpack("<i", jop_imm32)[0] # signed

                                    if jop_imm32 <= self.maxJopImm:
                                        self.retns.append( (ea, m.file))

                                else:
                                    self.retns.append( (ea, m.file))

        print "[idasploiter] Found %d returns" % len(self.retns)

    def search_gadgets(self):

        count_total = len(self.retns)
        count_notify = 0
        count_curr  = 0

        # BUG: A separate flag is used to track user canceling the search,
        #      because multiple calls to idaapi.wasBreak() do not properly
        #      detect cancellations.
        breakFlag = False

        # Show wait dialog
        if not self.debug: idaapi.show_wait_box("Searching gadgets: 00%%%%")

        for (ea_end, module) in self.retns:

            # Flush the gadgets cache for each new retn pointer
            self.gadgets_cache = dict()

            # Flush memory cache for each new retn pointer
            self.dbg_mem_cache = None

            # CACHE: It is faster to read as much memory in one blob than to make incremental reads backwards.
            #        Try to read and cache self.maxRopOffset bytes back. In cases where it is not possible,
            #        then simply try to read the largest chunk.
            
            # NOTE: Read a bit extra to cover correct decoding of RETN, RETN imm16, CALL /2, and JMP /4 instructions.

            for i in range(self.maxRopOffset):
                self.dbg_mem_cache = idaapi.dbg_read_memory(ea_end - self.maxRopOffset + i, self.maxRopOffset - i + self.dbg_read_extra)
                if self.dbg_mem_cache != None:
                    break

            # Search all possible gadgets up to maxoffset bytes back
            # NOTE: Try all byte combinations to capture longer/more instructions
            #       even with bad bytes in the middle.
            for i in range(1, len(self.dbg_mem_cache) - self.dbg_read_extra):

                ea = ea_end - i

                # Get pointer charset
                ptr_charset = self.sploiter.get_ptr_charset(ea)

                # Filter the pointer
                if ptr_charset == None:                                    continue
                if self.ptrNonull     and not "nonull"     in ptr_charset: continue
                if self.ptrUnicode    and not "unicode"    in ptr_charset: continue
                if self.ptrAscii      and not "ascii"      in ptr_charset: continue
                if self.ptrAsciiPrint and not "asciiprint" in ptr_charset: continue
                if self.ptrAlphaNum   and not "alphanum"   in ptr_charset: continue
                if self.ptrNum        and not "numeric"    in ptr_charset: continue
                if self.ptrAlpha      and not "alpha"      in ptr_charset: continue

                # Try to build a gadget at the pointer
                gadget = self.build_gadget(ea, ea_end)

                # Successfully built the gadget
                if gadget:

                    # Populate gadget object with more data
                    gadget.address = ea
                    gadget.module = module
                    gadget.ptr_charset = ptr_charset

                    # Filter gadgets with too many instruction
                    if gadget.size > self.maxRopSize: 
                        break

                    # Append newly built gadget
                    self.gadgets.append(gadget)
                    self.gadgets_cache[ea] = gadget

                    # Exceeded maximum number of gadgets
                    if self.maxRops and len(self.gadgets) >= self.maxRops:
                        breakFlag = True
                        print "[idasploiter] Maximum number of gadgets exceeded."
                        break
                else:
                    self.gadgets_cache[ea] = None

                if breakFlag or idaapi.wasBreak():
                    breakFlag = True
                    break

            # Canceled
            # NOTE: Only works when started from GUI not script.
            if breakFlag or idaapi.wasBreak():
                breakFlag = True
                print "[idasploiter] Canceled."
                break

            # Progress report
            if not self.debug and count_curr >= count_notify:

                # NOTE: Need to use %%%% to escape both Python and IDA's format strings
                idaapi.replace_wait_box("Searching gadgets: %02d%%%%" % (count_curr*100/count_total) ) 

                count_notify += 0.10 * count_total

            count_curr += 1            

        print "[idasploiter] Found %d gadgets." % len(self.gadgets)
        if not self.debug: idaapi.hide_wait_box()


    # Attempt to build a gadget at the provided start address
    # by verifying it properly terminates at the expected RETN.
    def build_gadget(self, ea, ea_end):

        instructions  = list()
        chg_registers = set()
        use_registers = set()
        operations    = set()
        pivot         = 0

        # Process each instruction in the gadget
        while ea <= ea_end:

            ###################################################################
            # Gadget Level Cache:
            #
            # Locate a gadget (failed or built) starting at this address.
            # If one is located, then we don't need to process any further
            # instructions and just get necessary data from the cached
            # gadget to never have to process the same address twice.
            if ea in self.gadgets_cache:

                # Check if the gadget was build successfully
                gadget_cache = self.gadgets_cache[ea]
                
                # Build the reset of the gadget from cache
                if gadget_cache:    

                    for insn in gadget_cache.instructions:
                        instructions.append(insn)

                    for reg in gadget_cache.chg_registers:
                        chg_registers.add(reg)

                    for reg in gadget_cache.use_registers:
                        use_registers.add(reg)

                    for op in gadget_cache.operations:
                        operations.add(op)

                    pivot += gadget_cache.pivot

                    gadget = Gadget(instructions, pivot, operations, chg_registers, use_registers)
                    return gadget

                # Previous attempt to build gadget at this address failed
                else:
                    return None

            # Process new instruction
            else:

                # Instruction length
                # NOTE: decode_insn also sets global idaapi.cmd
                #       which contains insn_t structure
                insn_size = idaapi.decode_insn(ea)

                # Check successful decoding of the instruction
                if insn_size:

                    # Decoded instruction is too big to be a RETN or RETN imm16
                    if ea + insn_size > ea_end + self.dbg_read_extra:
                        return None

                    ###############################################################
                    # Instruction Level Cache
                    #
                    # Most instructions are repetitive so we can just cache
                    # unique byte combinations to avoid costly decoding more
                    # than once
                    
                    # Read instruction from memory cache
                    dbg_mem_offset = ea - (ea_end - (len(self.dbg_mem_cache) - self.dbg_read_extra) )
                    dbg_mem = self.dbg_mem_cache[dbg_mem_offset:dbg_mem_offset + insn_size]

                    # Create instruction cache if it doesn't already exist
                    if not dbg_mem in self.insn_cache:

                        ###########################################################
                        # Decode instruction
                        ###########################################################

                        # Get global insn_t structure describing the instruction
                        # NOTE: copy() is expensive, so we keep this single-threaded
                        insn = idaapi.cmd                  

                        #######################################################
                        # Decode and Cache instruction characteristics
                        self.insn_cache[dbg_mem] = self.decode_instruction(insn, ea, ea_end)

                    ##################################################################
                    # Retrieve cached instruction and apply it to the gadget    

                    # Check that cached instruction contains valid data
                    if self.insn_cache[dbg_mem]:

                        # Retrieve basic instruction characteristics
                        insn_mnem = self.insn_cache[dbg_mem]["insn_mnem"]
                        insn_disas = self.insn_cache[dbg_mem]["insn_disas"]
                        instructions.append(insn_disas)

                        #######################################################
                        # Expected ending instruction of the gadget
                        if ea == ea_end:
                            gadget = Gadget(instructions, pivot, operations, chg_registers, use_registers)
                            return gadget

                        #######################################################
                        # Filter out of place ROP/JOP/COP terminators
                        # NOTE: retn/jmp/call are allowed, but only in the last position

                        # Unexpected return instruction
                        elif insn_mnem == "retn":
                            return None

                        # Unexpected call/jmp instruction
                        elif insn_mnem in ["jmp","call"]:
                            return None

                        #######################################################
                        # Add instruction instruction characteristics to the gadget
                        else:

                            for reg in self.insn_cache[dbg_mem]["insn_chg_registers"]:
                                chg_registers.add(reg)

                            for reg in self.insn_cache[dbg_mem]["insn_use_registers"]:
                                use_registers.add(reg)

                            for op in self.insn_cache[dbg_mem]["insn_operations"]:
                                operations.add(op)

                            pivot += self.insn_cache[dbg_mem]["insn_pivot"]

                    # Previous attempt to decode the instruction invalidated the gadget
                    else:                        
                        return None
                        
                    ###############################################################
                    # Next instruction
                    # NOTE: This is outside cache
                    ea += insn_size

                ###################################################################
                # Failed decoding of the instruction
                # NOTE: Gadgets may have bad instructions in the middle which
                #       can be tolerated as long as we can find a useful instruction
                #       further out.
                else:

                    # HACK: IDA does not disassemble "\x00\x00" unless you enable
                    #       "Disassemble zero opcode instructions" in Processor Options.
                    #       Since this option is normally disabled, I will attempt
                    #       to get this instruction manually.

                    # Read two bytes from memory cache at current instruction candidate
                    dbg_mem_offset = ea - (ea_end - self.maxRopOffset)
                    dbg_mem = self.dbg_mem_cache[dbg_mem_offset:dbg_mem_offset + 2]

                    # Compare to two zero bytes
                    if dbg_mem[:2] == "\x00\x00":

                        if self.sploiter.addr64:
                            instructions.append("add [rax],al")
                        else:
                            instructions.append("add [eax],al")

                        use_registers.add("al")
                        operations.add("reg-to-mem")

                        ea += 2

                    # "MOV Sreg, r/m16" instructions will result in illegal instruction exception: c000001d
                    # or the memory couldn't be read exception: c0000005 which we don't want in our gadgets.
                    elif dbg_mem[0] == "\x8E":
                        return None

                    # Record a "bad byte" if allowed
                    elif dbg_mem and not self.ropNoBadBytes:
                        byte = dbg_mem[0]

                        instructions.append("db %sh" % binascii.hexlify(byte))

                        ea += 1

                    # Invalidate the gadget
                    else:
                        return None

        # Failed to build a gadget, because RETN instruction was not found
        else:    
            return None

    ###############################################################
    # Decode instruction

    def decode_instruction(self, insn, ea, ea_end):

        # Instruction specific characteristics
        insn_chg_registers = set()
        insn_use_registers = set()
        insn_operations = set()
        insn_pivot = 0

        # Instruction feature
        #
        # instruc_t.feature
        #
        # CF_STOP = 0x00001 #  Instruction doesn't pass execution to the next instruction
        # CF_CALL = 0x00002 #  CALL instruction (should make a procedure here)
        # CF_CHG1 = 0x00004 #  The instruction modifies the first operand
        # CF_CHG2 = 0x00008 #  The instruction modifies the second operand
        # CF_CHG3 = 0x00010 #  The instruction modifies the third operand
        # CF_CHG4 = 0x00020 #  The instruction modifies 4 operand
        # CF_CHG5 = 0x00040 #  The instruction modifies 5 operand
        # CF_CHG6 = 0x00080 #  The instruction modifies 6 operand
        # CF_USE1 = 0x00100 #  The instruction uses value of the first operand
        # CF_USE2 = 0x00200 #  The instruction uses value of the second operand
        # CF_USE3 = 0x00400 #  The instruction uses value of the third operand
        # CF_USE4 = 0x00800 #  The instruction uses value of the 4 operand
        # CF_USE5 = 0x01000 #  The instruction uses value of the 5 operand
        # CF_USE6 = 0x02000 #  The instruction uses value of the 6 operand
        # CF_JUMP = 0x04000 #  The instruction passes execution using indirect jump or call (thus needs additional analysis)
        # CF_SHFT = 0x08000 #  Bit-shift instruction (shl,shr...)
        # CF_HLL  = 0x10000 #  Instruction may be present in a high level language function.
        insn_feature = insn.get_canon_feature()

        # Instruction mnemonic name
        insn_mnem = insn.get_canon_mnem()

        #if insn_mnem in self.mnems: self.mnems[insn_mnem] += 1
        #else:                       self.mnems[insn_mnem]  = 1

        # Get instruction operand types
        #
        # op_t.type
        #                    Description                          Data field
        # o_void     =  0 #  No Operand                           ----------
        # o_reg      =  1 #  General Register (al,ax,es,ds...)    reg
        # o_mem      =  2 #  Direct Memory Reference  (DATA)      addr
        # o_phrase   =  3 #  Memory Ref [Base Reg + Index Reg]    phrase
        # o_displ    =  4 #  Memory Reg [Base Reg + Index Reg + Displacement] phrase+addr
        # o_imm      =  5 #  Immediate Value                      value
        # o_far      =  6 #  Immediate Far Address  (CODE)        addr
        # o_near     =  7 #  Immediate Near Address (CODE)        addr
        insn_op1 = insn.Operands[0].type
        insn_op2 = insn.Operands[1].type

        ###############################################################
        # Filter gadget
        ###############################################################

        # Do not filter ROP, JOP, COP, always decode them
        # NOTE: A separate check must be done to check if they are out of place.
        if not insn_mnem in ["retn","jmp","call"]:

            # Filter gadgets with instructions that don't forward execution to the next address
            if insn_feature & idaapi.CF_STOP:
                return None

            # Filter gadgets with instructions in a bad list
            elif insn_mnem in self.ropBadMnems:
                return None

            # Filter gadgets with jump instructions
            # Note: conditional jumps may still be useful if we can
            #       set flags prior to calling them.
            elif not self.ropAllowJcc and insn_mnem[0] == "j":
                return None

        ###############################################################
        # Get disassembly
        ###############################################################
        # NOTE: GENDSM_FORCE_CODE ensures correct decoding
        #       of split instructions.
        insn_disas = idc.GetDisasmEx(ea, idaapi.GENDSM_FORCE_CODE)                
        insn_disas = insn_disas.partition(';')[0]       # Remove comments from disassembly                       
        insn_disas = ' '.join(insn_disas.split())       # Remove extraneous space from disassembly

        ###############################################################
        # Analyze instruction
        ###############################################################

        # Standalone instruction
        if insn_op1 == idaapi.o_void:

            # TODO: Determine and test how these instructions affect the stack
            #       in 32-bit and 64-bit modes.
            if insn_mnem in ["pusha","pushad","popa","popad","pushf","pushfd","pushfq","popf","popfd","popfq"]:
                insn_operations.add("stack")

                if insn_mnem in ["popa","popad"]:
                    insn_pivot += 7*4
                elif insn_mnem in ["pusha","pushad"]:
                    insn_pivot -= 8*4
                elif insn_mnem in ["popf","popfd"]:
                    insn_pivot += 4
                elif insn_mnem in ["pushf","pushfd"]:
                    insn_pivot -= 4
                elif insn_mnem == "popfq":  # TODO: Needs testing
                    insn_pivot += 8
                elif insn_mnem == "pushfq": # TODO: Needs testing
                    insn_pivot -= 8

        # Single operand instruction
        elif insn_op2 == idaapi.o_void:

            # Single operand register
            if insn_op1 == idaapi.o_reg:
                insn_operations.add("one-reg")

                if insn_feature & idaapi.CF_CHG1:
                    reg_name = self.get_o_reg_name(insn, 0)
                    insn_chg_registers.add(reg_name)

                    # Check for stack operation
                    if reg_name[1:] == "sp":
                        insn_operations.add("stack")

                        if insn_mnem == "inc":
                            insn_pivot += 1

                        elif insn_mnem == "dec":
                            insn_pivot -= 1

                elif insn_feature & idaapi.CF_USE1:
                    reg_name = self.get_o_reg_name(insn, 0)
                    insn_use_registers.add(reg_name)

            # Single operand immediate
            elif insn_op1 == idaapi.o_imm:
                insn_operations.add("one-imm")

            # Single operand reference
            # TODO: determine the [reg + ...] value if present
            elif insn_op1 == idaapi.o_phrase or insn_op1 == idaapi.o_displ:
                insn_operations.add("one-mem")

            # PUSH/POP mnemonic with a any operand type
            if insn_mnem in ["push","pop"]:
                insn_operations.add("stack")

                # Adjust pivot based on operand size (32bit vs 64bit)
                if insn_mnem == "pop":
                    if   insn.Operands[0].dtyp == idaapi.dt_dword: insn_pivot += 4
                    elif insn.Operands[0].dtyp == idaapi.dt_qword: insn_pivot += 8
                elif insn_mnem == "push":                   
                    if   insn.Operands[0].dtyp == idaapi.dt_dword: insn_pivot -= 4
                    elif insn.Operands[0].dtyp == idaapi.dt_qword: insn_pivot -= 8

            # Check for arithmetic operation:
            if insn_mnem in self.insn_arithmetic_ops:
                insn_operations.add("math")

            # Check for bit-wise operations:
            if insn_mnem in self.insn_bit_ops:
                insn_operations.add("bit")

        # Two operand instruction
        else:

            # Check for arithmetic operations
            if insn_mnem in self.insn_arithmetic_ops:
                insn_operations.add("math")

            # Check for bit-wise operations
            if insn_mnem in self.insn_bit_ops:
                insn_operations.add("bit")

            # Two operand instruction with the first operand a register
            if insn_op1 == idaapi.o_reg:

                reg_name = self.get_o_reg_name(insn, 0)

                if insn_feature & idaapi.CF_CHG1:
                    insn_chg_registers.add(reg_name)

                    # Check for stack operation
                    if reg_name[1:] == "sp":
                        insn_operations.add("stack")

                        # Determine stack pivot distance
                        if insn_op2 == idaapi.o_imm:

                            # NOTE: adb and sbb may also be useful, but let the user
                            #       determine their use by locating the operations "stack"
                            if insn_mnem == "add":
                                insn_pivot += insn.Operands[1].value

                            elif insn_mnem == "sub":
                                insn_pivot -= insn.Operands[1].value

                    # Check for operations
                    if insn_op2 == idaapi.o_reg:
                        insn_operations.add("reg-to-reg")
                    elif insn_op2 == idaapi.o_imm:
                        insn_operations.add("imm-to-reg")

                    # TODO: determine the [reg + ...] value if present
                    elif insn_op2 == idaapi.o_phrase or insn_op2 == idaapi.o_displ:
                        insn_operations.add("mem-to-reg")

                if insn_feature & idaapi.CF_USE1:
                    insn_use_registers.add(reg_name)


            # Two operand instruction with the second operand a register
            if insn_op2 == idaapi.o_reg:

                reg_name = self.get_o_reg_name(insn, 1)

                if insn_feature & idaapi.CF_CHG2:
                    insn_chg_registers.add(reg_name)

                    # Check for stack operation
                    if reg_name[1:] == "sp":
                        insn_operations.add("stack")

                if insn_feature & idaapi.CF_USE2:
                    insn_use_registers.add(reg_name)

                # Check for operations
                # TODO: determine the [reg + ...] value if present
                if insn_op1 == idaapi.o_phrase or insn_op1 == idaapi.o_displ:
                    insn_operations.add("reg-to-mem")

        # Build instruction dictionary
        insn = dict()
        insn["insn_mnem"] = insn_mnem
        insn["insn_disas"] = insn_disas
        insn["insn_operations"] = insn_operations
        insn["insn_chg_registers"] = insn_chg_registers
        insn["insn_use_registers"] = insn_use_registers
        insn["insn_pivot"] = insn_pivot

        return insn

###############################################################################
# Sploiter Engine

class Sploiter():

    def __init__(self):
        


        # Process modules list
        self.modules = list()
        self.rop     = None
        self.funcptr = None

        # Default patterns
        self.c4_list = string.punctuation
        self.c3_list = string.uppercase
        self.c2_list = string.lowercase
        self.c1_list = string.digits

        if not idaapi.ph.id == idaapi.PLFM_386:
            print "[idasploiter] Only Intel 80x86 processors are supported."
            sys.exit(1)

        # Check if processor supports 64-bit addressing
        if idaapi.ph.flag & idaapi.PR_USE64:
            self.addr64 = True
            self.addr_format = "%016X"
            self.pack_format_be = ">Q"
            self.pack_format_le = "<Q"
        else:
            self.addr64 = False
            self.addr_format = "%08X"
            self.pack_format_be = ">I"
            self.pack_format_le = "<I"

    def get_ptr_charset(self, ea):

        ptr_charset = []

        # Flags. True until proven otherwise.
        nonull     = True
        unicode    = True
        ascii      = True
        asciiprint = True        
        alphanum   = True
        alpha      = True
        numeric    = True
        
        ptr_bytes = pack(self.pack_format_be, ea)
        for i,b in enumerate( ptr_bytes ):

            b_int = ord(b)

            if b in self.ptrBadChars:
                return None

            # Locate any null bytes
            if nonull and b == "\x00":
                nonull = False

            # Unicode compatible address must have
            # null bytes at even points in the address
            if unicode and not i % 2:

                if not ptr_bytes[i:i+2] in ASCII_TO_UNICODE[ self.unicodeTable ]:
                    unicode = False

            # Find any non-ascii characters
            if ascii and not b_int > 127:

                # Find any non-ascii printable characters
                if asciiprint and not (b_int < 32 or b_int > 126):

                    # Find any non-alphanumeric characters
                    if alphanum and (b in string.ascii_letters or b in string.digits):

                        # Find any non-numeric characters
                        if numeric and not b in string.digits:
                            numeric = False

                        # Find any non-letter characters
                        if alpha and not b in string.ascii_letters:
                            alpha = False
                    else:                        
                        alphanum = False
                        numeric  = False
                        alpha     = False
                else:                    
                    asciiprint = False
                    alphanum   = False
                    numeric    = False
                    alpha      = False
            else:                
                ascii      = False
                asciiprint = False
                alphanum   = False
                numeric    = False
                alpha      = False

                # NOTE: You can continue to filter for upper/lower here if necessary
        
        if nonull:     ptr_charset.append("nonull")

        if unicode:    ptr_charset.append("unicode")

        if ascii:      
            ptr_charset.append("ascii")

            if asciiprint: 
                ptr_charset.append("asciiprint")

                if alphanum:   
                    ptr_charset.append("alphanum")

                    if numeric:    ptr_charset.append("numeric")
                    if alpha:      ptr_charset.append("alpha")
        

        return ptr_charset  

    def process_modules(self):

        # Reset modules list
        self.modules = list()

        for m in idautils.Modules():

            module = Module(m.name, m.size, m.base, m.rebase_to)
            self.modules.append(module)

    def show_modules_view(self):
        mod = ModuleView(self)
        mod.show()

    def process_rop(self, select_list = None):

        # Initialize ROP gadget search engine
        self.rop = Rop(self)

        # Prompt user for ROP search settings
        f = RopForm(self, select_list)
        ok = f.Execute()
        if ok == 1:

            # Configure ROP gadget search engine

            # Get selected modules
            self.rop.modules = [self.modules[i] for i in f.mod.GetEmbSelection()]

            if len(self.rop.modules) > 0:

                # Pointer filters
                self.rop.ptrNonull     = f.cPtrNonull.checked
                self.rop.ptrUnicode    = f.cPtrUnicode.checked
                self.rop.ptrAscii      = f.cPtrAscii.checked
                self.rop.ptrAsciiPrint = f.cPtrAsciiPrint.checked
                self.rop.ptrAlphaNum   = f.cPtrAlphaNum.checked
                self.rop.ptrAlpha      = f.cPtrAlpha.checked
                self.rop.ptrNum        = f.cPtrNum.checked

                # Filter bad characters
                buf                    = f.strBadChars.value
                buf = buf.replace(' ','')         # remove spaces
                buf = buf.replace('\\x','')       # remove '\x' prefixes
                buf = buf.replace('0x','')        # remove '0x' prefixes
                try:
                    buf = binascii.unhexlify(buf) # convert to bytes
                    self.ptrBadChars   = buf
                except Exception, e:
                    idaapi.warning("Invalid input: %s" % e)
                    self.ptrBadChars   = ""

                # Ascii_to_Unicode_transformation table
                # BUG: DropdownControl does not work on IDA 6.5
                self.unicodeTable      = f.radUnicode.value

                # ROP instruction filters
                self.rop.ropBadMnems   = [mnem.strip().lower() for mnem in f.strBadMnems.value.split(',')]
                self.rop.ropAllowJcc   = f.cRopAllowJcc.checked
                self.rop.ropNoBadBytes = f.cRopNoBadBytes.checked

                # Get ROP engine settings
                self.rop.maxRopSize    = f.intMaxRopSize.value
                self.rop.maxRopOffset  = f.intMaxRopOffset.value
                self.rop.maxRops       = f.intMaxRops.value
                self.rop.maxRetnImm    = f.intMaxRetnImm.value

                # Gadget search values
                self.rop.searchRop     = f.cRopSearch.checked
                self.rop.searchJop     = f.cJopSearch.checked

                # Search for returns and ROP gadgets
                self.rop.search_retns()
                self.rop.search_gadgets()

                # Show the ROP gadgets view
                ropView = RopView(self)
                ropView.show()

            else:
                idaapi.warning("No modules selected.")

        f.Free()

    def process_funcptr(self, select_list = None):

        # Prompt user for function pointer search settings
        f = PtrForm(self, select_list)
        ok = f.Execute()
        if ok == 1:

            # Initialize Function Pointer search engine
            self.funcptr = FuncPtr(self)

            # Configure Function Pointer search engine

            # Get selected modules
            self.funcptr.modules = [self.modules[i] for i in f.mod.GetEmbSelection()]

            if len(self.funcptr.modules) > 0:

                # Pointer filters
                self.funcptr.ptrNonull     = f.cPtrNonull.checked
                self.funcptr.ptrUnicode    = f.cPtrUnicode.checked
                self.funcptr.ptrAscii      = f.cPtrAscii.checked
                self.funcptr.ptrAsciiPrint = f.cPtrAsciiPrint.checked
                self.funcptr.ptrAlphaNum   = f.cPtrAlphaNum.checked
                self.funcptr.ptrAlpha      = f.cPtrAlpha.checked
                self.funcptr.ptrNum        = f.cPtrNum.checked

                # Apply pointer filters to the funcptr or ptr-to-ptr
                if f.rFilterP2P.selected:
                    self.funcptr.filterP2P = True
                else:
                    self.funcptr.filterP2P = False

                # Ascii_to_Unicode_transformation table
                # BUG: DropdownControl does not work on IDA 6.5
                self.unicodeTable      = f.radUnicode.value

                # Filter bad characters
                buf                    = f.strBadChars.value
                buf = buf.replace(' ','')         # remove spaces
                buf = buf.replace('\\x','')       # remove '\x' prefixes
                buf = buf.replace('0x','')        # remove '0x' prefixes
                try:
                    buf = binascii.unhexlify(buf) # convert to bytes
                    self.ptrBadChars = buf
                except Exception, e:
                    idaapi.warning("Invalid input: %s" % e)
                    self.ptrBadChars = ""

                # Offsets
                self.funcptr.ptrOffset = int(f.intPtrOffset.value, 16)
                self.funcptr.p2pOffset = int(f.intP2POffset.value, 16)

                # Pointer search engine settings
                self.funcptr.searchP2P = not f.rSearchPtrOnly.checked
                self.funcptr.maxPtrs   = f.intMaxPtrs.value

                # Search for writable pointers
                self.funcptr.search_pointers()

                # Show the writable pointers view
                ptrView = PtrView(self)
                ptrView.show()

            else:
                idaapi.warning("No modules selected.")

        f.Free()

    def pattern_create(self):

        f = PatternCreateForm(self)
        ok = f.Execute()
        if ok == 1:  
            pass

        f.Free()

    def pattern_detect(self, debugger=True):

        f = PatternDetectForm(self, debugger)
        ok = f.Execute()
        if ok == 1:  
            pass

        f.Free()

    def process_compare(self, debugger=True):

        f = CompareForm(self, debugger)
        ok = f.Execute()
        if ok == 1:
            pass

        f.Free()


###############################################################################
# Sploiter UI
###############################################################################

###############################################################################
# Module UI
###############################################################################

class ModuleView(Choose2):
    """
    Chooser class to display security characteristics of loaded modules.
    """
    def __init__(self, sploiter, embedded = False):

        self.sploiter = sploiter

        Choose2.__init__(self,
                         "Modules",
                         [ ["Address",  13 | Choose2.CHCOL_HEX], 
                           ["Name",     10 | Choose2.CHCOL_PLAIN], 
                           ["Size",     10 | Choose2.CHCOL_HEX],
                           ["SafeSEH",   6 | Choose2.CHCOL_PLAIN],
                           ["ASLR",      6 | Choose2.CHCOL_PLAIN], 
                           ["DEP",       6 | Choose2.CHCOL_PLAIN],
                           ["Canary",    6 | Choose2.CHCOL_PLAIN],
                           ["Path",     40 | Choose2.CHCOL_PLAIN], 
                         ],
                         flags = Choose2.CH_MULTI,  # Select multiple modules
                         embedded=embedded)

        self.icon = 150

        # Items for display
        self.items = list()

        # Initialize/Refresh the view
        self.refreshitems()

        # Selected items
        self.select_list = list()

        # Command callbacks
        self.cmd_load_module     = None
        self.cmd_search_gadgets  = None
        self.cmd_search_pointers = None

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

        # Add extra context menu commands
        # NOTE: Make sure you check for duplicates
        if self.cmd_load_module == None:
            self.cmd_load_module = self.AddCommand("Load module...", flags = idaapi.CHOOSER_POPUP_MENU, icon=135)
        if self.cmd_search_gadgets == None:
            self.cmd_search_gadgets = self.AddCommand("Search gadgets...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=182 )
        if self.cmd_search_pointers == None:
            self.cmd_search_pointers = self.AddCommand("Search function pointers...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=143 )


        return True

    def refreshitems(self):
        self.items = list()

        for m in self.sploiter.modules:
            self.items.append([ self.sploiter.addr_format % m.addr, 
                                m.file, 
                                "%08X" % m.size,
                                m.SafeSEH, m.ASLR,
                                m.NXCompat,
                                m.GS,
                                m.path])

    def OnCommand(self, n, cmd_id):

        # Search ROP gadgets
        if cmd_id == self.cmd_search_gadgets:

            # Empty selection
            if n == -1:

                # Initialize ROP gadget form with empty selection
                self.sploiter.process_rop(select_list = self.select_list)

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:

                # Initialize ROP gadget form with user selection
                self.sploiter.process_rop(select_list = self.select_list)
                self.select_list = list()

            # Selection number
            else:
                self.select_list.append(n)

        # Search function pointers
        elif cmd_id == self.cmd_search_pointers:

            # Empty selection
            if n == -1:

                # Initialize function pointers form with empty selection
                self.sploiter.process_funcptr(select_list = self.select_list)

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:

                # Initialize function pointers form with user selection
                self.sploiter.process_funcptr(select_list = self.select_list)
                self.select_list = list()

            # Selection number
            else:
                self.select_list.append(n)

        elif cmd_id == self.cmd_load_module:

            module_name = idaapi.askfile_c(0, "*.*", "Please select module to load")
            if module_name:
                print "[idasploiter] Loading module: %s" % module_name     
                loadlib = idaapi.Appcall.proto("kernel32_LoadLibraryA", "int __stdcall loadlib(const char *fn);")
                hmod = loadlib(module_name)
                if hmod:
                    print "[idasploiter] Finished loading module: %s" % module_name
                else:
                    print "[idasploiter] Could not load: %s" % module_name

                self.refreshitems()

        return 1

    def OnSelectLine(self, n):
        pass

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

    def OnGetIcon(self, n):

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

        m = self.sploiter.modules[n]
        if m.SafeSEH != "No" and m.ASLR == "Yes" and m.NXCompat == "Yes" and m.GS != "No":
            return 61
        elif m.ASLR == "Yes":
            return 60
        else:
            return 59

    def OnClose(self):
        self.cmd_search_gadgets = None
        self.cmd_search_pointers    = None

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

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

    def OnActivate(self):
        self.refreshitems()

###############################################################################
# ROP UI
###############################################################################

###############################################################################
# ROP/JOP/COP Form

class RopForm(Form):

    def __init__(self, sploiter, select_list = None):

        self.sploiter = sploiter
        self.select_list = select_list

        self.mod = ModuleView(self.sploiter, embedded=True)

        Form.__init__(self, 
r"""BUTTON YES* Search
Search ROP gadgets

{FormChangeCb}<Modules:{cEChooser}>

Pointer Charset:                  Search Settings:
<nonull:{cPtrNonull}>        <Bad Chars        :{strBadChars}>     
<unicode:{cPtrUnicode}>       Unicode Table    <ANSI:{rUnicodeANSI}><OEM:{rUnicodeOEM}><UTF7:{rUnicodeUTF7}><UTF8:{rUnicodeUTF8}>{radUnicode}>
<ascii:{cPtrAscii}>         <Bad Instructions :{strBadMnems}>
<asciiprint:{cPtrAsciiPrint}>    <Max gadget size  :{intMaxRopSize}>
<alphanum:{cPtrAlphaNum}>      <Max gadget offset:{intMaxRopOffset}>
<alpha:{cPtrAlpha}>         <Max RETN imm16   :{intMaxRetnImm}>
<numeric:{cPtrNum}>{ptrGroup}>       <Max JOP imm8/32  :{intMaxJopImm}>
                <Max gadgets      :{intMaxRops}>
                Other settings   <Allow conditional jumps:{cRopAllowJcc}>
                <Do not allow bad bytes:{cRopNoBadBytes}>
                <Search for ROP gadgets:{cRopSearch}>
                <Search for JOP gadgets:{cJopSearch}>{ropGroup}>


""", {
                'cEChooser'       : Form.EmbeddedChooserControl(self.mod, swidth=131),
                'ptrGroup'        : Form.ChkGroupControl(("cPtrNonull", "cPtrAscii", "cPtrAsciiPrint", "cPtrUnicode",'cPtrAlphaNum','cPtrAlpha','cPtrNum')),
                'ropGroup'        : Form.ChkGroupControl(('cRopAllowJcc','cRopNoBadBytes','cRopSearch','cJopSearch')),
                'intMaxRopSize'   : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=self.sploiter.rop.maxRopSize),
                'intMaxRopOffset' : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=self.sploiter.rop.maxRopOffset),
                'intMaxRops'      : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=self.sploiter.rop.maxRops),
                'intMaxRetnImm'   : Form.NumericInput(swidth=4,tp=Form.FT_HEX,value=self.sploiter.rop.maxRetnImm),
                'intMaxJopImm'    : Form.NumericInput(swidth=4,tp=Form.FT_HEX,value=self.sploiter.rop.maxJopImm),
                'strBadChars'     : Form.StringInput(swidth=70,tp=Form.FT_ASCII),
                'radUnicode'      : Form.RadGroupControl(("rUnicodeANSI","rUnicodeOEM","rUnicodeUTF7","rUnicodeUTF8")),
                'strBadMnems'     : Form.StringInput(swidth=80,tp=Form.FT_ASCII,value="leave, int, into, enter, syscall, sysenter, sysexit, sysret, in, out, loop, loope, loopne, lock, rep, repe, repz, repne, repnz"),
                'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.cEChooser)

            # Preselect non-ASLR modules on startup if none were already specified
            if self.select_list == None:

                self.select_list = list()
                for i, m in enumerate(self.sploiter.modules):
                    if m.ASLR == "No":
                        self.select_list.append(i)

            self.SetControlValue(self.cEChooser, self.select_list)

            # Enable both ROP and JOP search by default
            self.SetControlValue(self.cRopSearch, True)
            self.SetControlValue(self.cJopSearch, True)

            # Skip bad instructions by default
            self.SetControlValue(self.cRopNoBadBytes, True)

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

        return 1

###############################################################################
# ROP Viewer

class RopView(Choose2):
    """
    Chooser class to display security characteristics of loaded modules.
    """
    def __init__(self, sploiter):

        self.sploiter = sploiter

        Choose2.__init__(self,
                         "ROP gadgets",
                         [ ["Address",           13 | Choose2.CHCOL_HEX], 
                           ["Gadget",            30 | Choose2.CHCOL_PLAIN], 
                           ["Module",            10 | Choose2.CHCOL_PLAIN],
                           ["Size",               3 | Choose2.CHCOL_DEC],
                           ["Pivot",              4 | Choose2.CHCOL_DEC],
                           ["Operations",         12 | Choose2.CHCOL_PLAIN],
                           ["Changed Registers", 12 | Choose2.CHCOL_PLAIN],
                           ["Used Registers",    12 | Choose2.CHCOL_PLAIN],
                           ["Charset",           12 | Choose2.CHCOL_PLAIN],
                         ],
                         flags = Choose2.CH_MULTI)

        self.icon = 182

        # Items for display
        self.items = []

        # Initialize/Refresh the view
        self.refreshitems()

        # Selected items
        self.gadget_chain = list()

        # Command callbacks
        self.cmd_chain_add   = None
        self.cmd_chain_build = None
        self.cmd_chain_clear = None
        self.cmd_export_csv  = None


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

        if self.cmd_chain_add == None:
            self.cmd_chain_add   = self.AddCommand("Add to chain",     flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=50)
        if self.cmd_chain_build == None:
            self.cmd_chain_build = self.AddCommand("Build chain",      flags = idaapi.CHOOSER_POPUP_MENU, icon=156)
        if self.cmd_chain_clear == None:
            self.cmd_chain_clear = self.AddCommand("Clear chain",      flags = idaapi.CHOOSER_POPUP_MENU, icon=32)
        if self.cmd_export_csv == None:
            self.cmd_export_csv  = self.AddCommand("Export as csv...", flags = idaapi.CHOOSER_POPUP_MENU, icon=40)

        return True

    def refreshitems(self):
        self.items = []

        for g in self.sploiter.rop.gadgets:

            self.items.append([ self.sploiter.addr_format % g.address, 
                                " # ".join(g.instructions), 
                                g.module, 
                                "%d" % g.size, 
                                "%d" % g.pivot, 
                                ", ".join(g.operations), 
                                ", ".join(g.chg_registers), 
                                ", ".join(g.use_registers), 
                                ", ".join(g.ptr_charset)
                                ])

    def OnCommand(self, n, cmd_id):

        # Build ROP chain
        if cmd_id == self.cmd_chain_build:

            # Display ROP chain builder form
            f = RopChainForm(self)
            ok = f.Execute()
            if ok == 1:
                pass
            f.Free()

        # Clear ROP chain
        elif cmd_id == self.cmd_chain_clear:
            self.gadget_chain = list()

        # Add gadget to ROP chain
        elif cmd_id == self.cmd_chain_add:
            if n >= 0:
                self.gadget_chain.append( self.sploiter.rop.gadgets[n] )

        # Export CSV
        elif cmd_id == self.cmd_export_csv:

            file_name = idaapi.askfile_c(1, "*.csv", "Please enter CSV file name")
            if file_name:
                print "[idasploiter] Exporting gadgets to %s" % file_name
                with open(file_name, 'wb') as csvfile:
                    csvwriter = csv.writer(csvfile, delimiter=',',
                                            quotechar='"', quoting=csv.QUOTE_MINIMAL)
                    csvwriter.writerow(["Address","Gadget","Module","Size","Pivot","Operations","Changed Registers","Used Registers","Charset"])
                    for item in self.items:
                        csvwriter.writerow(item)

        return 1

    def OnSelectLine(self, n):
        idaapi.jumpto( self.sploiter.rop.gadgets[n].address )

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

    def OnClose(self):
        self.cmd_chain_add   = None
        self.cmd_chain_build = None
        self.cmd_chain_clear = None
        self.cmd_export_csv  = None

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

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

    def OnActivate(self):
        self.refreshitems()

###############################################################################
# ROP Chain Builder UI
###############################################################################

###############################################################################
# ROP Chain Viewer

class RopChainView(Choose2):
    """
    Chooser class to display security characteristics of loaded modules.
    """
    def __init__(self, ropview):

        self.ropview = ropview

        Choose2.__init__(self,
                         "ROP gadgets",
                         [ ["Address",           13 | Choose2.CHCOL_HEX], 
                           ["Gadget",            30 | Choose2.CHCOL_PLAIN], 
                           ["Module",            10 | Choose2.CHCOL_PLAIN],
                           ["Size",               3 | Choose2.CHCOL_DEC],
                           ["Pivot",              4 | Choose2.CHCOL_DEC],
                           ["Operations",        12 | Choose2.CHCOL_PLAIN],
                           ["Changed Registers", 12 | Choose2.CHCOL_PLAIN],
                           ["Used Registers",    12 | Choose2.CHCOL_PLAIN],
                           ["Charset",           12 | Choose2.CHCOL_PLAIN],
                         ],
                         flags = Choose2.CH_MULTI,
                         embedded = True)

        self.icon = 182

        # Items for display and corresponding data
        self.items = []

        # Selected items
        self.select_list = list()

        # Initialize/Refresh the view
        self.refreshitems()

    def refreshitems(self):
        self.items = []

        for g in self.ropview.gadget_chain:

            self.items.append([ self.ropview.sploiter.addr_format % g.address, 
                                " # ".join(g.instructions), 
                                g.module, 
                                "%d" % g.size if g.size else "", 
                                "%d" % g.pivot if g.pivot else "", 
                                ", ".join(g.operations), 
                                ", ".join(g.chg_registers), 
                                ", ".join(g.use_registers), 
                                ", ".join(g.ptr_charset)
                                ])

        return

    def OnCommand(self, n, cmd_id):

        if cmd_id == self.cmd_move_up:
            if n > 0:
                self.ropview.gadget_chain.insert(n-1, self.ropview.gadget_chain.pop(n))

        # NOTE: Moving items down should be done in one go in reverse
        #       to avoid self.items from desyncing from gadget_chain list.
        elif cmd_id == self.cmd_move_down:

            # Empty selection
            if n == -1:
                pass

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:
                for i in sorted(self.select_list, reverse=True):
                    self.ropview.gadget_chain.insert(i+1, self.ropview.gadget_chain.pop(i))
                self.select_list = list()

            else:
                self.select_list.append(n)

        # NOTE: Duplicating items should be done in one go in reverse
        #       to avoid self.items desyncing from gadget_chain list.
        elif cmd_id == self.cmd_duplicate:
            
            # Empty selection
            if n == -1:
                pass

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:

                # Duplicate selection after the last index
                last_i = self.select_list[-1]

                # Duplicate selected items in reverse
                for i in sorted(self.select_list, reverse=True):

                    if last_i == len(self.ropview.gadget_chain) - 1:
                        self.ropview.gadget_chain.append(self.ropview.gadget_chain[i])
                    else:
                        self.ropview.gadget_chain.insert(last_i + 1, self.ropview.gadget_chain[i])

                self.select_list = list()

            else:
                self.select_list.append(n)

        # NOTE: Deleting items should be done in one go in reverse
        #       to avoid self.items desyncing from gadget_chain list.
        elif cmd_id == self.cmd_delete:
            
            # Empty selection
            if n == -1:
                pass

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:
                for i in sorted(self.select_list, reverse=True):
                    del self.ropview.gadget_chain[i]
                self.select_list = list()

            else:
                self.select_list.append(n)



        return 1

    def OnSelectLine(self, n):
        pass

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

    def OnClose(self):
        pass

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

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

    def OnRefreshed(self):
        self.refreshitems()
        return len(self.items)

    def OnActivate(self):
        self.refreshitems()


###############################################################################
# ROPChain Form

class RopChainForm(Form):

    def __init__(self, ropview):

        self.ropview = ropview

        self.ropchainview = RopChainView(self.ropview)

        Form.__init__(self, 
r"""BUTTON YES* NONE
BUTTON CANCEL NONE
ROP Chain Builder

{FormChangeCb}<ROP gadgets:{cEChooser}>
<:{strAddr}><:{strComment}><Insert:{iButtonInsert}>
<:{strRopChain}>
Formats<Plain:{rFmtPlain}><Python:{rFmtPython}><Ruby:{rFmtRuby}><Perl:{rFmtPerl}><JavaScript:{rFmtJS}>{rGroup}> 
{spacer}<Generate:{iButtonGenerate}>
""", {
                'spacer'          : Form.StringLabel("                                "),
                'cEChooser'       : Form.EmbeddedChooserControl(self.ropchainview, swidth=125),
                'strAddr'         : Form.StringInput(value="0x4141414141414141" if self.ropview.sploiter.addr64 else "0x41414141", swidth=10),
                'strComment'      : Form.StringInput(value="padding", swidth=20),
                'iButtonInsert'   : Form.ButtonInput(self.OnButtonInsert, swidth=5),
                'strRopChain'     : Form.MultiLineTextControl(text="", swidth=125, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT),
                'rGroup'          : Form.RadGroupControl(("rFmtPlain","rFmtPython", "rFmtRuby", "rFmtPerl", "rFmtJS") ),
                'iButtonGenerate' : Form.ButtonInput(self.OnButtonGenerate, swidth=5),
                'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.cEChooser)

            # Add commands to embedded view
            self.ropchainview.cmd_move_up   = self.cEChooser.AddCommand("Move Up", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION | idaapi.CHOOSER_HOTKEY, icon=86)
            self.ropchainview.cmd_move_down = self.cEChooser.AddCommand("Move Down", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=85)
            self.ropchainview.cmd_duplicate = self.cEChooser.AddCommand("Duplicate", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=111)
            self.ropchainview.cmd_delete    = self.cEChooser.AddCommand("Delete", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=112)

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

        return 1

    def OnButtonInsert(self, code=0):
        # Get user input
        address = self.GetControlValue(self.strAddr)
        address = int(address,16)

        comment = self.GetControlValue(self.strComment)

        # Create a dummy gadget
        gadget = Gadget(instructions = [comment,], pivot = 0, operations = [], chg_registers = [], use_registers = [])
        gadget.address = address
        gadget.size = None
        gadget.pivot = None

        # Insert dummy gadget at selected positions
        select_list = self.GetControlValue(self.cEChooser)
        if select_list:
            for i in sorted(select_list,reverse=True):
                self.ropview.gadget_chain.insert(i + 1, gadget)
        else:
            self.ropview.gadget_chain.append(gadget)

        # Update the view
        self.RefreshField(self.cEChooser)

    # TODO: Generate relative addresses to a dynamically determined module base
    #       (e.g. ROP chain built based on a leaked or calculated base address)
    def OnButtonGenerate(self, code=0):

        ropchain = ""

        if self.ropview.sploiter.addr64:
            addr_format = "%016X"
            pack_format_python = "<Q"
            pack_format_ruby = "Q<"
            pack_format_perl = "Q<"
            pack_words_js = 4
            pack_bytes_c = 8

        else:
            addr_format = "%08X"
            pack_format_python = "<I"
            pack_format_ruby = "V"
            pack_format_perl = "V"
            pack_words_js = 2
            pack_bytes_c = 4

        # Plain output
        if self.GetControlValue(self.rGroup) == 0:

            ropchain += "// Generated by IDA Sploiter\n"

            for g in self.ropview.gadget_chain:

                bytes =  ["\\x%02x" % b for b in struct.unpack("B"*pack_bytes_c, struct.pack(pack_format_python, g.address))]

                ropchain += "\"%s\" // %s # %s\n" % (
                                ''.join(bytes),
                                addr_format % g.address,
                                " # ".join(g.instructions))

        # Python
        elif self.GetControlValue(self.rGroup) == 1:

            ropchain += "# Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                ropchain += "rop_chain %s= struct.pack('%s', 0x%s) # %s\n" % (
                                '+' if i > 0 else ' ',
                                pack_format_python,
                                addr_format % g.address,
                                " # ".join(g.instructions) )

        # Ruby
        elif self.GetControlValue(self.rGroup) == 2:

            ropchain += "# Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                ropchain += "rop_chain %s [0x%s].pack('%s') # %s\n" % (
                                '<<' if i > 0 else ' =',
                                addr_format % g.address,
                                pack_format_ruby,
                                " # ".join(g.instructions) )

        # Perl
        elif self.GetControlValue(self.rGroup) == 3:

            ropchain += "# Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                ropchain += "%s$rop_chain %s= pack('%s', 0x%s); # %s\n" % (
                                'my ' if i == 0 else '',
                                '.' if i > 0 else ' ',
                                pack_format_perl,
                                addr_format % g.address,
                                " # ".join(g.instructions) )

        # JavaScript
        elif self.GetControlValue(self.rGroup) == 4:

            ropchain += "// Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                words_js = ["%%u%04x" % w for w in struct.unpack("H"*pack_words_js, struct.pack(pack_format_python, g.address))]

                ropchain += "rop_chain %s= unescape(\"%s\"); // %s # %s\n" % (
                                '+' if i > 0 else ' ',
                                ''.join(words_js),
                                addr_format % g.address,
                                " # ".join(g.instructions) )

        strRopChain = idaapi.textctrl_info_t(text=ropchain, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT)
        self.SetControlValue(self.strRopChain, strRopChain)

###############################################################################
# Function Pointer UI
###############################################################################

###############################################################################
# Function Pointer Form

class PtrForm(Form):

    def __init__(self, sploiter, select_list = None):

        self.sploiter = sploiter
        self.select_list = select_list
        self.mod = ModuleView(self.sploiter, embedded=True)

        Form.__init__(self, 
r"""BUTTON YES* Search
Search writable function pointers

{FormChangeCb}<Modules:{cEChooser}>

Pointer Charset:                  Search Settings:
<nonull:{cPtrNonull}>        <Bad chars     :{strBadChars}>
<unicode:{cPtrUnicode}>       Unicode Table <ANSI:{rUnicodeANSI}><OEM:{rUnicodeOEM}><UTF7:{rUnicodeUTF7}><UTF8:{rUnicodeUTF8}>{radUnicode}>
<ascii:{cPtrAscii}>         <PTR Offset    :{intPtrOffset}> 
<asciiprint:{cPtrAsciiPrint}>    <P2P Offset    :{intP2POffset}>
<alphanum:{cPtrAlphaNum}>      <Max Pointers  :{intMaxPtrs}>
<alpha:{cPtrAlpha}>         Filter        <Function Pointers:{rFilterPtr}><Pointers-to-Pointers:{rFilterP2P}>{rGroup}>
<numeric:{cPtrNum}>{ptrGroup}>       Other settings<Do not search for pointers-to-pointers:{rSearchPtrOnly}>{rSearchGroup}>


""", {
                'cEChooser'       : Form.EmbeddedChooserControl(self.mod, swidth=131),
                'ptrGroup'        : Form.ChkGroupControl(("cPtrNonull", "cPtrAscii", "cPtrAsciiPrint", "cPtrUnicode",'cPtrAlphaNum','cPtrAlpha','cPtrNum')),
                'rGroup'          : Form.RadGroupControl(("rFilterPtr","rFilterP2P") ),
                'rSearchGroup'    : Form.ChkGroupControl(("rSearchPtrOnly",) ),
                'strBadChars'     : Form.StringInput(swidth=70,tp=Form.FT_ASCII),
                'radUnicode'      : Form.RadGroupControl(("rUnicodeANSI","rUnicodeOEM","rUnicodeUTF7","rUnicodeUTF8")),
                'intPtrOffset'    : Form.StringInput(tp=Form.FT_ASCII,value="0x0", swidth=10),
                'intP2POffset'    : Form.StringInput(tp=Form.FT_ASCII,value="0x0", swidth=10),
                'intMaxPtrs'      : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=0),
                'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.cEChooser)

            # Preselect non-ASLR modules on startup if none were already specified
            if self.select_list == None:
                
                self.select_list = list()
                for i, m in enumerate(self.sploiter.modules):
                    if m.ASLR == "No":
                        self.select_list.append(i)

            self.SetControlValue(self.cEChooser, self.select_list)

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

        return 1

###############################################################################
# Function Pointer Viewer

class PtrView(Choose2):
    """
    Chooser class to display writable function pointers.
    """
    def __init__(self, sploiter):

        self.sploiter = sploiter

        Choose2.__init__(self,
                         "Writeable function pointers",
                         [ ["Pointer",            13 | Choose2.CHCOL_HEX], 
                           ["Offset",              4 | Choose2.CHCOL_HEX], 
                           ["Name",               10 | Choose2.CHCOL_PLAIN],
                           ["Module",             10 | Choose2.CHCOL_PLAIN],
                           ["Charset",            12 | Choose2.CHCOL_PLAIN],
                           ["Caller Address",     13 | Choose2.CHCOL_HEX], 
                           ["Caller Name",        10 | Choose2.CHCOL_PLAIN],
                           ["Instruction",        15 | Choose2.CHCOL_PLAIN],
                           ["Ptr-to-Ptr",         13 | Choose2.CHCOL_PLAIN],
                           ["Ptr-to-Ptr Offset",   4 | Choose2.CHCOL_HEX], 
                           ["Ptr-to-Ptr Charset", 12 | Choose2.CHCOL_PLAIN],
                         ],
                         flags = Choose2.CH_MULTI_EDIT)

        self.icon = 143

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

        # Initialize/Refresh the view
        self.refreshitems()

        # Selected items
        self.select_caller_breakpoint_list = list()

        # Command callbacks
        self.cmd_select_caller_breakpoint = None
        self.cmd_jump_ptr = None
        self.cmd_jump_p2p = None
        self.cmd_jump_caller = None
        self.cmd_export_csv  = None

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

        # Add extra context menu commands
        # NOTE: Make sure you check for duplicates
        if self.cmd_select_caller_breakpoint == None:
            self.cmd_select_caller_breakpoint = self.AddCommand("Add breakpoint", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=120)

        if self.cmd_jump_ptr == None:
            self.cmd_jump_ptr = self.AddCommand("Jump to pointer", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=141)

        if self.cmd_jump_p2p == None:
            self.cmd_jump_p2p = self.AddCommand("Jump to pointer-to-pointer", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=142)

        if self.cmd_jump_caller == None:
            self.cmd_jump_caller = self.AddCommand("Jump to caller", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=143)

        if self.cmd_export_csv == None:
            self.cmd_export_csv = self.AddCommand("Export as csv...", flags = idaapi.CHOOSER_POPUP_MENU, icon=40)

        return True

    def refreshitems(self):
        self.items = []

        for ptr in self.sploiter.funcptr.ptrs:

            self.items.append([ self.sploiter.addr_format % ptr.ptr_ea, 
                                hex(ptr.ptr_offset),
                                ptr.name, 
                                ptr.module, 
                                ", ".join(ptr.ptr_charset),
                                self.sploiter.addr_format % ptr.call_ea,
                                ptr.call_name,
                                ptr.insn_disas,      
                                self.sploiter.addr_format % ptr.p2p_ea if ptr.p2p_ea else "",
                                hex(ptr.p2p_offset) if ptr.p2p_ea else "",
                                ", ".join(ptr.p2p_charset) if ptr.p2p_ea else "",                          
                                ])

    def OnCommand(self, n, cmd_id):

        if cmd_id == self.cmd_select_caller_breakpoint:

            # Empty selection
            if n == -1:
                pass                

            # Selection start
            elif n == -2:
                self.select_caller_breakpoint_list = list()

            # Selection end
            elif n == -3:               

                # Unique set of caller addresses
                eas = set()

                for i in self.select_caller_breakpoint_list:
                    eas.add( self.sploiter.funcptr.ptrs[i].call_ea )

                print "[idasploiter] Settings %d breakpoints on selected caller addresses." % len(eas)

                for ea in eas:
                    # Set default software breakpoint
                    idaapi.add_bpt(ea, 0, idaapi.BPT_DEFAULT)

                self.select_caller_breakpoint_list = list()

            # Selection number
            else:
                self.select_caller_breakpoint_list.append(n)

        # Jump to Pointer address
        elif cmd_id == self.cmd_jump_ptr:
            if n >= 0 and self.sploiter.funcptr.ptrs[n].ptr_ea:
                idaapi.jumpto( self.sploiter.funcptr.ptrs[n].ptr_ea )

        # Jump to Pointer-to-Pointer address
        elif cmd_id == self.cmd_jump_p2p:
            if n >= 0 and self.sploiter.funcptr.ptrs[n].p2p_ea:
                idaapi.jumpto( self.sploiter.funcptr.ptrs[n].p2p_ea )

        # Jump to Caller address
        elif cmd_id == self.cmd_jump_caller:
            if n >= 0 and self.sploiter.funcptr.ptrs[n].call_ea:
                idaapi.jumpto( self.sploiter.funcptr.ptrs[n].call_ea )

        # Export CSV
        elif cmd_id == self.cmd_export_csv:

            file_name = idaapi.askfile_c(1, "*.csv", "Please enter CSV file name")
            if file_name:
                print "[idasploiter] Exporting function pointers to %s" % file_name
                with open(file_name, 'wb') as csvfile:
                    csvwriter = csv.writer(csvfile, delimiter=',',
                                            quotechar='"', quoting=csv.QUOTE_MINIMAL)
                    csvwriter.writerow(["Pointer","Offset","Name","Module","Charset","Caller Address","Caller Name","Instruction","Ptr-to-Ptr","Ptr-to-Ptr Offset","Ptr-to-Ptr Charset"])
                    for item in self.items:
                        csvwriter.writerow(item)
        return 1

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

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

    def OnClose(self):
        self.cmd_select_caller_breakpoint = None
        self.cmd_jump_ptr = None
        self.cmd_jump_p2p = None
        self.cmd_jump_caller = None
        self.cmd_export_csv  = None

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

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

    def OnActivate(self):
        self.refreshitems()

###############################################################################
# Pattern UI
###############################################################################

###############################################################################
# Pattern Create Form

class PatternCreateForm(Form):

    def __init__(self, sploiter):

        self.sploiter = sploiter
        self.pattern_complete = ""

        Form.__init__(self, 
r"""BUTTON YES NONE
BUTTON CANCEL NONE
Create pattern
{FormChangeCb}<:{strPattern}>
<Size   :{intSize}>
Formats<Regular:{rFmtReg}><Hex:{rFmtHex}><JavaScript:{rFmtJS}>{rGroup}>
Charset<C4:{cC4}><##:{strC4}>
       <C3:{cC3}><##:{strC3}>
       <C2:{cC2}><##:{strC2}>
       <C1:{cC1}>{cGroup}><##:{strC1}>

{spacer}<Generate:{iButtonGenerate}>
""", {
            'spacer'          : Form.StringLabel("            "),
            'intSize'         : Form.NumericInput(swidth=30,tp=Form.FT_DEC,value=1000),
            'iButtonGenerate' : Form.ButtonInput(self.OnButtonGenerate,swidth=5),
            'rGroup'          : Form.RadGroupControl(("rFmtReg","rFmtHex", "rFmtJS") ),
            'cGroup'          : Form.ChkGroupControl(("cC4","cC3","cC2","cC1")),
            'strC4'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC3'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC2'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC1'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strPattern'      : Form.MultiLineTextControl(text="", width=400, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_SELECTED | Form.MultiLineTextControl.TXTF_READONLY),
            'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:

            # Set initial checkboxes for the classic metasploit 3-byte pattern
            self.SetControlValue(self.cC4, False)
            self.SetControlValue(self.cC3, True)
            self.SetControlValue(self.cC2, True)
            self.SetControlValue(self.cC1, True)

            # Set initial charsets values
            self.SetControlValue(self.strC4, self.sploiter.c4_list)
            self.SetControlValue(self.strC3, self.sploiter.c3_list)
            self.SetControlValue(self.strC2, self.sploiter.c2_list)
            self.SetControlValue(self.strC1, self.sploiter.c1_list)

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

        return 1

    def OnButtonGenerate(self, code=0):

        size = self.GetControlValue(self.intSize) 
        print "[idasploiter] Creating a pattern of size %d." % size

        # Update global charset lists
        self.sploiter.c4_list = self.GetControlValue(self.strC4)
        self.sploiter.c3_list = self.GetControlValue(self.strC3)
        self.sploiter.c2_list = self.GetControlValue(self.strC2)
        self.sploiter.c1_list = self.GetControlValue(self.strC1)

        # Generate pattern only using selected lists
        self.update_complete_pattern(self.GetControlValue(self.cC4), self.GetControlValue(self.cC3), self.GetControlValue(self.cC2), self.GetControlValue(self.cC1))

        pattern = self.pattern_create(size) 

        if pattern:

            # Hex output
            if self.GetControlValue(self.rGroup) == 1:
                pattern = binascii.hexlify(pattern)

            # JavaScript Unicode output
            elif self.GetControlValue(self.rGroup) == 2:
                
                pattern_js = ""

                # Number of WORDs in the pattern
                num_words = len(pattern) / 2

                # Chop off extra characters
                num_extra = len(pattern) % 2
                if num_extra:
                    pattern = pattern[:-num_extra]

                # Unpack individual WORDs from the pattern
                words = struct.unpack("H"*num_words, pattern)
                for word in words:
                    pattern_js += "%%u%s" % hex(word)

                pattern = pattern_js

            strPattern = idaapi.textctrl_info_t(text=pattern, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_SELECTED | Form.MultiLineTextControl.TXTF_READONLY)
            self.SetControlValue(self.strPattern, strPattern )


    def pattern_create(self, size):

        # Make sure at least one custom charset was specified
        if not len(self.pattern_complete):
            idaapi.warning("Cannot generate an empty pattern.")
            return False

        # Requested pattern size exceeds maximum allowed pattern size for the selected characterset
        elif size > len(self.pattern_complete):
            idaapi.warning("Requested size %d exceeds maximum unique pattern size %d." % (size, len(self.pattern_complete)))
            return False

        # Get a size substring from the complete pattern
        else:
            return self.pattern_complete[:size]

    def update_complete_pattern(self, c4, c3, c2, c1):

        c4_list = self.sploiter.c4_list if c4 else ["",]
        c3_list = self.sploiter.c3_list if c3 else ["",]
        c2_list = self.sploiter.c2_list if c2 else ["",]
        c1_list = self.sploiter.c1_list if c1 else ["",]

        # Generate complete pattern
        # NOTE: Not as pretty as calculating offset based on positions, but
        #       allows custom charactersets with varying positions
        self.pattern_complete = ""

        # Produce a cartesian product of all character lists and make a single string by flattening the list of lists
        self.pattern_complete = ''.join([i for j in itertools.product(c4_list, c3_list, c2_list, c1_list) for i in j])


###############################################################################
# Pattern Offset Form

class PatternDetectForm(Form):

    def __init__(self, sploiter, debugger = True):

        self.sploiter = sploiter
        self.debugger = debugger

        self.pattern_complete = ""

        form_dict = {
            'spacer'         : Form.StringLabel("           "),
            'strPattern'      : Form.StringInput(tp=Form.FT_ASCII, swidth=17),
            'strOffset'       : Form.StringInput(tp=Form.FT_ASCII, swidth=13),
            'iButtonDetect'   : Form.ButtonInput(self.OnButtonDetect, swidth=5),
            'rGroup'          : Form.RadGroupControl(('rFmtAuto', 'rFmtStr', 'rFmtAddr')),
            'cGroup'          : Form.ChkGroupControl(("cC4","cC3","cC2","cC1")),
            'strC4'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC3'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC2'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC1'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            }

        reg_form = ""
        self.reg_ctrls = list()

        # Display register values when executed with a debugger running
        if debugger:

            # Select general purpose registers for a given architecture
            if self.sploiter.addr64:
                self.reg_list = ["RAX", "RBX", "RCX", "RDX", "RSP", "RBP", "RSI", "RDI", "RIP", "R8", "R9", "R10", "R11", "R12", "R13", "R14", "R15"]
            else:
                self.reg_list = ["EAX", "EBX", "ECX", "EDX", "ESP", "EBP", "ESI", "EDI", "EIP"]

            # Populate register values and store form controls in a list
            for reg_name in self.reg_list:
                reg_value = idc.GetRegValue(reg_name)
                reg_value = self.sploiter.addr_format % reg_value

                ctrlRegValue = Form.StringInput(tp=Form.FT_ASCII, value=reg_value,swidth=17)
                form_dict["str%s" % reg_name] = ctrlRegValue

                ctrlRegOffset = Form.StringInput(tp=Form.FT_ASCII, swidth=13)
                form_dict["str%sOffset" % reg_name] = ctrlRegOffset

                reg_form += "  <%(reg)03s  :{str%(reg)s}><Offset:{str%(reg)sOffset}>\n" % {'reg': reg_name}

                self.reg_ctrls.append( (ctrlRegValue,  ctrlRegOffset) )

        Form.__init__(self, 
r"""BUTTON YES NONE
BUTTON CANCEL NONE
Detect pattern
{FormChangeCb}%s<Pattern:{strPattern}><Offset:{strOffset}>

Formats<Auto:{rFmtAuto}><Address:{rFmtAddr}><String:{rFmtStr}>{rGroup}>
Charset<C4:{cC4}><:{strC4}>
       <C3:{cC3}><:{strC3}>
       <C2:{cC2}><:{strC2}>
       <C1:{cC1}>{cGroup}><:{strC1}>

{spacer}<Detect:{iButtonDetect}>
""" % reg_form, form_dict)

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:

            # Set initial checkboxes for the classic metasploit 3-byte pattern
            self.SetControlValue(self.cC4, False)
            self.SetControlValue(self.cC3, True)
            self.SetControlValue(self.cC2, True)
            self.SetControlValue(self.cC1, True)

            # Set initial charsets values
            self.SetControlValue(self.strC4, self.sploiter.c4_list)
            self.SetControlValue(self.strC3, self.sploiter.c3_list)
            self.SetControlValue(self.strC2, self.sploiter.c2_list)
            self.SetControlValue(self.strC1, self.sploiter.c1_list)

            # NOTE: Trying to detect offset for registry values during
            #       form initialization appears to be very slow and does
            #       not provide an option for a user to specify a custom 
            #       characterset.

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

        return 1

    def OnButtonDetect(self, code=0):

        # Update global charset lists
        self.sploiter.c4_list = self.GetControlValue(self.strC4)
        self.sploiter.c3_list = self.GetControlValue(self.strC3)
        self.sploiter.c2_list = self.GetControlValue(self.strC2)
        self.sploiter.c1_list = self.GetControlValue(self.strC1)

        # Generate pattern only using selected lists
        self.update_complete_pattern(self.GetControlValue(self.cC4), self.GetControlValue(self.cC3), self.GetControlValue(self.cC2), self.GetControlValue(self.cC1))

        for ctrlRegValue, ctrlRegOffset in self.reg_ctrls:
            reg_value = self.GetControlValue(ctrlRegValue)
            self.SetControlValue(ctrlRegOffset, self.verbalize_pattern(reg_value, address=True) )

        pattern = self.GetControlValue(self.strPattern)

        # Force address format
        if self.GetControlValue(self.rFmtAddr):
            self.SetControlValue(self.strOffset, self.verbalize_pattern(pattern, address=True) )

        # Force string format
        elif self.GetControlValue(self.rFmtStr):
            self.SetControlValue(self.strOffset, self.verbalize_pattern(pattern, string=True) )

        # Auto format
        else:
            self.SetControlValue(self.strOffset, self.verbalize_pattern(pattern) )

    def verbalize_pattern(self, pattern, address=False, string=False):

        # Forced address format
        if address:

            # Cleanup the pattern
            pattern = pattern.replace(' ','')       # remove spaces
            pattern = pattern.replace('\\x','')     # remove '\x' prefixes
            pattern = pattern.replace('0x','')      # remove '0x' prefixes

            if self.sploiter.addr64 and len(pattern) == 8*2:
                try:
                    pattern = struct.pack("Q", int(pattern, 16))
                except Exception, e:
                    return "Invalid address"

            elif not self.sploiter.addr64 and len(pattern) == 4*2:
                try:
                    pattern = struct.pack("I", int(pattern, 16))
                except Exception, e:
                    return "Invalid address"

            else:
                return "Invalid address"

        # Forced string format
        elif string:
            pass

        # Autodetect format
        else:

            # Cleanup the pattern
            pattern = pattern.replace(' ','')       # remove spaces
            pattern = pattern.replace('\\x','')     # remove '\x' prefixes
            pattern = pattern.replace('0x','')      # remove '0x' prefixes

            # Attempt to detect and unpack pattern
            if self.sploiter.addr64 and len(pattern) == 8*2:
                try:
                    pattern = struct.pack("Q", int(pattern, 16))
                except Exception, e:
                    pass

            elif not self.sploiter.addr64 and len(pattern) == 4*2:
                try:
                    pattern = struct.pack("I", int(pattern, 16))
                except Exception, e:
                    pass

        if len(pattern) > 0:
            offset         = self.pattern_offset(pattern)
            offset_reverse = self.pattern_offset(pattern[::-1])
        else:
            return "Blank pattern"

        if offset >= 0:
            return "%d" % offset
        elif offset_reverse >= 0:
            return "%d (reverse)" % offset_reverse
        else:
            return "Not detected"

    def pattern_offset(self, pattern):

        offset = -1

        # Make sure at least one custom charset was specified
        if not len(self.pattern_complete):
            idaapi.warning("Cannot detect an empty pattern.")
            return offset

        # Locate the pattern
        if pattern in self.pattern_complete:
            offset = self.pattern_complete.index(pattern)

        return offset

    def update_complete_pattern(self, c4, c3, c2, c1):

        c4_list = self.sploiter.c4_list if c4 else ["",]
        c3_list = self.sploiter.c3_list if c3 else ["",]
        c2_list = self.sploiter.c2_list if c2 else ["",]
        c1_list = self.sploiter.c1_list if c1 else ["",]

        # Generate complete pattern
        # NOTE: Not as pretty as calculating offset based on positions, but
        #       allows custom character sets with varying positions
        self.pattern_complete = ""

        # Produce a cartesian product of all character lists and make a single string by flattening the list of lists
        self.pattern_complete = ''.join([i for j in itertools.product(c4_list, c3_list, c2_list, c1_list) for i in j])

###############################################################################
# Compare UI
###############################################################################

###############################################################################
# Compare Form

class CompareForm(Form):

    def __init__(self, sploiter, debugger = True):

        self.sploiter = sploiter
        self.debugger = debugger

        Form.__init__(self, 
r"""BUTTON YES NONE
BUTTON CANCEL NONE
Compare file to memory
{FormChangeCb}<File     :{impFile}>
<Memory   :{intAddr}>

<:{strCompare}>

<Bad Chars:{strBadChars}>
<Mem Holes:{strMemHoles}>

{spacer}<Compare:{iButtonCompare}>
""", {
            'spacer'        : Form.StringLabel("                                "),
            'impFile'        : Form.FileInput(swidth=62, open=True),
            'intAddr'        : Form.NumericInput(swidth=62, tp=Form.FT_ADDR, value=idaapi.get_screen_ea()),
            'strCompare'     : Form.MultiLineTextControl(text="", swidth=72, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_READONLY),
            'strBadChars'    : Form.StringInput(swidth=62,tp=Form.FT_ASCII),
            'strMemHoles'    : Form.StringInput(swidth=62,tp=Form.FT_ASCII),
            'iButtonCompare' : Form.ButtonInput(self.OnButtonCompare,),
            'FormChangeCb'   : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            pass

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

        return 1

    def OnButtonCompare(self, code=0):
        
        #######################################################################
        # Read file into buffer
        file_name = self.GetControlValue(self.impFile)

        # Verify file name was provided
        if not len(file_name):
            idaapi.warning("Please provide file name.")
            return

        # Attempt to read the file
        print "[idasploiter] Reading file %s" % file_name
        
        try:
            f = open(file_name,'rb')
        except Exception, e:
            idaapi.warning("File I/O error({0}): {1}.".format(e.errno, e.strerror) )
            return
        else:
            file_buf = f.read()
            f.close()

        # Check file content
        if not len(file_buf) > 0:
            idaapi.warning("Provided file was blank.")
            return

        # 64k ought to be enough for anybody =)
        elif len(file_buf) > 0xffff:
            idaapi.warning("File is too big.")
            return

        # Covert file bytes into printable hex view
        file_buf_str = " ".join(["%02x" % ord(b) for b in file_buf])
        file_buf_list = textwrap.wrap(file_buf_str, 48)

        #######################################################################
        # Read memory into buffer and mark non-matching bytes
        address = self.GetControlValue(self.intAddr)

        print "[idasploiter] Reading %d bytes from %s." % ( len(file_buf), self.sploiter.addr_format % address)

        if self.debugger:
            mem_buf = idaapi.dbg_read_memory(address, len(file_buf))
        else:
            mem_buf = idaapi.get_many_bytes(address, len(file_buf))

        if not mem_buf:
            idaapi.warning("Could not read memory.")
            return

        # Convert memory bytes into printable hex view
        mem_buf_str = ""

        # Badchars tracking
        badchars = list()
        badchars_nonseq = set()
        badchars_hole = list()

        # Locate all of the non-sequential badchars. The assumption is that
        # after the first badchar (e.g. 0x0A) the rest may be corrupted.
        for i, b in enumerate(file_buf):
            if file_buf[i] == mem_buf[i]:
                mem_buf_str+= ".. "

            else:
                mem_buf_str+= "%02x " % ord( mem_buf[i])

            # Do differential processing if the two bytes are not the same or
            # if the matching byte is in the bad characters list

            if file_buf[i] != mem_buf[i] or "%02x" % ord(file_buf[i]) in badchars_nonseq:

                ###############################################################
                # Memory holes

                # Check if the badchar appears anywhere else in memory to
                # detect positional holes. If this is the first discrepancy
                # then it is likely it is a badchar (e.g. 0x00, 0x0a, 0x0d)

                if len(badchars_hole):
                    hole_i, hole_size = badchars_hole[-1]

                    # Continuation of a previously discovered hole
                    if i == hole_i + hole_size:
                        badchars_hole[-1] = (hole_i, hole_size + 1)

                    # New hole discovered
                    else:
                        badchars_hole.append( (i,1) )

                # New hole discovered
                else:
                    badchars_hole.append( (i,1) )

                ###############################################################
                # Bad characters

                # Check if previous byte in the file is already a badchar, so
                # this new discrepancy doesn't matter:

                if not "%02x" % ord(file_buf[i-1]) in  badchars_nonseq:

                    # Skip sequential corruption
                    if len(badchars_nonseq):
                        bad_i, bad_b = badchars[-1]

                        # Check if the previous byte was also corrupted
                        # indicating a sequential corruption
                        if i - 1 != bad_i:
                            badchars_nonseq.add( "%02x" % ord(b) )

                    else:                      
                        badchars_nonseq.add( "%02x" % ord(b) )

                # Store badchars in a list
                badchars.append( (i,b) )


        mem_buf_list = textwrap.wrap(mem_buf_str, 48)

        if len(badchars) == 0:
            self.SetControlValue(self.strBadChars, "Identical.")  
        elif len(badchars) == len(file_buf):
            self.SetControlValue(self.strBadChars, "Completely different.")          
        else:
            self.SetControlValue(self.strBadChars, " ".join( sorted(badchars_nonseq) ) )
            self.SetControlValue(self.strMemHoles, ", ".join( [ "0x%02X: %d byte(s)" % h for h in badchars_hole ]  ) )

        #######################################################################
        # Populate the text box

        buf_formatted = "      00 01 02 03 04 05 06 07  08 09 0A 0B 0C 0D 0E 0F\n\n"
        for i in range( len(file_buf_list) ):

            file_line = file_buf_list[i].ljust(48)
            file_line = file_line[:23]+" "+file_line[23:]

            mem_line = mem_buf_list[i].ljust(48)
            mem_line = mem_line[:23]+" "+mem_line[23:]

            buf_formatted += "%04X  %s  F\n      %s  M\n\n" % (i*0x10, file_line, mem_line)

        strCompare = idaapi.textctrl_info_t(text=buf_formatted, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_READONLY)
        self.SetControlValue(self.strCompare, strCompare)

###############################################################################
# Plugin Manager
###############################################################################

class SploitManager():
    """ Class that manages GUI forms and exploitation methods of the plugin. """
    
    def __init__(self): 
        self.addmenu_item_ctxs = list()
        self.sploiter = Sploiter()

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

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

    def add_menu_items(self):
        if self.add_menu_item_helper("View/Open subviews/Segments", "Modules", "Shift+F6", 0, self.show_modules_view, None): return 1

        if self.add_menu_item_helper("Search/all error operands", "function pointers...", "Alt+f", 1, self.show_funcptr_view, None): return 1
        if self.add_menu_item_helper("Search/all error operands", "gadgets...", "Alt+r", 1, self.show_rop_view, None): return 1

        if self.add_menu_item_helper("Edit/Begin selection", "Create pattern...", "Shift+c", 0, self.show_pattern_create, None): return 1
        if self.add_menu_item_helper("Edit/Begin selection", "Detect pattern...", "Shift+d", 0, self.show_pattern_detect, None): return 1
        if self.add_menu_item_helper("Edit/Begin selection", "Compare file to memory...", "Shift+f", 0, self.show_compare, None): return 1
        
        return 0

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

    ###########################################################################
    # Utility Functions

    # Check debugger is ready for querying
    def check_debugger(self):
        if idaapi.dbg_can_query() and idaapi.get_process_state() < 0:
            return True
        else:
            return False

    ###########################################################################
    # View Callbacks
    
    # Modules View
    def show_modules_view(self):
        if self.check_debugger():
            self.sploiter.process_modules()
            self.sploiter.show_modules_view()
        else:            
            idaapi.warning("Please run and suspend the debugger.")
            

    # ROP View
    def show_rop_view(self):
        if self.check_debugger():
            self.sploiter.process_modules()
            self.sploiter.process_rop()
        else:            
            idaapi.warning("Please run and suspend the debugger.")
            

    # Function Pointer View
    def show_funcptr_view(self):     
        if self.check_debugger():   
            self.sploiter.process_modules()
            self.sploiter.process_funcptr()
        else:            
            idaapi.warning("Please run and suspend the debugger.")

    # Create Pattern Form
    def show_pattern_create(self):
        self.sploiter.pattern_create()

    # Detect Pattern Form
    def show_pattern_detect(self):

        # Debugger Version
        if self.check_debugger():
            self.sploiter.pattern_detect()

        # Static Version
        else:
            self.sploiter.pattern_detect(debugger=False)

    # Compare File to Memory Form
    def show_compare(self):

        # Debugger Version
        if self.check_debugger():
            self.sploiter.process_compare()

        # Static Version
        else:
            self.sploiter.process_compare(debugger=False)


###############################################################################

class idasploiter_t(plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = "Exploit development and vulnerability research toolkit."
    help = "Exploit development and vulnerability research toolkit."
    wanted_name = "IDA Sploiter"
    wanted_hotkey = ""

    def init(self):  

        # Only Intel x86/x86-64 are supported
        if idaapi.ph_get_id() == idaapi.PLFM_386:

            global idasploiter_manager

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

                idasploiter_manager = SploitManager()
                if idasploiter_manager.add_menu_items():
                    print "Failed to initialize IDA Sploiter."
                    idasploiter_manager.del_menu_items()
                    del idasploiter_manager
                    return idaapi.PLUGIN_SKIP
                else:
                    print("Initialized IDA Sploiter v%s (c) Peter Kacherginsky <iphelix@thesprawl.org>" % IDASPLOITER_VERSION)

            return idaapi.PLUGIN_KEEP
        else:
            return idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass

    def term(self):
        pass        

def PLUGIN_ENTRY():
    return idasploiter_t()

###############################################################################
# Script / Testing
###############################################################################

def idasploiter_main():
    global idasploiter_manager

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

    idasploiter_manager = SploitManager()
    idasploiter_manager.add_menu_items()

    sploiter = idasploiter_manager.sploiter

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


