<s> from collections import defaultdict <EOL> import idaapi <EOL> import sark <EOL> import idc <EOL> import sark . exceptions <EOL> def apply_enum_by_name ( enum , member_name ) : <EOL> member_value = enum . members [ member_name ] . value <EOL> for line in sark . lines ( * sark . get_selection ( ) ) : <EOL> for operand in line . insn . operands : <EOL> if operand . type . is_imm : <EOL> if operand . imm == member_value : <EOL> idc . OpEnumEx ( line . ea , operand . n , enum . eid , enum . members [ member_name ] . serial ) <EOL> elif operand . type . is_displ or operand . type . is_phrase : <EOL> if operand . addr == member_value : <EOL> idc . OpEnumEx ( line . ea , operand . n , enum . eid , enum . members [ member_name ] . serial ) <EOL> def get_common_value ( desired = None ) : <EOL> values = defaultdict ( int ) <EOL> for line in sark . lines ( * sark . get_selection ( ) ) : <EOL> for operand in line . insn . operands : <EOL> if operand . type . is_imm : <EOL> if desired is not None : <EOL> if desired == operand . imm : <EOL> return desired <EOL> values [ operand . imm ] += <NUM_LIT:1> <EOL> elif operand . type . is_displ or operand . type . is_phrase : <EOL> if desired is not None : <EOL> if desired == operand . addr : <EOL> return desired <EOL> values [ operand . addr ] += <NUM_LIT:1> <EOL> values [ <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> common_value = max ( values . iteritems ( ) , key = lambda x : x [ <NUM_LIT:1> ] ) [ <NUM_LIT:0> ] <EOL> return common_value <EOL> def const_name ( enum , value ) : <EOL> return "<STR_LIT>" . format ( enum . name . upper ( ) , value ) <EOL> def rename_immediate ( ) : <EOL> highlighted = idaapi . get_highlighted_identifier ( ) <EOL> try : <EOL> desired = int ( highlighted , <NUM_LIT:0> ) <EOL> except ( ValueError , TypeError ) : <EOL> desired = None <EOL> value = idc . AskLong ( get_common_value ( desired ) , "<STR_LIT>" ) <EOL> if value is None : <EOL> return <EOL> name = idc . AskStr ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> if name is None : <EOL> return <EOL> try : <EOL> enum = sark . Enum ( '<STR_LIT>' ) <EOL> except sark . exceptions . EnumNotFound : <EOL> enum = sark . add_enum ( '<STR_LIT>' ) <EOL> enum . members . add ( name , value ) <EOL> apply_enum_by_name ( enum , name ) <EOL> class RenameImmediateHandler ( idaapi . action_handler_t ) : <EOL> def activate ( self , ctx ) : <EOL> try : <EOL> rename_immediate ( ) <EOL> except : <EOL> import traceback <EOL> traceback . print_exc ( ) <EOL> return <NUM_LIT:1> <EOL> def update ( self , ctx ) : <EOL> if ctx . form_type == idaapi . BWN_DISASM : <EOL> return idaapi . AST_ENABLE_FOR_FORM <EOL> return idaapi . AST_DISABLE_FOR_FORM <EOL> class AutoEnumHandler ( idaapi . action_handler_t ) : <EOL> def activate ( self , ctx ) : <EOL> try : <EOL> rename_immediate ( ) <EOL> except : <EOL> import traceback <EOL> traceback . print_exc ( ) <EOL> return <NUM_LIT:1> <EOL> def update ( self , ctx ) : <EOL> if ctx . form_type == idaapi . BWN_DISASM : <EOL> return idaapi . AST_ENABLE_FOR_FORM <EOL> return idaapi . AST_DISABLE_FOR_FORM <EOL> class AutoEnum ( 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 . _last_enum = "<STR_LIT>" <EOL> self . rename_action_desc = idaapi . action_desc_t ( '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> RenameImmediateHandler ( ) , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> - <NUM_LIT:1> ) <EOL> idaapi . register_action ( self . rename_action_desc ) <EOL> self . autoenum_action_desc = idaapi . action_desc_t ( '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> AutoEnumHandler ( ) , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> - <NUM_LIT:1> ) <EOL> idaapi . register_action ( self . autoenum_action_desc ) <EOL> return idaapi . PLUGIN_KEEP <EOL> def term ( self ) : <EOL> idaapi . unregister_action ( self . autoenum_action_desc . name ) <EOL> idaapi . unregister_action ( self . rename_action_desc . name ) <EOL> def run ( self , arg ) : <EOL> pass <EOL> def autoenum ( self ) : <EOL> common_value = get_common_value ( ) <EOL> enum_name = idc . AskStr ( self . _last_enum , "<STR_LIT>" ) <EOL> if enum_name is None : <EOL> return <EOL> if not enum_name : <EOL> enum_name = None <EOL> self . _last_enum = enum_name <EOL> if common_value >> ( ( <NUM_LIT:8> * sark . core . get_native_size ( ) ) - <NUM_LIT:1> ) : <EOL> common_value = <NUM_LIT:0> <EOL> const_value = idc . AskLong ( common_value , "<STR_LIT>" ) <EOL> if const_value is None : <EOL> return <EOL> modify = True <EOL> try : <EOL> enum = sark . add_enum ( enum_name ) <EOL> except sark . exceptions . EnumAlreadyExists : <EOL> enum = sark . Enum ( enum_name ) <EOL> yes_no_cancel = idc . AskYN ( idaapi . ASKBTN_NO , <EOL> "<STR_LIT>" ) <EOL> if yes_no_cancel == idaapi . ASKBTN_CANCEL : <EOL> return <EOL> elif yes_no_cancel == idaapi . ASKBTN_YES : <EOL> modify = True <EOL> else : <EOL> modify = False <EOL> member_name = const_name ( enum , const_value ) <EOL> if modify : <EOL> try : <EOL> enum . members . add ( member_name , const_value ) <EOL> except sark . exceptions . SarkErrorAddEnumMemeberFailed as ex : <EOL> idaapi . msg ( "<STR_LIT>" . format ( ex . message ) ) <EOL> else : <EOL> for member in enum . members : <EOL> if member . value == const_value : <EOL> member_name = member . name <EOL> break <EOL> else : <EOL> return <EOL> apply_enum_by_name ( enum , member_name ) <EOL> def PLUGIN_ENTRY ( ) : <EOL> return AutoEnum ( ) </s>
