Add programmable undo support to IDAPython

ida_undo.py

"""
IDAPython Extension library (c) Elias Bachaalany.

Undo utilities
"""

import idaapi

# ---------------------------------------------------------------
class undo_handler_t(idaapi.action_handler_t):
    """Helper internal class to execute the undo-able user function"""
    id = 0
    def __init__(self, callable, *args, **kwargs):
        idaapi.action_handler_t.__init__(self)
        self.id += 1
        self.callable   = callable
        self.args       = args
        self.kwargs     = kwargs
        self.result     = None


    def activate(self, ctx):
        self.result = self.callable(*self.args, **self.kwargs)
        return 0


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



# ---------------------------------------------------------------
class undoable_t:
    """Callable class that invokes the user's function via
    process_ui_actions(). This will create an undo point and
    hence making the function 'undoable'
    """
    def __init__(self, callable):
        self.callable = callable

    
    def __call__(self, *args, **kwargs):
        ah   = undo_handler_t(self.callable, *args, **kwargs)
        desc = idaapi.action_desc_t(
                    f"ida_undo_{self.callable.__name__}_{ah.id}",
                    f"IDAPython: {self.callable.__name__}",
                    ah)

        if not idaapi.register_action(desc):
            raise(f'Failed to register action {desc.name}')

        idaapi.process_ui_action(desc.name)
        idaapi.unregister_action(desc.name)

        return ah.result


    @staticmethod
    def undo():
        idaapi.process_ui_action('Undo')


undoable = lambda callable: undoable_t(callable)

Example use

from ida_undo import undoable

@undoable
def my_del_function(ea, banner):
    idaapi.del_func(ea)
    print(f'{ea:x}: {banner}')
    return True


def main():
    idaapi.msg_clear()
    
    if my_del_function(here(), 'script: undoable delete function'):
        warning('about to undo!')
        my_del_function.undo()


main()