Skip to content

Instantly share code, notes, and snippets.

@clayne
Forked from mostobriv/case_chooser.py
Created April 30, 2022 17:45
Show Gist options
  • Save clayne/0265691e1c2b81bc9724592e98c9a889 to your computer and use it in GitHub Desktop.
Save clayne/0265691e1c2b81bc9724592e98c9a889 to your computer and use it in GitHub Desktop.
import idaapi
import ida_hexrays
import ida_lines
class MyChoose(idaapi.Choose):
def __init__(self, items, title, cols, icon=-1):
idaapi.Choose.__init__(self, title, cols, flags=idaapi.Choose.CH_MODAL, icon=icon)
self.items = items
def OnClose(self):
pass
def OnGetLine(self, n):
return self.items[n]
def OnGetSize(self):
return len(self.items)
def get_line_items(vdui, lnnum):
items = []
pc = vdui.cfunc.get_pseudocode()
if lnnum >= len(pc):
return items
line = pc[lnnum].line
tag = ida_lines.COLOR_ON + chr(ida_lines.COLOR_ADDR)
pos = line.find(tag)
while pos != -1 and len(line[pos+len(tag):]) >= ida_lines.COLOR_ADDR_SIZE:
addr = line[pos+len(tag):pos+len(tag)+ida_lines.COLOR_ADDR_SIZE]
idx = int(addr, 16)
a = ida_hexrays.ctree_anchor_t()
a.value = idx
if a.is_valid_anchor() and a.is_citem_anchor():
item = vdui.cfunc.treeitems.at(a.get_index())
if item:
items.append(item)
pos = line.find(tag, pos+len(tag)+ida_lines.COLOR_ADDR_SIZE)
return items
class GetCasesOfSwitch(idaapi.action_handler_t):
name = "GetCasesOfSwitch"
description = "Get cases of switch"
hotkey = ""
def __init__(self):
super().__init__()
def activate(self, ctx):
vdui = idaapi.get_widget_vdui(ctx.widget)
if vdui.item.citype == idaapi.VDI_EXPR:
if vdui.item.it.op == idaapi.cit_switch:
switch_insn = vdui.item.it.cinsn
else:
switch_insn = self.find_parent_switch_of(vdui.item.it, vdui.cfunc)
else:
items = get_line_items(vdui, vdui.cpos.lnnum)
if items[-1].op == idaapi.cit_switch:
switch_insn = items[-1].cinsn
else:
switch_insn = self.find_parent_switch_of(items[-1], vdui.cfunc)
if switch_insn is None:
return
switch = switch_insn.cswitch
concrete_cases, default_case = self.fetch_cases(switch)
chooser_cases = [['%#x' % address, 'case %d: // %#x' % (value, value)] for address, value in concrete_cases]
chooser_cases.append(['%#x' % default_case[0], "default:"])
case_chooser = MyChoose(chooser_cases,
"Cases of switch at %#x" % switch_insn.ea,
[["Address", 10 | MyChoose.CHCOL_HEX], ["Case value", 20]]
)
idx = case_chooser.Show(True)
if idx != -1:
idaapi.jumpto((concrete_cases + [default_case])[idx][0])
def fetch_cases(self, cswitch) -> list:
concrete_cases = list()
default_case = None
for case in cswitch.cases:
# concrete-value case
if len(case.values) > 0:
for value in case.values:
concrete_cases.append((case.ea, value))
# default case
else:
default_case = (case.ea, None)
return concrete_cases, default_case
def check(self, vdui) -> bool:
if vdui.item.citype == idaapi.VDI_EXPR and vdui.item.i.op == idaapi.cit_switch:
return True
items = get_line_items(vdui, vdui.cpos.lnnum)
if len(items) == 0:
return False
return self.find_parent_switch_of(items[-1], vdui.cfunc) is not None or items[-1].op == idaapi.cit_switch
def find_parent_switch_of(self, item, cfunc):
while item != cfunc.body:
item = cfunc.body.find_parent_of(item).cinsn
if item.op == idaapi.cit_switch:
return item
return None
def update(self, ctx) -> int:
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
return idaapi.AST_ENABLE_FOR_WIDGET
return idaapi.AST_DISABLE_FOR_WIDGET
_case_switch_action = GetCasesOfSwitch()
def _popup_hexrays_handler(event, *args):
if event == idaapi.hxe_populating_popup:
form, popup, hx_view = args
if _case_switch_action.check(hx_view):
idaapi.attach_action_to_popup(form, popup, _case_switch_action.name, None)
return 0
class CaseChooserPlugin_t(idaapi.plugin_t):
flags = idaapi.PLUGIN_HIDE
comment = "Navigation between cases of switch in pseudocode"
help = ""
wanted_name = "CaseChooserPlugin"
wanted_hotkey = ""
def init(self):
# Some initialization
global hexlight_cb_info, hexlight_cb
if idaapi.init_hexrays_plugin():
idaapi.register_action(
idaapi.action_desc_t(_case_switch_action.name, _case_switch_action.description, _case_switch_action, _case_switch_action.hotkey)
)
idaapi.install_hexrays_callback(_popup_hexrays_handler)
return idaapi.PLUGIN_KEEP
return idaapi.PLUGIN_SKIP
def run(self, arg=0):
return
def term(self):
return
def PLUGIN_ENTRY():
return CaseChooserPlugin_t()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment