respond to button clicks in IDA Pro

clicker.py

import re
import collections

import idaapi
import ida_kernwin


class button_hooks_t(ida_kernwin.View_Hooks):
    def __init__(self, v):
        '''
        respond to "button" clicks for the given widget.

        example:

            h = button_hooks_t(ida_kernwin.find_widget("IDA View-A"))
            h.add_button(0x401050, "click me!", lambda ea, x, y: print("clicked!"))
            
        renders to:
            .text:00401050
            .text:00401050 push    ebp             ; ˂click me!˃
            .text:00401051 mov     ebp, esp
        '''
        ida_kernwin.View_Hooks.__init__(self)
        self.hook()
        self.v = v
        self.callbacks = collections.defaultdict(lambda: list())

    def view_click(self, view, event):
        # only respond to click events for the view we're targetting.
        if view != self.v:
            return

        # only care about addresses that we're tracking.
        ea = idaapi.get_screen_ea()
        if ea not in self.callbacks:
            return

        for cb in self.callbacks[ea]:
            cb(ea, event.x, event.y)

    # note, these are not < and > characters.
    # these are actually unicode lookalikes.
    # this makes it less likely that someone accidentally types these in as comments.
    BUTTON_START_SYMBOL = "˂"
    BUTTON_END_SYMBOL = "˃"
    CLEAR_RE = re.compile(r'˂[^˃]*˃ ')

    def Button(tag, callback):
        def _handle(ea, x, y):
            # like:
            #
            #     .text:00401050     push    ebp    ; <foo>

            line = idaapi.tag_remove(ida_kernwin.get_curline())

            # col: the index of the position within the line.
            # row: the index of the row, within the current view. *not* global row.
            _, col, row = ida_kernwin.get_cursor()

            # everything after the cursor, until the end of button
            after = line[col:].partition(button_hooks_t.BUTTON_END_SYMBOL)[0]

            # everything from start of button until the cursor
            before = line[:col].rpartition(button_hooks_t.BUTTON_START_SYMBOL)[2]

            # this should be the tag contents (or junk)
            found_tag = before + after
            if tag == found_tag:
                callback(ea, x, y)

        return _handle

    def add_button(self, ea, name, callback):
        '''
        add a "button" as a comment at the given address. 
        when the button is clicked by the user, invoke the given callback.
        the button looks like `<name>`; however, those are not < and > characters - they're unicode lookalikes.
        '''
        tag = button_hooks_t.BUTTON_START_SYMBOL + name + button_hooks_t.BUTTON_END_SYMBOL

        existing_cmt = idaapi.get_cmt(ea, False) or ''
        if existing_cmt:
            existing_cmt += ' '

        if tag not in existing_cmt:
            idaapi.set_cmt(ea, existing_cmt + tag, False)
        self.callbacks[ea].append(button_hooks_t.Button(name, callback))

    def clear_buttons(self):
        '''
        remove the "buttons" from comments written into the IDB.
        '''
        for ea in self.callbacks.keys():
            idaapi.set_cmt(ea, button_hooks_t.CLEAR_RE.sub(idaapi.get_cmt(ea, False), ''), False)
        self.callbacks = collections.defaultdict(lambda: list())