IDA Pro Plugin - know what your function is doing by ollama LLM

In this document, I tried to develop a simple IDA Pro python plugin which is able to interpret what a decompiled function is doing by using LLM (ollama)

Experiment Environment:

1. Windows 10
2. IDA Pro 7.2
3. idapython
4. ollama environment (Please use these two files!!!)
   ollama model: https://huggingface.co/AverageBusinessUser/aidapal/blob/main/aidapal-8k.Q4_K_M.gguf
   ollama model file: https://huggingface.co/AverageBusinessUser/aidapal/resolve/main/aidapal.modelfile

.sh

ollama create aidapal -f aidapal.modelfile

Below is the plugin script that fetches IDA-Pro decompiled pseudo code and sending to remote server where ollama is deployed.

analyze.py

import ida_kernwin,ida_funcs,idaapi,idautils,idc,ida_hexrays,ida_name
import threading,requests,json,zlib
from functools import partial

# Define callbacks which are the actions I actually
# want to perform

# this list holds the list of models registered with Ollama to be accessible from the plugin.
models = ['aidapal']
# where ollama service is running
ollama_url = "http://<IP Adress>:11434/api/generate"

# this function sends request to ollama API and get response as json
def do_analysis(code,model_name):
    url = ollama_url
    headers = {"Content-Type": "application/json"}
    payload = {"model": model_name, "prompt": code, "stream": False,"format":"json"}
    res = requests.post(url, headers=headers, json=payload)
    try:
        t = res.json()['response']
        t = json.loads(t)
        return t
    except:
        # rarely this occurs, leftover from early on
        print('aidapal: error unpacking response')
        print(res.json()['response'])

def hello_mars():

    print('Hello Mars!')

def hello_earth():

    # Create an instance and install
    cur_func = ida_hexrays.decompile(idaapi.get_screen_ea())
    print "[*] Get Function:"
    print str(cur_func)

    res = do_analysis(str(cur_func), 'aidapal')
    print "[*] Get Analyzed Result:"
    print res['comment']


# Define the action_handler_t object that fires the
# callback function when each action is activated

class ActionHandler(idaapi.action_handler_t):

    def __init__(self, callback):
        
        idaapi.action_handler_t.__init__(self)
        self.callback = callback
    
    def activate(self, ctx):

        self.callback()
        return 1

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


# Define a method to register all the actions when
# the plugin is initialized

def register_actions():   

    actions = [
        {
            'id': 'hello:earth',
            'name': 'Know What is doing',
            'hotkey': 'Ctrl+Alt+E',
            'comment': 'Know What is doing',
            'callback': hello_earth,
            'menu_location': 'Edit/Ollma/Know What is doing'
        },
    ]


    for action in actions:

        if not idaapi.register_action(idaapi.action_desc_t(
            action['id'], # Must be the unique item
            action['name'], # The name the user sees
            ActionHandler(action['callback']), # The function to call
            action['hotkey'], # A shortcut, if any (optional)
            action['comment'] # A comment, if any (optional)
        )):

            print('Failed to register ' + action['id'])

        if not idaapi.attach_action_to_menu(
            action['menu_location'], # The menu location
            action['id'], # The unique function ID
            0):

            print('Failed to attach to menu '+ action['id'])


# Define the plugin class itself which is returned by
# the PLUGIN_ENTRY method that scriptable plugins use
# to be recognized within IDA

class FunctionDecompilationAnalysis(idaapi.plugin_t):

    # Use the HIDE flag to avoid the entry in 
    # Edit/Plugins since this plugin's run() 
    # method has no functionality...it's all 
    # in the actions.

    flags = idaapi.PLUGIN_HIDE
    comment = 'A plugin to call ollama'
    help = 'No help - this is just a test'
    wanted_name = 'Know what this function is doing'
    wanted_hotkey = ''

    def init(self):
        print('FunctionDecompilationAnalysis init')
        
        register_actions()

        # Return KEEP instead of OK to keep the
        # plugin loaded since it registers
        # callback actions and hotkeys
        #
        # Use OK if the functionality is 
        # all in the plugin and it does
        # one thing then completes.
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        print('FunctionDecompilationAnalysis run')

    def term(self):
        print('HelloWorldPlugin term')


# The PLUGIN_ENTRY method is what IDA calls when
# scriptable plugins are loaded. It needs to
# return a plugin of type idaapi.plugin_t

def PLUGIN_ENTRY():
    
    try:
        return FunctionDecompilationAnalysis()
    
    except Exception, err:
        import traceback
        print('Error: %s\n%s' % str((err), traceback.format_exc()))
        raise

Copy analyze.py into %IDAPRO%/plugins/ and restart IDA Pro to launch the plugin

Click "Ollma" -> "Know What is doing"

Get Result from console