# IDAxLM

[![GitHub Release](https://img.shields.io/github/release/harlamism/idaxlm?color=blue)](https://github.com/harlamism/idaxlm/releases)
[![GitHub Origin](https://img.shields.io/github/created-at/harlamism/idaxlm?color=blue)](https://github.com/harlamism/idaxlm/tree/4b92a29a9ece743b863230e424f7040fbd2c5459)
[![Last Commit](https://img.shields.io/github/last-commit/harlamism/idaxlm/master?color=blue)](https://github.com/harlamism/idaxlm/commits)

***IDAxLM*** is an ***IDA Pro*** plugin for engineering and consuming the large language model prompts.

**Key Features:**

- supplying context information alongside the function pseudocode
- processing functions in batch with a defined processing order
- extensible context and prompt template system
- making requests to multiple models simultaneously
- mapping the results to different parts of ***IDA*** database

![shot](./docs/main_dialog.png)  
*Figure 1. IDAxLM GUI*

> The plugin follows the same philosophy as [**IdaClu**](https://github.com/harlamism/IdaClu) by the same author. It does so by separating the plugin logic from user-generated content. The *UI* navigation model retains a familiar design. While ***IDAxLM*** and ***IdaClu*** complement each other, they stay distinct.

## The Purpose

If a reverse engineer lacks expertise in a given field, disassembled or even well reconstructed decompiled code may not add much to their understanding. This is where ***LLMs (Large Language Models)*** can help by identifying known algorithms, commonly used formulas, or at the very least, explaining custom-written code step by step.  

Given that ***LLM*** topic (including their use in *RE*) has been around for a while, the limitations of current application practices are becoming more apparent. At the time of this writing, some promising *PoC* solutions are beginning to emerge, though they address specific aspects and are not yet been organized into a systematic framework. **IDAxLM** is a step in that direction.

### Limitations

1. **Inconsistent Results**  
   *Analysts often interact with language models in an unstructured manner, leading to poorly reproducible results. Alternatively, they may be restricted to specific requests defined by certain plugins.*
2. **Lack of Context Information**  
   *Typically, only stripped pseudocode is sent for processing. There is no standardized method for including additional information such as reconstructed data types, observed input/output values, or similar function code text.*
3. **Filtering Generic Responses**  
   *Software, whether malware or benign, contains a significant amount of custom business logic. Without a specialized model, LLMs provide generic answers for most functions, requiring engineers to filter them out.*
4. **Chaotic Function Analysis**  
   *There is no standard method for selecting functions to send to LLMs based on attributes like prefix, folder, or color. The order is usually either random or based on manually choosing specific functions.*
5. **Limited LLM Provider Choice**  
   *Typically if the LLM plugin is used, you are limited to the list of providers supported by the developer. Adding a new provider requires studying plugin architecture and writing some code.*

### Solutions

1. **Prompt Templates**  
   *Community-driven database of building block sentences for creating consistent prompt.*
2. **Extracting Functions**  
   *Supporting functions that extract context data from **IDA** and insert it into prompt templates.*
3. **LLM Self-check**  
   *Evaluating if an LLM's answer is generic and flagging it accordingly*
4. **Function Analysis Order**  
   *Sending functions in an order that reflects bottom-up propagation in the call graph.*
5. **LLM Dialing Heuristics**  
   *Adding the provider's info to the config file allows the plugin's heuristic to likely call it correctly, making custom coding a last resort.*

Please note that **IDAxLM** is not a panacea. Despite our best efforts in providing context and crafting nearly perfect prompts, the results may sometimes be no better than simply copy-pasting a random stripped function into one of the major LLM chat windows.  
However, it is designed to organize and systematize information about what works well and what doesn’t, aiming to improve the LLM prompting workflow over time.

## Setup

### 1. Plugin Setup

#### IDAxLM Releases

- Hex-Rays Plugin Repository

  [**https://plugins.hex-rays.com/idaxlm**](https://plugins.hex-rays.com/idaxlm)

- GitHub Repository



### 2. Required Dependencies Installation

```shell
<python_path> -m pip install -r <plugin_path>/requirements.txt
```

If you're unsure about the ***Python*** version that ***IDA*** is using and need to find the path to ***Python***, type the following platform-neutral command in the input line at the bottom of ***IDA's*** screen (ensure ***IdaPython*** is working):

```shell
Python > python(os.__file__)
```

### 3. Config File Modification

Set the API tokens for the currently supported providers/models in the `./idaxlm/data/.env` config file:

- `API_TOKEN_OPENAI` / [**OpenAI**](https://playground.openai.com/) (gpt-3.5-turbo, gpt-4)
- `API_TOKEN_MISTRAL` / [**Mistral**](https://console.mistral.ai/) (mistral-large-2407)
- `API_TOKEN_GROQ` / [**Groq**](https://console.groq.com/playground) (llama3-70b-8192)


## Main Concepts

The plugin introduces 2 concepts - **templates** and **mappers**.

**Templates** - a set of static building block sentences for creating consistent final prompts to ***LMs***.  
There are 2 types of templates:

- Context Templates

  ```text
  The algorithm that is used in this function is {}.  
  This function's variables use the following type definitions: {}.  
  Here is/are an alternative implementation(s) of this function: {}
  ...
  ```

- Query Templates

  ```text
  Identify inefficient loops, redundant calculations, or unnecessary memory usage.
  Rewrite the following function in Python.
  Can this function be used in a malware context?
  ...
  ```

**Mappers** - a set of helper functions that **IDA** interface and **LM**-related content.  
There are 3 types of mappers:

- Pickers (***IDA*** interface $\longrightarrow$ extracted value / plugin user $\longrightarrow$ input value)
- Mutants (raw value $\longrightarrow$ value representation for ***LM***)
- Linkers (***LM*** completion $\longrightarrow$ ***IDA*** interface)

## User Interface

### Function Picker Menu

The context menu recognizes the selected function attributes such as prefix, folder, and color.  
Based on the selection, it can offer 1 to 5 picking options.

1. Pick same folder functions
2. Pick same prefix functions
3. Pick same color functions
4. Pick functions recursively
5. Pick selected functions

A function set chosen in a single pick is referred to as a task.

![shot](./docs/func_menu.png)  
*Figure 2. IDAxLM "Function Picker" context menu*

### Task Tool

The **Task Tool** reflects all picks made so far in the current database. The task header displays the function grouping principle, the originally picked function, and the total number of functions in this task. Expanding the header reveals the functions included in this task.

![shot](./docs/task_tool.png)  
*Figure 3. IDAxLM "Task Tool" widget*

A single click on a function switches the **current function context**.  
Double-clicking on a specific function shifts the focus of ***IDA*** to that function in the currently active view.  

Tasks and individual functions can be removed from the task using the `<Del>` key.  
They can also be disabled or re-enabled via the `Disable` option in the context menu.

### Prompt-Response Tabs

When a function is selected in the **Task Tool**, the controls on the **Prompt-Response Tabs** become enabled.  
There are 4 tabs with self-explanatory names that help manage communication with LLMs.

![shot](./docs/prompt_tabs.png)  
*Figure 4. IDAxLM "Prompt-Response Tabs"*

#### Context Tab

Selecting a template is mandatory if you want to provide context information about a function. Choosing a template will automatically fill in other fields, except for the **Handle**, which can be either filled by the corresponding **Picker** mapper or entered manually by the user.

![shot](./docs/context_tab.png)  
*Figure 5. IDAxLM Context Tab*

It is advised to follow these syntax rules for the **Handle** field:

- raw text is treated as a single value
  - `Exploit64`
  - `LU decomposition algorithm`
- colons `<:>` are used to separate keys and values
  - `v5:{{1,2,3},{4,5,6},{7,8,9}}`
  - `a1:PDEVICE_OBJECT`
- commas `<,>` are used to separate values of the same type
  - `0xFE,0xED,0xFA,0xCF`
  - `C:\\lab\\idb\\sample.i64:0x004522A0,0x00452490`
- semicolons `<;>` are used to denote consecutive key-value pairs
  - `a1:Mx3;a2:Mx3;a3:Vec3`
  - `C:\\idb\\7820b663.i64:0x00403A80;C:\\idb\\4bf11f56.i64:0x00403710,0x00403850;`

#### Prompt Tab

As with the **Context Tab**, selecting a template name on **Prompt Tab** will automatically fill in other fields. The template is divided into four typical prompting parts, each affecting how the resulting response will appear.

- **Text** - the question to ask or the action to be performed
- **Type** - the text format in which the answer should be provided *(optional)*
- **Case** - example demonstrating how the ideal prompt-response should look *(optional)*
- **Safe** - instructions on how to ensure consistent responses *(optional)*

![shot](./docs/prompt_tab.png)  
*Figure 6. IDAxLM Prompt Tab*

If the template does not meet your needs, you can either modify the existing one or create new by selecting the **\<New Template\>** entry in the template selection menu. The templates are assumed to be well-configured, so the small input fields are intended for lookup purposes or minor edits.

#### Completion Tab

The **Completion Tab** displays the answers to prompts selected in the Prompt tab for the current function. It allows you to take action with the gathered information.
The actions that can be performed with specific data are defined by the **Linker** list referenced in the prompt settings *(see "Extensions" section)*.

![shot](./docs/response_mapping.png)  
*Figure 7. IDAxLM Completion Tab*

#### Chat Tab

The **Chat Tab** provides a familiar chat interface with the added feature of receiving responses from multiple LLMs, allowing users to compare results. The expectation is that in a particularly difficult situation, one of the models might offer a better solution.

![shot](./docs/live_chat.png)  
*Figure 8. IDAxLM Chat Tab*

An icon in the top right of each response allows you to copy the response data as a comment into the pseudocode of the function being discussed. If the code language was specified, fenced Markdown code snippets are syntax-highlighted within chat bubbles.

**Fragments** can be excluded from the conversation context, allowing them to remain visible without being sent repeatedly in subsequent requests. To toggle this setting, right-click on the empty space to the left or right of the conversation bubble and select either `Enable` or `Disable`.

For brevity and reusability, special **markers** are available to add related function and type definitions:

- functions: `<func:func_ea/func_name/screen>`
- types: `<type:type_name>`

### Provider Tool

#### Provider Picker

**Provider Picker**, as the name suggests, allows you to select **services** to be polled. A **service** here refers to the combination of the provider and the model. The service list is populated from an internal configuration file. At least one service must be selected to make any further queue or live requests.

![shot](./docs/provider_tool.png)  
*Figure 9. IDAxLM Service Picker*

#### Pricing Calc

**Pricing Calc** is a calculator that estimates the overall input token count and price if all tasks are executed. It becomes active only if paid language models are listed. To show it, click on the **Provider Tool** header.

![shot](./docs/pricing_calc.png)  
*Figure 10. IDAxLM Pricing Calculator*

## Data Flow

### Mappers

While **templates** store the static structure, **mappers** handle the interactive part. As mentioned earlier, the former reference the latter. The **context templates** and **mutant mappers** serve as good examples of how templates and mappers work together with data.

#### Example #1: Context Template / Mutant Mapper

Consider **types** template that was selected on the **Context Tab**. This template is used to extract data type descriptions and map them to the specified variables. Here is full text description of it at the moment of submission:

```json
{
  "template": "types",
  "picker": "get_idle",
  "handle": "v1:Vec3;v2:Vec3",
  "mutant": "fix_types",
  "prompt": "This function's variables use the following type definitions: {}."
}
```

Here, the `get_idle` picker mapper instructs the plugin to expect custom values from the user. These are `v1:Vec3;v2:Vec3`, which translates to "variables names `v1` and `v2` are of the type `Vec3`".

When user clicks the `Submit` button another mutant mapper `fix_types` comes into play. It processes the given input and inserts the results into the prompt string.

The overall result of applying this template appears as follows:

```text
This function's variables use the following type definitions:
The variables `v1, v2` are of type `Vec3`.
The definition of the type `Vec3` is as follows:

    ```
    struct Vec3
    {
        float x;
        float y;
        float z;
    };
    ```
```

#### Example #2: Context Template / Mutant Mapper

Let’s turn our attention to the **variants** context template. This template retrieves function text either from the current ***IDA*** database or from another database if a path is provided. Below is the text description of the template at the time of submission:

```json
{
  "template": "variants",
  "picker": "get_idle",
  "handle": "C:\\lab\\idb\\sample.i64:0x004522A0",
  "mutant": "fix_funcs",
  "prompt": "Here is/are an alternative implementation(s) of this function:\n{}"
}

```

The `get_idle` picker allowed the user to input custom values. Since a path to the ***IDA*** ***i64*** database was provided, the `fix_funcs` mapper will run another instance of ***IDA*** in headless mode to fetch the function text at a given address.

The prompt outcome will be:

```text
Here is/are an alternative implementation(s) of this function:
    ```cpp
    double __cdecl f1d(float a1)
    {
      _BYTE *v1; // ecx
      <redacted for brevity>
      v1 = &v7[4 * dword_4595A8];
      if ( v7 < v1 )
      {
        v2 = 0;
        do
        {
          v2 += 4;
          result = a1 * *(v2 + 4560276);
          *&v7[v2 - 4] = result;
        }
        while ( &v7[v2] < v1 );
      }
      <redacted for brevity>
      return result;
    }
    ```
```

### Workflow

A simple illustrative workflow example is provided with code but without comments, assuming the reader is now familiar with how templates and mappers function.

It is assumed that a function is selected in the **Tasks** widget, and at least one provider is chosen in the **Provider** tool.

#### Sequence #1: switch to "Context" tab > select "xrefs" template > click "Submit" button

1. `./data/context.json/xrefs`

    ```json
    "xrefs": {
      "picker": "get_xrefs",
      "handle": "text",
      "default": "",
      "mutant": "fix_count",
      "prompt": "This function is called by {} function(s)."
    }
    ```

2. `./data/pickers.py/get_xrefs()`

    ```python
    def get_xrefs(ea, env=None):
      xref_count = len(list(idautils.XrefsTo(ea)))
      return xref_count
    ```

3. `./data/mutant.py/fix_count()`

    ```python
    def fix_count(val, env=None):
        return 'many' if int(val) > 1 else 'single'
    ```

4. `assembled prompt part`

   ```text
   This function is called by many function(s).
   ```

#### Sequence #2: switch to "Prompt" tab > select "func_name" template

1. `./data/prompt.json/func_name()`

    ```json
    "func_name": {
        "text": "Suggest new better function name.",
        "type": "Reply with a regular text.",
        "case": "",
        "safe": "Do not explain anything, only print the function name.",
        "cond": "Add the `gen_` prefix to this name if the formula or algorithm \
                "used in this function is not known and doesn't have a specific name.",
        "link": ["set_func_comment", "set_func_name"]
    }
    ```

#### Sequence #3: click "Submit" button -> switch to "Completion" tab -> select "Set Function Name" from menu

1. `./data/linker.py/set_func_name()`

    ```python
    def set_func_name(loc, val, org, env=None):
        idc.set_name(loc, val, idaapi.SN_CHECK)
        update_ui(loc)
    ```

## Compatibility

**Recommended specs:** *IDA Pro v8.3+* with *Python v3.x*.  
**Minimum specs:** *IDA Pro v7.7+* with *Python v3.x+*.  

The scripts bundled with **IDAxLM** are currently tailored to work with ***Intel x86/x64 > PE32/PE32+ > C/C++*** binaries.  
To ensure cross-compatibility **IDAxLM** relies on ***PyQt*** and ***IdaPython*** shims.

### Test environment

  1. *IDA Pro v8.4, Windows 11*
  2. *IDA Pro v7.6, macOS Sonoma v14.6.1 (Intel)*
  3. *IDA Pro v7.0, macOS High Sierra v10.13.6*
  4. *IDA Pro v9.0 Beta, Windows 11/macOS Sonoma v14.6.1 (Intel)/Linux Debian v12.7*

Other ***IDA*** versions for *Windows* can be added as test environments upon request.  

## Upcoming Features

- Saving/loading user sessions;
- Indicator in **Task Tool** to show if context, prompt, completion, or chat was set for the function;
- Recognition of whether the examined binary was written in a programming language other than *C/C++*;
- Fine-tuning of the language model from the plugin interface;
- Hierarchical types as context elements;
- Same class methods as context elements;
- Importing context data from other sources into ***IDA***;
- Additional tooltips explaining the use and purpose of *UI* elements;
- Conditions in prompt templates;
- Parallel requests where/when applicable;
- Code optimizations and data caching.

## Known Issues

- Token count and estimated price calculations for *LM* requests are significantly inaccurate;
- Function names don't update in plugin views if changed externally;
- Function processing order in bottom-up propagation is sometimes incorrect;
- No support for environment variables in `.env` file starting with `PLUGIN_` prefix;
- Screen space is critical for both ***IDA*** and the plugin, so control choice/placement are subjects of discussion;
- Minor *UI* inconsistency on *macOS*;
- Cross-compatibility issues with older ***IDA*** versions around ***v7.0***.

## Acknowledgements

[***Tim 'mr.phrazer' Blazytko, Moritz Schloegel***](https://mschloegel.me/slides/slides_recon24_llm_reverse_engineering.pdf) - for putting a bit more weight on my inferences with your reasoning.  
[***Guys from Atredis Partners***](https://www.atredis.com/blog/2024/6/3/how-to-train-your-large-language-model) - for demonstrating the workflow and usage patterns when working with language models locally.

## Version History

- 2024-10-08 - Release of more stable *v1.0* *(UPCOMING)*
- 2024-12-30 - The original release of IDAxLM v0.9!
