# Directory Structure
```
├── .env_example
├── .gitignore
├── images
│ ├── pic.png
│ └── pic2.png
├── pyproject.toml
├── README.md
├── src
│ └── headless_ida_mcp_server
│ ├── __init__.py
│ ├── __main__.py
│ ├── helper.py
│ ├── logger.py
│ └── server.py
└── test
├── heap
│ ├── main
│ └── main.c
├── stack
│ ├── main
│ └── main.c
└── test_ida_agent.py
```
# Files
--------------------------------------------------------------------------------
/.env_example:
--------------------------------------------------------------------------------
```
IDA_PATH=/home/ubuntu/idapro-9.0/idat
PORT=8888
HOST=0.0.0.0
TRANSPORT=sse
# TRANSPORT=stdio
```
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
```
# Python-generated files
__pycache__/
*.py[oc]
build/
dist/
wheels/
*.egg-info
# Virtual environments
.venv/
.vscode/
.python-version
.cursor/
.env
exploit.py
```
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
```markdown
# Acknowledgments
This project builds upon the work of:
- Tools code adapted from [ida-pro-mcp](https://github.com/mrexodia/ida-pro-mcp) by mrexodia
- Utilizes the [headless-ida](https://github.com/DennyDai/headless-ida) library by DennyDai
# Headless IDA MCP Server
If you want to run the server directly as a cli app, rather than an IDA plugin interactively,you can chose it.
## Project Description
This project uses IDA Pro's headless mode to analyze binary files and provides a suite of tools via MCP to manage and manipulate functions, variables, and more.
## Prerequisites
- Python 3.12 or higher
- IDA Pro with headless support (idat) https://github.com/DennyDai/headless-ida
## Installation
1. Clone the project locally:
```bash
git clone https://github.com/cnitlrt/headless-ida-mcp-server.git
cd headless-ida-mcp-server
```
2. Install dependencies:
```bash
uv python install 3.12
uv venv --python 3.12
uv pip install -e .
```
## Configuration
1. Copy the example environment file:
```bash
cp .env_example .env
```
2. Configure the following environment variables in `.env`:
- `IDA_PATH`: Path to IDA Pro's headless executable (idat), e.g., `/home/ubuntu/idapro/idat`
- `PORT`: Port number for the MCP server, e.g., `8888`
- `HOST`: Host address for the MCP server, e.g., `127.0.0.1`
- `TRANSPORT`: MCP transport mode (`sse` or `stdio`)
## Usage
1. Start the server:
```bash
uv run headless_ida_mcp_server
```
2. Connect to the server using an MCP client:
Debug it:
```bash
npx -y @modelcontextprotocol/inspector
```
or
```json
{
"mcpServers": {
"ida": {
"command": "/path/to/uv",
"args": ["--directory","path/to/headless-ida-mcp-server","run","headless_ida_mcp_server"]
}
}
}
```
```
--------------------------------------------------------------------------------
/src/headless_ida_mcp_server/__main__.py:
--------------------------------------------------------------------------------
```python
from .server import main
if __name__ == "__main__":
main()
```
--------------------------------------------------------------------------------
/test/stack/main.c:
--------------------------------------------------------------------------------
```cpp
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
// 栈溢出漏洞,english
void vuln() {
char buf[100];
printf("Please enter a string: ");
read(0,buf,0x200);
}
int main() {
setvbuf(stdout, 0,2,0);
setvbuf(stdin, 0,2,0);
setvbuf(stderr, 0,2,0);
vuln();
return 0;
}
```
--------------------------------------------------------------------------------
/src/headless_ida_mcp_server/logger.py:
--------------------------------------------------------------------------------
```python
import logging
#### LOGGING ####
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger('headless_ida_mcp_server')
logger.setLevel(logging.DEBUG)
logging.getLogger("mcp.server.lowlevel.server").setLevel(logging.ERROR)
```
--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
```toml
[project]
name = "headless-ida-mcp-server"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"dotenv>=0.9.9",
"fastmcp>=0.4.1",
"headless-ida>=0.6.1",
"mcp>=1.6.0",
"pytest>=8.3.5",
"pytest-asyncio>=0.26.0",
]
[project.scripts]
headless_ida_mcp_server = "headless_ida_mcp_server.server:main"
```
--------------------------------------------------------------------------------
/src/headless_ida_mcp_server/__init__.py:
--------------------------------------------------------------------------------
```python
"""
headless-ida-mcp-server
"""
from dotenv import load_dotenv,find_dotenv,dotenv_values
import os
from .logger import*
__all__ = ['PORT', 'HOST', 'TRANSPORT', 'BINARY_PATH', 'IDA_PATH']
load_dotenv(find_dotenv(),override=True)
PORT = os.environ.get("PORT", 8888)
HOST = os.environ.get("HOST", "0.0.0.0")
TRANSPORT = os.environ.get("TRANSPORT", "sse")
IDA_PATH = os.environ.get("IDA_PATH", "")
if IDA_PATH == "":
raise ValueError("IDA_PATH is not set")
```
--------------------------------------------------------------------------------
/test/test_ida_agent.py:
--------------------------------------------------------------------------------
```python
import pytest
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
async def test_ida_function_query():
model = ChatOpenAI(model="gpt-4")
async with MultiServerMCPClient(
{
"ida": {
"url": "http://localhost:8888/sse",
"transport": "sse",
}
}
) as client:
agent = create_react_agent(model, client.get_tools())
response = await agent.ainvoke(
{"messages": HumanMessage(content="get function in 4905")}
)
assert response is not None
assert "messages" in response
assert len(response["messages"]) > 0
content = response["messages"][-1].content
assert isinstance(content, str)
assert len(content) > 0
async def test_ida_server_connection():
async with MultiServerMCPClient(
{
"ida": {
"url": "http://localhost:8888/sse",
"transport": "sse",
}
}
) as client:
assert client is not None
tools = client.get_tools()
assert len(tools) > 0
```
--------------------------------------------------------------------------------
/test/heap/main.c:
--------------------------------------------------------------------------------
```cpp
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#define MAX_NOTES 100
#define MAX_TITLE_LEN 50
#define MAX_CONTENT_LEN 500
#define FILENAME "notes.dat"
// Note structure
typedef struct {
int id;
char title[MAX_TITLE_LEN];
char content[MAX_CONTENT_LEN];
char date[20];
int isDeleted; // Flag to mark if note is deleted
} Note;
// Global variables
Note notes[MAX_NOTES];
int noteCount = 0;
int lastId = 0;
// Function declarations
void loadNotes();
void saveNotes();
void addNote();
void deleteNote();
void updateNote();
void searchNote();
void listAllNotes();
void getCurrentDate(char *dateStr);
void clearInputBuffer();
int findNoteById(int id);
void printNote(Note note);
void printMenu();
int main() {
loadNotes(); // Load existing notes
int choice;
do {
printMenu();
scanf("%d", &choice);
clearInputBuffer();
switch(choice) {
case 1:
addNote();
break;
case 2:
deleteNote();
break;
case 3:
updateNote();
break;
case 4:
searchNote();
break;
case 5:
listAllNotes();
break;
case 0:
saveNotes();
printf("Thank you for using Note Manager! Goodbye!\n");
break;
default:
printf("Invalid choice, please try again.\n");
}
} while(choice != 0);
return 0;
}
// Print menu
void printMenu() {
printf("\n===== Note Manager =====\n");
printf("1. Add Note\n");
printf("2. Delete Note\n");
printf("3. Update Note\n");
printf("4. Search Note\n");
printf("5. List All Notes\n");
printf("0. Exit\n");
printf("Please select an operation: ");
}
// Clear input buffer
void clearInputBuffer() {
int c;
while ((c = getchar()) != '\n' && c != EOF);
}
// Get current date
void getCurrentDate(char *dateStr) {
time_t now = time(NULL);
struct tm *t = localtime(&now);
sprintf(dateStr, "%04d-%02d-%02d",
t->tm_year + 1900, t->tm_mon + 1, t->tm_mday);
}
// Load notes data
void loadNotes() {
FILE *file = fopen(FILENAME, "rb");
if (file == NULL) {
printf("Note file not found, creating new file.\n");
return;
}
noteCount = 0;
lastId = 0;
while (fread(¬es[noteCount], sizeof(Note), 1, file) == 1 && noteCount < MAX_NOTES) {
if (notes[noteCount].id > lastId) {
lastId = notes[noteCount].id;
}
noteCount++;
}
fclose(file);
printf("Successfully loaded %d notes.\n", noteCount);
}
// Save notes data
void saveNotes() {
FILE *file = fopen(FILENAME, "wb");
if (file == NULL) {
printf("Cannot open file to save notes.\n");
return;
}
for (int i = 0; i < noteCount; i++) {
if (!notes[i].isDeleted) {
fwrite(¬es[i], sizeof(Note), 1, file);
}
}
fclose(file);
printf("Notes have been saved.\n");
}
// Add note
void addNote() {
if (noteCount >= MAX_NOTES) {
printf("Maximum number of notes reached, cannot add more.\n");
return;
}
Note newNote;
newNote.id = ++lastId;
newNote.isDeleted = 0;
printf("Enter note title: ");
fgets(newNote.title, MAX_TITLE_LEN, stdin);
newNote.title[strcspn(newNote.title, "\n")] = 0; // Remove newline
printf("Enter note content: ");
fgets(newNote.content, MAX_CONTENT_LEN, stdin);
newNote.content[strcspn(newNote.content, "\n")] = 0; // Remove newline
getCurrentDate(newNote.date);
notes[noteCount++] = newNote;
printf("Note added, ID: %d\n", newNote.id);
saveNotes();
}
// Find note
int findNoteById(int id) {
for (int i = 0; i < noteCount; i++) {
if (notes[i].id == id && !notes[i].isDeleted) {
return i;
}
}
return -1;
}
// Print note
void printNote(Note note) {
printf("ID: %d\n", note.id);
printf("Title: %s\n", note.title);
printf("Content: %s\n", note.content);
printf("Date: %s\n", note.date);
printf("-----------------------\n");
}
// Delete note
void deleteNote() {
int id;
printf("Enter the ID of note to delete: ");
scanf("%d", &id);
clearInputBuffer();
int index = findNoteById(id);
if (index == -1) {
printf("Note does not exist.\n");
return;
}
notes[index].isDeleted = 1;
printf("Note has been deleted.\n");
saveNotes();
}
// Update note
void updateNote() {
int id;
printf("Enter the ID of note to update: ");
scanf("%d", &id);
clearInputBuffer();
int index = findNoteById(id);
if (index == -1) {
printf("Note does not exist.\n");
return;
}
printf("Current note content:\n");
printNote(notes[index]);
printf("Enter new title (leave empty to keep unchanged): ");
char newTitle[MAX_TITLE_LEN];
fgets(newTitle, MAX_TITLE_LEN, stdin);
newTitle[strcspn(newTitle, "\n")] = 0;
printf("Enter new content (leave empty to keep unchanged): ");
char newContent[MAX_CONTENT_LEN];
fgets(newContent, MAX_CONTENT_LEN, stdin);
newContent[strcspn(newContent, "\n")] = 0;
if (strlen(newTitle) > 0) {
strcpy(notes[index].title, newTitle);
}
if (strlen(newContent) > 0) {
strcpy(notes[index].content, newContent);
}
printf("Note has been updated.\n");
saveNotes();
}
// Search note
void searchNote() {
printf("1. Search by ID\n");
printf("2. Search by title keyword\n");
printf("Choose search method: ");
int choice;
scanf("%d", &choice);
clearInputBuffer();
if (choice == 1) {
int id;
printf("Enter note ID: ");
scanf("%d", &id);
clearInputBuffer();
int index = findNoteById(id);
if (index == -1) {
printf("Note does not exist.\n");
return;
}
printf("\nMatching note found:\n");
printNote(notes[index]);
}
else if (choice == 2) {
char keyword[MAX_TITLE_LEN];
printf("Enter title keyword: ");
fgets(keyword, MAX_TITLE_LEN, stdin);
keyword[strcspn(keyword, "\n")] = 0;
int found = 0;
printf("\nFound notes:\n");
for (int i = 0; i < noteCount; i++) {
if (!notes[i].isDeleted && strstr(notes[i].title, keyword) != NULL) {
printNote(notes[i]);
found = 1;
}
}
if (!found) {
printf("No matching notes found.\n");
}
}
else {
printf("Invalid choice.\n");
}
}
// List all notes
void listAllNotes() {
int count = 0;
printf("\nAll notes:\n");
for (int i = 0; i < noteCount; i++) {
if (!notes[i].isDeleted) {
printNote(notes[i]);
count++;
}
}
if (count == 0) {
printf("No notes available.\n");
} else {
printf("Total %d notes.\n", count);
}
}
```
--------------------------------------------------------------------------------
/src/headless_ida_mcp_server/server.py:
--------------------------------------------------------------------------------
```python
# -*- coding: utf-8 -*-
from mcp.server import FastMCP
from mcp.server.fastmcp.prompts import base
from functools import wraps
from typing import Any, Callable, get_type_hints, TypedDict, Optional, Annotated
import struct
from headless_ida_mcp_server.helper import IDA
from headless_ida_mcp_server.logger import logger
from headless_ida_mcp_server import PORT,TRANSPORT
mcp = FastMCP("IDA MCP Server", port=PORT)
ida = None
@mcp.tool()
def set_binary_path(path: Annotated[str, "Path to the binary file"]):
"""Set the path to the binary file"""
global ida
ida = IDA(path)
return "Binary path set"
@mcp.tool()
def get_function(address: Annotated[int, "Address of the function"]):
"""Get a function by address"""
if ida is None:
raise ValueError("Binary path not set")
return ida.get_function(address)
@mcp.tool()
def get_function_by_name(name: Annotated[str, "Name of the function"]):
"""Get a function by name"""
if ida is None:
raise ValueError("Binary path not set")
return ida.get_function_by_name(name)
@mcp.tool()
def get_function_by_address(address: Annotated[int, "Address of the function"]):
"""Get a function by address"""
if ida is None:
raise ValueError("Binary path not set")
return ida.get_function_by_address(address)
@mcp.tool()
def get_current_address():
"""Get the current address"""
if ida is None:
raise ValueError("Binary path not set")
return ida.get_current_address()
@mcp.tool()
def get_current_function():
"""Get the current function"""
if ida is None:
raise ValueError("Binary path not set")
return ida.get_current_function()
@mcp.tool()
def convert_number(text: Annotated[str, "Textual representation of the number to convert"],size: Annotated[Optional[int], "Size of the variable in bytes"]):
"""Convert a number to a different representation"""
if ida is None:
raise ValueError("Binary path not set")
return ida.convert_number(text, size)
@mcp.tool()
def list_functions():
"""List all functions"""
if ida is None:
raise ValueError("Binary path not set")
return ida.list_functions()
@mcp.tool()
def decompile_checked(address: Annotated[int, "Address of the function to decompile"]):
"""Decompile a function at the given address"""
if ida is None:
raise ValueError("Binary path not set")
return ida.decompile_checked(address)
@mcp.tool()
def decompile_function(address: Annotated[int, "Address of the function to decompile"]):
"""Decompile a function at the given address"""
if ida is None:
raise ValueError("Binary path not set")
return ida.decompile_function(address)
@mcp.tool()
def disassemble_function(address: Annotated[int, "Address of the function to disassemble"]):
"""Disassemble a function at the given address"""
if ida is None:
raise ValueError("Binary path not set")
return ida.disassemble_function(address)
@mcp.tool()
def get_xrefs_to(address: Annotated[int, "Address to get cross references to"]):
"""Get cross references to a given address"""
if ida is None:
raise ValueError("Binary path not set")
return ida.get_xrefs_to(address)
@mcp.tool()
def get_entry_points():
"""Get all entry points of the binary"""
if ida is None:
raise ValueError("Binary path not set")
return ida.get_entry_points()
@mcp.tool()
def set_decompiler_comment(address: Annotated[int, "Address in the function to set the comment for"],comment: Annotated[str, "Comment text (not shown in the disassembly)"]):
"""Set a comment for a given address in the function pseudocode"""
if ida is None:
raise ValueError("Binary path not set")
return ida.set_decompiler_comment(address, comment)
@mcp.tool()
def set_disassembly_comment(address: Annotated[int, "Address in the function to set the comment for"],comment: Annotated[str, "Comment text (not shown in the pseudocode)"]):
"""Set a comment for a given address in the function disassembly"""
if ida is None:
raise ValueError("Binary path not set")
return ida.set_disassembly_comment(address, comment)
@mcp.tool()
def refresh_decompiler_widget():
"""Refresh the decompiler widget"""
if ida is None:
raise ValueError("Binary path not set")
return ida.refresh_decompiler_widget()
@mcp.tool()
def refresh_decompiler_ctext(function_address: Annotated[int, "Address of the function to refresh the decompiler ctext for"]):
"""Refresh the decompiler ctext for a given function"""
if ida is None:
raise ValueError("Binary path not set")
return ida.refresh_decompiler_ctext(function_address)
@mcp.tool()
def rename_local_variable(function_address: Annotated[int, "Address of the function containing the variable"],old_name: Annotated[str, "Current name of the variable"],new_name: Annotated[str, "New name for the variable"]):
"""Rename a local variable in a function"""
if ida is None:
raise ValueError("Binary path not set")
return ida.rename_local_variable(function_address, old_name, new_name)
@mcp.tool()
def rename_function(function_address: Annotated[int, "Address of the function to rename"],new_name: Annotated[str, "New name for the function"]):
"""Rename a function"""
if ida is None:
raise ValueError("Binary path not set")
return ida.rename_function(function_address, new_name)
@mcp.tool()
def set_function_prototype(function_address: Annotated[int, "Address of the function"],prototype: Annotated[str, "New function prototype"]):
"""Set a function's prototype"""
if ida is None:
raise ValueError("Binary path not set")
return ida.set_function_prototype(function_address, prototype)
@mcp.tool()
def save_idb_file(save_path: Annotated[str, "Path to save the IDB file"]):
"""Save the IDB file"""
if ida is None:
raise ValueError("Binary path not set")
return ida.save_idb_file(save_path)
@mcp.prompt()
def exploit_prompt():
"""Exploit prompt"""
messages = [
base.UserMessage("You are a helpful assistant that can help me with my exploit."),
base.UserMessage("""
You need to follow these steps to complete the exploit:
1. Reverse analyze the binary file to locate vulnerabilities and analyze vulnerability types
- Locate vulnerabilities: Use IDA Pro's analysis features to find vulnerability locations
- Need to first gather binary file information, such as using checksec tool to check binary protection mechanisms
- If you find NX protection is disabled, you can use ret2shellcode method to get a shell
2. Choose appropriate exploitation method based on vulnerability type
- For stack overflow vulnerabilities, analyze the overflow pattern and check if there are backdoor functions in the binary. If there are backdoor functions, modify the return address to point to the backdoor function address
- If there are no backdoor functions, need to use ret2libc method. Don't assume the binary contains the gadgets you want - you need to use ROPgadget to find gadgets and combine them to construct the pop chain.
"""),
]
return messages
def main():
mcp.run(transport = TRANSPORT)
if __name__ == "__main__":
main()
```
--------------------------------------------------------------------------------
/src/headless_ida_mcp_server/helper.py:
--------------------------------------------------------------------------------
```python
"""
Helper module for IDA Python integration.
Provides configuration and IDA interface initialization.
"""
from headless_ida import HeadlessIda
import logging
from headless_ida_mcp_server import IDA_PATH
import struct
from typing import Optional, TypedDict, Annotated
class Function(TypedDict):
start_address: int
end_address: int
name: str
prototype: Optional[str]
class IDAError(Exception):
def __init__(self, message: str):
super().__init__(message)
@property
def message(self) -> str:
return self.args[0]
class Xref(TypedDict):
address: int
type: str
function: Optional[Function]
class ConvertedNumber(TypedDict):
decimal: str
hexadecimal: str
bytes: str
ascii: Optional[str]
binary: str
class IDA():
def __init__(self, binary_path: Annotated[str, "Path to the binary file"]):
self.headlessida = HeadlessIda(IDA_PATH,binary_path)
self.idaapi = self.headlessida.import_module("idaapi")
self.idautils = self.headlessida.import_module("idautils")
self.ida_entry = self.headlessida.import_module("ida_entry")
self.ida_funcs = self.headlessida.import_module("ida_funcs")
self.ida_hexrays = self.headlessida.import_module("ida_hexrays")
self.ida_kernwin = self.headlessida.import_module("ida_kernwin")
self.ida_lines = self.headlessida.import_module("ida_lines")
self.ida_nalt = self.headlessida.import_module("ida_nalt")
self.ida_name = self.headlessida.import_module("ida_name")
self.ida_typeinf = self.headlessida.import_module("ida_typeinf")
self.ida_xref = self.headlessida.import_module("ida_xref")
self.idc = self.headlessida.import_module("idc")
self.ida_loader = self.headlessida.import_module("ida_loader")
######## COPY from https://github.com/mrexodia/ida-pro-mcp ########
def get_image_size(self):
try:
# https://www.hex-rays.com/products/ida/support/sdkdoc/structidainfo.html
info = self.idaapi.get_inf_structure()
omin_ea = info.omin_ea
omax_ea = info.omax_ea
except AttributeError:
import ida_ida
omin_ea = ida_ida.inf_get_omin_ea()
omax_ea = ida_ida.inf_get_omax_ea()
# Bad heuristic for image size (bad if the relocations are the last section)
image_size = omax_ea - omin_ea
# Try to extract it from the PE header
header = self.idautils.peutils_t().header()
if header and header[:4] == b"PE\0\0":
image_size = struct.unpack("<I", header[0x50:0x54])[0]
return image_size
def get_prototype(self,fn) -> Optional[str]:
try:
if isinstance(fn, self.ida_funcs.func_t):
func = fn
else:
func = self.idaapi.get_func(fn)
if func is None:
return None
tif = self.ida_typeinf.tinfo_t()
if self.ida_nalt.get_tinfo(tif, func.start_ea):
return str(tif)
return self.idc.get_type(func.start_ea)
except Exception as e:
print(f"Error getting function prototype: {e}")
return None
def get_function(self,address: int, *, raise_error=True) -> Optional[Function]:
fn = self.idaapi.get_func(address)
if fn is None:
if raise_error:
raise IDAError(f"No function found at address {address}")
return None
try:
name = fn.get_name()
except AttributeError:
name = self.ida_funcs.get_func_name(fn.start_ea)
return {
"address": fn.start_ea,
"end_address": fn.end_ea,
"name": name,
"prototype": self.get_prototype(fn),
}
def get_function_by_name(self,name: Annotated[str, "Name of the function to get"]) -> Function:
"""Get a function by its name"""
function_address = self.idaapi.get_name_ea(self.idaapi.BADADDR, name)
if function_address == self.idaapi.BADADDR:
raise IDAError(f"No function found with name {name}")
return self.get_function(function_address)
def get_function_by_address(self,address: Annotated[int, "Address of the function to get"]) -> Function:
"""Get a function by its address"""
return self.get_function(address)
def get_current_address(self) -> int:
"""Get the address currently selected by the user"""
return self.idaapi.get_screen_ea()
def get_current_function(self) -> Optional[Function]:
"""Get the function currently selected by the user"""
return self.get_function(self.idaapi.get_screen_ea())
def convert_number(self,text: Annotated[str, "Textual representation of the number to convert"],size: Annotated[Optional[int], "Size of the variable in bytes"]) -> ConvertedNumber:
"""Convert a number (decimal, hexadecimal) to different representations"""
try:
value = int(text, 0)
except ValueError:
raise IDAError(f"Invalid number: {text}")
# Estimate the size of the number
if not size:
size = 0
n = abs(value)
while n:
size += 1
n >>= 1
size += 7
size //= 8
# Convert the number to bytes
try:
bytes = value.to_bytes(size, "little", signed=True)
except OverflowError:
raise IDAError(f"Number {text} is too big for {size} bytes")
# Convert the bytes to ASCII
ascii = ""
for byte in bytes.rstrip(b"\x00"):
if byte >= 32 and byte <= 126:
ascii += chr(byte)
else:
ascii = None
break
return {
"decimal": str(value),
"hexadecimal": hex(value),
"bytes": bytes.hex(" "),
"ascii": ascii,
"binary": bin(value)
}
def list_functions(self) -> list[Function]:
"""List all functions in the database"""
return [self.get_function(address) for address in self.idautils.Functions()]
def decompile_checked(self,address: int):
if not self.ida_hexrays.init_hexrays_plugin():
raise IDAError("Hex-Rays decompiler is not available")
error = self.ida_hexrays.hexrays_failure_t()
cfunc: self.ida_hexrays.cfunc_t = self.ida_hexrays.decompile_func(address, error, self.ida_hexrays.DECOMP_WARNINGS)
if not cfunc:
message = f"Decompilation failed at {address}"
if error.str:
message += f": {error.str}"
if error.errea != self.idaapi.BADADDR:
message += f" (address: {error.errea})"
raise IDAError(message)
return cfunc
def decompile_function(self,address: Annotated[int, "Address of the function to decompile"]) -> str:
"""Decompile a function at the given address"""
cfunc = self.decompile_checked(address)
sv = cfunc.get_pseudocode()
pseudocode = ""
for i, sl in enumerate(sv):
sl: self.ida_kernwin.simpleline_t
item = self.ida_hexrays.ctree_item_t()
addr = None if i > 0 else cfunc.entry_ea
if cfunc.get_line_item(sl.line, 0, False, None, item, None):
ds = item.dstr().split(": ")
if len(ds) == 2:
try:
addr = int(ds[0], 16)
except ValueError:
pass
line = self.ida_lines.tag_remove(sl.line)
if len(pseudocode) > 0:
pseudocode += "\n"
if addr is None:
pseudocode += f"/* line: {i} */ {line}"
else:
pseudocode += f"/* line: {i}, address: {addr} */ {line}"
return pseudocode
def disassemble_function(self,address: Annotated[int, "Address of the function to disassemble"]) -> str:
"""Get assembly code (address: instruction; comment) for a function"""
func = self.idaapi.get_func(address)
if not func:
raise IDAError(f"No function found at address {address}")
# TODO: add labels
disassembly = ""
for address in self.ida_funcs.func_item_iterator_t(func):
if len(disassembly) > 0:
disassembly += "\n"
disassembly += f"{address}: "
disassembly += self.idaapi.generate_disasm_line(address, self.idaapi.GENDSM_REMOVE_TAGS)
comment = self.idaapi.get_cmt(address, False)
if not comment:
comment = self.idaapi.get_cmt(address, True)
if comment:
disassembly += f"; {comment}"
return disassembly
def get_xrefs_to(self,address: Annotated[int, "Address to get cross references to"]) -> list[Xref]:
"""Get all cross references to the given address"""
xrefs = []
xref: self.ida_xref.xrefblk_t
for xref in self.idautils.XrefsTo(address):
xrefs.append({
"address": xref.frm,
"type": "code" if xref.iscode else "data",
"function": self.get_function(xref.frm, raise_error=False),
})
return xrefs
def get_entry_points(self) -> list[Function]:
"""Get all entry points in the database"""
result = []
for i in range(self.ida_entry.get_entry_qty()):
ordinal = self.ida_entry.get_entry_ordinal(i)
address = self.ida_entry.get_entry(ordinal)
func = self.get_function(address, raise_error=False)
if func is not None:
result.append(func)
return result
def set_decompiler_comment(self,address: Annotated[int, "Address in the function to set the comment for"],comment: Annotated[str, "Comment text (not shown in the disassembly)"]):
"""Set a comment for a given address in the function pseudocode"""
# Reference: https://cyber.wtf/2019/03/22/using-ida-python-to-analyze-trickbot/
# Check if the address corresponds to a line
cfunc = self.decompile_checked(address)
# Special case for function entry comments
if address == cfunc.entry_ea:
self.idc.set_func_cmt(address, comment, True)
cfunc.refresh_func_ctext()
return
eamap = cfunc.get_eamap()
if address not in eamap:
raise IDAError(f"Failed to set comment at {address}")
nearest_ea = eamap[address][0].ea
# Remove existing orphan comments
if cfunc.has_orphan_cmts():
cfunc.del_orphan_cmts()
cfunc.save_user_cmts()
# Set the comment by trying all possible item types
tl = self.idaapi.treeloc_t()
tl.ea = nearest_ea
for itp in range(self.idaapi.ITP_SEMI, self.idaapi.ITP_COLON):
tl.itp = itp
cfunc.set_user_cmt(tl, comment)
cfunc.save_user_cmts()
cfunc.refresh_func_ctext()
if not cfunc.has_orphan_cmts():
return
cfunc.del_orphan_cmts()
cfunc.save_user_cmts()
raise IDAError(f"Failed to set comment at {address}")
def set_disassembly_comment(self,address: Annotated[int, "Address in the function to set the comment for"],comment: Annotated[str, "Comment text (not shown in the pseudocode)"]):
"""Set a comment for a given address in the function disassembly"""
if not self.idaapi.set_cmt(address, comment, False):
raise IDAError(f"Failed to set comment at {address}")
def refresh_decompiler_widget(self):
widget = self.ida_kernwin.get_current_widget()
if widget is not None:
vu = self.ida_hexrays.get_widget_vdui(widget)
if vu is not None:
vu.refresh_ctext()
def refresh_decompiler_ctext(self,function_address: int):
error = self.ida_hexrays.hexrays_failure_t()
cfunc: self.ida_hexrays.cfunc_t = self.ida_hexrays.decompile_func(function_address, error, self.ida_hexrays.DECOMP_WARNINGS)
if cfunc:
cfunc.refresh_func_ctext()
def rename_local_variable(self,function_address: Annotated[int, "Address of the function containing the variable"],old_name: Annotated[str, "Current name of the variable"],new_name: Annotated[str, "New name for the variable"]):
"""Rename a local variable in a function"""
func = self.idaapi.get_func(function_address)
if not func:
raise IDAError(f"No function found at address {function_address}")
if not self.ida_hexrays.rename_lvar(func.start_ea, old_name, new_name):
raise IDAError(f"Failed to rename local variable {old_name} in function {func.start_ea}")
self.refresh_decompiler_ctext(func.start_ea)
def rename_function(self,function_address: Annotated[int, "Address of the function to rename"],new_name: Annotated[str, "New name for the function"]):
"""Rename a function"""
fn = self.idaapi.get_func(function_address)
if not fn:
raise IDAError(f"No function found at address {function_address}")
if not self.idaapi.set_name(fn.start_ea, new_name):
raise IDAError(f"Failed to rename function {fn.start_ea} to {new_name}")
self.refresh_decompiler_ctext(fn.start_ea)
def set_function_prototype(self,function_address: Annotated[int, "Address of the function"],prototype: Annotated[str, "New function prototype"]) -> str:
"""Set a function's prototype"""
fn = self.idaapi.get_func(function_address)
if not fn:
raise IDAError(f"No function found at address {function_address}")
try:
tif = self.ida_typeinf.tinfo_t(prototype, None, self.ida_typeinf.PT_SIL)
if not tif.is_func():
raise IDAError(f"Parsed declaration is not a function type")
if not self.ida_typeinf.apply_tinfo(fn.start_ea, tif, self.ida_typeinf.PT_SIL):
raise IDAError(f"Failed to apply type")
self.refresh_decompiler_ctext(fn.start_ea)
except Exception as e:
raise IDAError(f"Failed to parse prototype string: {prototype}")
def save_idb_file(self,save_path: Annotated[str, "Path to save the IDB file"]):
self.ida_loader.save_database(save_path, 0)
```