<s> from awesome . context import ignored <EOL> import idaapi <EOL> import sark <EOL> from ida_settings import IDASettings <EOL> HIGHLIGHT_COLOR = <NUM_LIT> <EOL> def highlight_calls_in_function ( ea ) : <EOL> highlighted_lines = set ( ) <EOL> for line in sark . Function ( ea ) . lines : <EOL> if not line . insn . is_call : <EOL> continue <EOL> if line . color is None : <EOL> line . color = HIGHLIGHT_COLOR <EOL> highlighted_lines . add ( line . ea ) <EOL> return highlighted_lines <EOL> class UiHooks ( idaapi . UI_Hooks ) : <EOL> def __init__ ( self , lines ) : <EOL> super ( UiHooks , self ) . __init__ ( ) <EOL> self . lines = lines <EOL> def updating_actions ( self , ctx ) : <EOL> if ctx . form_type == idaapi . BWN_DISASM : <EOL> with ignored ( sark . exceptions . SarkNoFunction ) : <EOL> self . lines . update ( highlight_calls_in_function ( ctx . cur_ea ) ) <EOL> return super ( UiHooks , self ) . updating_actions ( ctx ) <EOL> class ToggleHandler ( idaapi . action_handler_t ) : <EOL> def __init__ ( self , state = False ) : <EOL> idaapi . action_handler_t . __init__ ( self ) <EOL> self . _state = state <EOL> def on_enable ( self , ctx ) : <EOL> raise NotImplementedError ( ) <EOL> def on_disable ( self , ctx ) : <EOL> raise NotImplementedError ( ) <EOL> def activate ( self , ctx = idaapi . action_activation_ctx_t ( ) ) : <EOL> if self . _state : <EOL> self . on_disable ( ctx ) <EOL> else : <EOL> self . on_enable ( ctx ) <EOL> self . _state = not self . _state <EOL> class ToggleHighlightHandler ( ToggleHandler ) : <EOL> def __init__ ( self , enable , disable ) : <EOL> ToggleHandler . __init__ ( self ) <EOL> self . enable = enable <EOL> self . disable = disable <EOL> def on_enable ( self , ctx ) : <EOL> self . enable ( ) <EOL> def on_disable ( self , ctx ) : <EOL> self . disable ( ) <EOL> def update ( self , ctx ) : <EOL> return idaapi . AST_ENABLE_ALWAYS <EOL> class SelectColorHandler ( idaapi . action_handler_t ) : <EOL> def __init__ ( self , set_color ) : <EOL> idaapi . action_handler_t . __init__ ( self ) <EOL> self . _set_color = set_color <EOL> def activate ( self , ctx ) : <EOL> color = sark . ui . ask_color ( initial = HIGHLIGHT_COLOR ) <EOL> if color is not None : <EOL> self . _set_color ( color ) <EOL> return <NUM_LIT:1> <EOL> def update ( self , ctx ) : <EOL> return idaapi . AST_ENABLE_ALWAYS <EOL> class HighlightCalls ( idaapi . plugin_t ) : <EOL> flags = idaapi . PLUGIN_PROC <EOL> comment = '<STR_LIT>' <EOL> help = '<STR_LIT>' <EOL> wanted_name = '<STR_LIT>' <EOL> wanted_hotkey = '<STR_LIT>' <EOL> def init ( self ) : <EOL> self . lines = set ( ) <EOL> self . settings = IDASettings ( '<STR_LIT>' ) <EOL> try : <EOL> self . set_color ( self . settings [ '<STR_LIT>' ] ) <EOL> except KeyError : <EOL> self . settings . user [ '<STR_LIT>' ] = HIGHLIGHT_COLOR <EOL> self . set_color ( HIGHLIGHT_COLOR ) <EOL> self . ui_hooks = UiHooks ( self . lines ) <EOL> self . toggle_action_desc = idaapi . action_desc_t ( '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> ToggleHighlightHandler ( self . enable_highlights , <EOL> self . disable_highlights ) , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> - <NUM_LIT:1> ) <EOL> idaapi . register_action ( self . toggle_action_desc ) <EOL> self . color_selector = idaapi . action_desc_t ( '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> SelectColorHandler ( set_color = self . set_color ) , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> - <NUM_LIT:1> ) <EOL> idaapi . register_action ( self . color_selector ) <EOL> idaapi . attach_action_to_menu ( '<STR_LIT>' , self . toggle_action_desc . name , idaapi . SETMENU_APP ) <EOL> idaapi . attach_action_to_menu ( '<STR_LIT>' , self . color_selector . name , idaapi . SETMENU_APP ) <EOL> return idaapi . PLUGIN_KEEP <EOL> def set_color ( self , color ) : <EOL> self . color = color <EOL> self . settings . user [ '<STR_LIT>' ] = color <EOL> global HIGHLIGHT_COLOR <EOL> HIGHLIGHT_COLOR = color <EOL> self . clear_all_highlights ( ) <EOL> def enable_highlights ( self ) : <EOL> self . ui_hooks . hook ( ) <EOL> def term ( self ) : <EOL> self . disable_highlights ( ) <EOL> idaapi . unregister_action ( self . toggle_action_desc . name ) <EOL> idaapi . unregister_action ( self . color_selector . name ) <EOL> def disable_highlights ( self ) : <EOL> self . ui_hooks . unhook ( ) <EOL> self . clear_all_highlights ( ) <EOL> def clear_all_highlights ( self ) : <EOL> for ea in self . lines : <EOL> sark . Line ( ea = ea ) . color = None <EOL> self . lines . clear ( ) <EOL> def run ( self , arg ) : <EOL> pass <EOL> def PLUGIN_ENTRY ( ) : <EOL> return HighlightCalls ( ) </s>
