├── .gitignore ├── imgs └── README │ ├── image-20250416205310491.png │ ├── image-20250416205344433.png │ ├── image-20250416205421973.png │ ├── image-20250416205428185.png │ ├── image-20250416205552353.png │ └── image-20250416205722302.png ├── config_sample.json ├── ida-plugin.json ├── README.md └── ComprehendAI.py /.gitignore: -------------------------------------------------------------------------------- 1 | config.json 2 | -------------------------------------------------------------------------------- /imgs/README/image-20250416205310491.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kvancy/ComprehendAI/HEAD/imgs/README/image-20250416205310491.png -------------------------------------------------------------------------------- /imgs/README/image-20250416205344433.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kvancy/ComprehendAI/HEAD/imgs/README/image-20250416205344433.png -------------------------------------------------------------------------------- /imgs/README/image-20250416205421973.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kvancy/ComprehendAI/HEAD/imgs/README/image-20250416205421973.png -------------------------------------------------------------------------------- /imgs/README/image-20250416205428185.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kvancy/ComprehendAI/HEAD/imgs/README/image-20250416205428185.png -------------------------------------------------------------------------------- /imgs/README/image-20250416205552353.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kvancy/ComprehendAI/HEAD/imgs/README/image-20250416205552353.png -------------------------------------------------------------------------------- /imgs/README/image-20250416205722302.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Kvancy/ComprehendAI/HEAD/imgs/README/image-20250416205722302.png -------------------------------------------------------------------------------- /config_sample.json: -------------------------------------------------------------------------------- 1 | { 2 | "openai": { 3 | "model": "qwq-32b", 4 | "api_key": "sk-", 5 | "base_url": "https://" 6 | } 7 | } -------------------------------------------------------------------------------- /ida-plugin.json: -------------------------------------------------------------------------------- 1 | { 2 | "IDAMetadataDescriptorVersion": 1, 3 | "plugin": { 4 | "name": "IDA ComprehendAI", 5 | "entryPoint": "ComprehendAI.py" 6 | } 7 | } -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ComprehendAI 2 | 3 | An AI plugin for assisting IDA reverse analysis, which facilitates quickly summarizing the functions of code and accelerates the analysis efficiency. 4 | 5 | ## Features 6 | 7 | **Non - blocking AI Analysis** 8 | 9 | - Description: This feature enables non - blocking AI analysis. While the analysis is in progress, you can continue with your work uninterrupted. Once the analysis is completed, the results will be printed in the output window. 10 | 11 | **Customizable Function Analysis Depth** 12 | 13 | - Description: You have the flexibility to set the depth of function analysis according to your needs. 14 | 15 | **Manual Interaction with AI** 16 | 17 | - Description: You can manually ask the AI various questions and perform any operations you prefer. 18 | 19 | **Support for Streaming Output with Inference Models** 20 | 21 | - Description: Supports the use of inference models with streaming output. Get continuous results as the model processes data. 22 | 23 | 24 | 25 | ## Usage 26 | 27 | #### 1. Project Retrieval 28 | 29 | First, you need to pull the project to your local machine. Open your terminal or command prompt and use the following command to clone the project repository: 30 | 31 | ```bash 32 | git clone https://github.com/Kvancy/ComprehendAI.git 33 | ``` 34 | 35 | #### 2. File Placement 36 | 37 | Navigate to the directory of the cloned project. Locate the `config.json` and `ComprehendAI.py` files. Then, place these two files into the `plugins` folder of IDA . 38 | 39 | #### 3. Configuration File Setup 40 | 41 | Open the `config.json` file. You will see a content structure similar to the following: 42 | 43 | ```json 44 | { 45 | "openai"{ 46 | "model":"", 47 | "api_key": "", 48 | "base_url": "" 49 | } 50 | } 51 | ``` 52 | 53 | Replace the content within the double - quotes with your own `api_key` and `base_url`. For example: 54 | 55 | ```json 56 | { 57 | "openai"{ 58 | "model":"", 59 | "api_key": "your_actual_api_key", 60 | "base_url": "your_actual_base_url" 61 | } 62 | } 63 | ``` 64 | 65 | Save and close the `config.json` file. 66 | 67 | #### 4. Dependencies 68 | 69 | You need to install the following Python libraries using `pip`. You can install the `openai` library with the following command: 70 | 71 | ```py 72 | pip install openai 73 | ``` 74 | 75 | #### 4. Launch IDA and Load the Plugin 76 | 77 | It has been successfully tested in IDA 9.1 and IDA 7.7. 78 | 79 | ### Example 80 | 81 | Right - click on the disassembly window to pop up the menu and select a function. 82 | 83 | Analysis allows you to continue your work, and then the results will be printed in the output window. 84 | 85 | ![image-20250416205310491](./imgs/README/image-20250416205310491.png) 86 | 87 | You can set the depth of function analysis by yourself. If you're only interested in the current function or you have to consider the tokens consumption, just set the depth to 0. 88 | 89 | ![image-20250416205344433](./imgs/README/image-20250416205344433.png) 90 | 91 | You can also manually ask the AI some questions and do whatever you like,Or ask questions while referring to the code. 92 | 93 | ![image-20250416205428185](./imgs/README/image-20250416205428185.png) 94 | 95 | You can easily pause the output whenever you want. 96 | 97 | ![image-20250416205722302](./imgs/README/image-20250416205722302.png) 98 | 99 | You have the freedom to customize the prompt according to your needs for the moment. 100 | 101 | ![image-20250416205552353](./imgs/README/image-20250416205552353.png) 102 | 103 | -------------------------------------------------------------------------------- /ComprehendAI.py: -------------------------------------------------------------------------------- 1 | import traceback 2 | import idaapi 3 | import idc 4 | import idautils 5 | import ida_xref 6 | import json 7 | import os 8 | import asyncio 9 | 10 | from idaapi import action_handler_t, UI_Hooks 11 | from threading import Lock,Thread,Event 12 | from openai import OpenAI 13 | from enum import Enum 14 | 15 | class TaskType(Enum): 16 | ANALYSIS = 1 17 | CUSTOM_QUERY = 2 18 | CUSTON_QUERY_WITH_CODE = 3 19 | 20 | class QueryStatus(Enum): 21 | SUCCESS = 1 22 | FAILED = 2 23 | STOPPED = 3 24 | 25 | #处理配置文件 26 | class ConfigManager: 27 | _instance = None 28 | _lock = Lock() 29 | 30 | def __new__(cls): 31 | with cls._lock: 32 | if not cls._instance: 33 | cls._instance = super().__new__(cls) 34 | cls._instance._initialize() 35 | return cls._instance 36 | 37 | def _initialize(self): 38 | self.script_dir = os.path.dirname(os.path.abspath(__file__)) 39 | self.config_path = os.path.join(self.script_dir, 'config.json') 40 | self.config = self._load_config() 41 | self.openai_client = self._create_openai_client() 42 | 43 | def _load_config(self): 44 | try: 45 | with open(self.config_path, "r") as f: 46 | return json.load(f) 47 | except Exception as e: 48 | raise RuntimeError(f"Failed to load config: {str(e)}") 49 | 50 | def _create_openai_client(self): 51 | return OpenAI( 52 | api_key=self.config["openai"]["api_key"], 53 | base_url=self.config["openai"]["base_url"] 54 | ) 55 | 56 | @property 57 | def model_name(self): 58 | return self.config["openai"]["model"] 59 | 60 | @property 61 | def client(self): 62 | return self.openai_client 63 | #处理反汇编代码提取 64 | class DisassemblyProcessor: 65 | def __init__(self, max_depth=2): 66 | self.max_depth = max_depth 67 | self._lock = Lock() 68 | self._reset_state() 69 | 70 | def _reset_state(self): 71 | with self._lock: 72 | self.processed_funcs = set() 73 | self.func_disasm_list = [] 74 | 75 | def get_current_function_disasm(self): 76 | self._reset_state() 77 | 78 | current_ea = idc.get_screen_ea() 79 | func_start = idc.get_func_attr(current_ea, idc.FUNCATTR_START) 80 | 81 | if func_start == idaapi.BADADDR: 82 | raise ValueError("Failed to locate function start address") 83 | 84 | self._process_function(func_start, self.max_depth) 85 | return "\n".join(self.func_disasm_list) 86 | 87 | def _process_function(self, func_ea, depth): 88 | if func_ea in self.processed_funcs or depth < 0: 89 | return 90 | 91 | with self._lock: 92 | self.processed_funcs.add(func_ea) 93 | 94 | try: 95 | decompiled = str(idaapi.decompile(func_ea)) 96 | with self._lock: 97 | self.func_disasm_list.append(decompiled) 98 | except Exception as e: 99 | print(f"Decompilation failed for {hex(func_ea)}: {str(e)}") 100 | 101 | for callee in self._get_callees(func_ea): 102 | self._process_function(callee, depth - 1) 103 | 104 | def _get_callees(self, func_ea): 105 | callees = set() 106 | for ea in range(func_ea, idc.get_func_attr(func_ea, idc.FUNCATTR_END)): 107 | for xref in idautils.XrefsFrom(ea): 108 | if xref.type in [ida_xref.fl_CN, ida_xref.fl_CF]: 109 | callee_ea = xref.to 110 | if idc.get_func_attr(callee_ea, idc.FUNCATTR_START) == callee_ea: 111 | callees.add(callee_ea) 112 | return callees 113 | #处理openai接口相关 114 | class AIService: 115 | def __init__(self): 116 | self.config = ConfigManager() 117 | self.stop_event = Event() 118 | 119 | def ask_ai(self, prompt, ai_isRunning:Lock): 120 | messages = [{"role": "user", "content": prompt}] 121 | print("ComprehendAI output:") 122 | self.stop_event.clear() #初始化事件 123 | 124 | result = self._request_openai(messages) 125 | ai_isRunning.release() # 分析完成，无论成功失败都需释放锁 126 | 127 | match result: 128 | case QueryStatus.SUCCESS: 129 | print("\r✅ 分析完成！") 130 | case QueryStatus.FAILED: 131 | print("\r❌ 分析失败，请重试") 132 | case QueryStatus.STOPPED: 133 | print("\r✅ 分析暂停") 134 | 135 | 136 | def _request_openai(self,messages): 137 | reasoning_content = "" 138 | answer_content = "" 139 | is_answering = False 140 | try: 141 | completion = self.config.client.chat.completions.create( 142 | model=self.config.model_name, 143 | messages=messages, 144 | stream=True, 145 | ) 146 | for chunk in completion: 147 | if self.stop_event.is_set(): 148 | raise StopIteration("任务被停止") 149 | 150 | # 如果chunk.choices为空，则打印usage 151 | if not chunk.choices: 152 | print("\nUsage:") 153 | print(chunk.usage) 154 | else: 155 | delta = chunk.choices[0].delta 156 | # 打印思考过程 157 | if hasattr(delta, 'reasoning_content') and delta.reasoning_content != None: 158 | print(delta.reasoning_content, end='', flush=True) 159 | reasoning_content += delta.reasoning_content 160 | else: 161 | # 开始回复 162 | if delta.content != "" and is_answering is False: 163 | print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n") 164 | is_answering = True 165 | answer_content += delta.content 166 | 167 | print(answer_content) 168 | return QueryStatus.SUCCESS 169 | 170 | except StopIteration as e: 171 | print(f"Error occurred: {e}") 172 | return QueryStatus.STOPPED 173 | 174 | except Exception as e: 175 | print(f"Error occurred: {e}") 176 | traceback.print_exc() 177 | return QueryStatus.FAILED 178 | 179 | #处理用户接口 180 | class AnalysisHandler: 181 | 182 | def __init__(self): 183 | self.disassembler = DisassemblyProcessor() 184 | self.ai_service = AIService() 185 | self.ai_isRunning = Lock() 186 | self.prompt = """ 187 | 你是一名人工智能逆向工程专家。 188 | 我会提供你一些反汇编代码，其中首个函数是你需要分析并总结成报告的函数， 189 | 其余函数是该函数调用的一些子函数。 190 | 分析要求： 191 | 重点描述主函数功能，并对核心行为进行推测； 192 | 简要描述子函数功能 193 | 194 | 输出要求： 195 | 主函数功能：... 196 | 行为推测：... 197 | 子函数功能：... 198 | 纯文本输出。 199 | 200 | 下面是你要分析的反汇编代码： 201 | """ 202 | def set_analysis_depth(self, depth): 203 | self.disassembler.max_depth = depth 204 | 205 | def _create_analysis_prompt(self, disassembly): 206 | return self.prompt + disassembly 207 | 208 | def _create_analysis_custom_query(self, disassembly,question): 209 | return question + disassembly 210 | 211 | def create_ai_task(self,taskType,question=""): 212 | 213 | match taskType: 214 | case TaskType.ANALYSIS: 215 | disassembly = self.disassembler.get_current_function_disasm() 216 | prompt = self._create_analysis_prompt(disassembly) 217 | self.async_task(prompt) 218 | case TaskType.CUSTOM_QUERY: 219 | self.async_task(question) 220 | case TaskType.CUSTON_QUERY_WITH_CODE: 221 | disassembly = self.disassembler.get_current_function_disasm() 222 | prompt = self._create_analysis_custom_query(disassembly,question) 223 | self.async_task(prompt) 224 | 225 | 226 | def async_task(self,question): 227 | print(question) 228 | if self.ai_isRunning.acquire(blocking=False): 229 | task = Thread(target=self.ai_service.ask_ai,args=(question,self.ai_isRunning,)) 230 | task.start() 231 | 232 | else: 233 | print("\r❌ 当前AI正在处理任务,请稍后尝试") 234 | 235 | def stop(self): 236 | self.ai_service.stop_event.set() 237 | 238 | #处理插件框架 239 | class ComprehendAIPlugin(idaapi.plugin_t): 240 | flags = idaapi.PLUGIN_HIDE 241 | comment = "AI-based Reverse Analysis Plugin" 242 | help = "Perform AI-based analysis on binary code" 243 | wanted_name = "ComprehendAI" 244 | wanted_hotkey = "Ctrl+Shift+A" 245 | 246 | ACTION_DEFINITIONS = [ 247 | ("AI_analysis:Analysis", "Analysis", "执行非阻塞型AI分析"), 248 | ("AI_analysis:SetDepth", "Set analysis depth", "设置分析深度"), 249 | ("AI_analysis:SetPrompt", "Set your own prompt", "自定义prompt"), 250 | ("AI_analysis:CustomQueryWithCode", "Ask AI with code", "结合代码自定义提问"), 251 | ("AI_analysis:CustomQuery", "Ask AI", "自定义提问"), 252 | ("AI_analysis:Stop", "Stop", "停止"), 253 | ] 254 | 255 | def init(self): 256 | self.ui_hook = self.MenuHook() 257 | self.ui_hook.hook() 258 | 259 | self.handler = AnalysisHandler() 260 | self._register_actions() 261 | 262 | print("ComprehendAI initialized") 263 | return idaapi.PLUGIN_KEEP 264 | 265 | def run(self, arg): 266 | pass 267 | 268 | def term(self): 269 | self.ui_hook.unhook() 270 | self._unregister_actions() 271 | print("ComprehendAI unloaded") 272 | 273 | def _register_actions(self): 274 | for action_id, label, tooltip in self.ACTION_DEFINITIONS: 275 | action_desc = idaapi.action_desc_t( 276 | action_id, 277 | label, 278 | self.MenuCommandHandler(action_id,self.handler), 279 | None, 280 | tooltip, 281 | 0 282 | ) 283 | idaapi.register_action(action_desc) 284 | 285 | def _unregister_actions(self): 286 | for action_id, _, _ in self.ACTION_DEFINITIONS: 287 | idaapi.unregister_action(action_id) 288 | 289 | class MenuHook(UI_Hooks): 290 | def finish_populating_widget_popup(self, form, popup): 291 | if idaapi.get_widget_type(form) in (idaapi.BWN_DISASM, idaapi.BWN_PSEUDOCODE): 292 | for action_id, _, _ in ComprehendAIPlugin.ACTION_DEFINITIONS: 293 | idaapi.attach_action_to_popup(form, popup, action_id, "ComprehendAI/", idaapi.SETMENU_APP) 294 | 295 | class MenuCommandHandler(action_handler_t): 296 | def __init__(self, action_id,handler:AnalysisHandler): 297 | super().__init__() 298 | self.action_id = action_id 299 | self.handler = handler 300 | 301 | def activate(self, ctx): 302 | 303 | match self.action_id: 304 | case "AI_analysis:Analysis": 305 | self.handler.create_ai_task(TaskType.ANALYSIS) 306 | case "AI_analysis:CustomQuery": 307 | question = idaapi.ask_text(0, "", "输入问题") 308 | if question: 309 | self.handler.create_ai_task(TaskType.CUSTOM_QUERY, question) 310 | case "AI_analysis:SetDepth": 311 | new_depth = idaapi.ask_long(2, "设置分析深度 (默认2):") 312 | if new_depth is not None: 313 | self.handler.set_analysis_depth(new_depth) 314 | case "AI_analysis:SetPrompt": 315 | yourPrompt = idaapi.ask_text(0, f"{self.handler.prompt}", "输入问题") 316 | self.handler.prompt = yourPrompt 317 | case "AI_analysis:CustomQueryWithCode": 318 | question = idaapi.ask_text(0, "", "输入问题") 319 | if question: 320 | self.handler.create_ai_task(TaskType.CUSTON_QUERY_WITH_CODE, question) 321 | case "AI_analysis:Stop": 322 | self.handler.stop() 323 | return 1 324 | 325 | def update(self, ctx): 326 | return idaapi.AST_ENABLE_ALWAYS 327 | 328 | def PLUGIN_ENTRY(): 329 | return ComprehendAIPlugin() --------------------------------------------------------------------------------