reverse-engineering-assistant 1.0.0__tar.gz

reverse-engineering-assistant-1.0.0/PKG-INFO

@@ -0,0 +1,152 @@
+Metadata-Version: 2.1
+Name: reverse-engineering-assistant
+Version: 1.0.0
+Summary: An AI assistant for reverse engineering tasks
+Author: サイバーカイダ (cyberkaida)
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Description-Content-Type: text/markdown
+Requires-Dist: langchain
+Requires-Dist: llama-cpp-python
+Requires-Dist: prompt_toolkit
+Requires-Dist: sentence_transformers
+Requires-Dist: PyYAML
+Requires-Dist: pydantic
+Requires-Dist: rich
+Requires-Dist: Flask
+
+# ReVA - Reverse Engineering Assistant
+
+[✨ A quick demo! ✨](https://asciinema.org/a/626197)
+
+The reverse engineering assistant (ReVA) is a project to build a disassembler agnostic AI assistant for
+reverse engineering tasks. This includes both _offline_ and online inference and a simple architecture.
+
+RevA is different to other efforts at building AI assistants for RE tasks because it uses a technique
+called [embedding](https://openai.com/blog/introducing-text-and-code-embeddings)
+to give the AI assistant a sort of "long term memory". The model also is given access to a number of tools
+that are tweaked to perform well with queries provided by the LLM. This allows the model to reason about the whole
+program, rather than just a single function. The tools are tweaked to lead the AI to examine deeper.
+
+Using this technique you can ask general questions and get relevant answers. The model prioritises
+information from the embeddings and tools, but when there is no information it can still respond to generic
+questions from its training.
+
+You can ask questions like:
+- Does this program use encryption?
+- Draw a class diagram using plantuml syntax.
+- Rename all the variables in main with descriptive names.
+- Explain the purpose of the `__mod_init` segment.
+- What does `mmap` return?
+- What does the function at address 0x80000 do?
+
+## Large Language Model Support
+
+RevA is based on [llama-index](https://github.com/jerryjliu/llama_index),
+which supports a number of models.
+
+Built in support is provided for:
+- [OpenAI](https://platform.openai.com/overview) for online inference and easy setup (Needs an OpenAI API key)
+- [Ollama](https://ollama.ai) and any model it supports for local on-device inference or connecting to a self hosted remote inference server.
+
+Limited support is provided for:
+- [llama-cpp](https://llama-cpp-python.readthedocs.io/en/latest/) and any model it supports for local on-device inference
+- [text-generation-webui](https://github.com/oobabooga/text-generation-webui) and any model it supports for self-hosted remote inference
+
+Adding additional inference servers is easy if it is supported by llama-index or langchain (on which llama-index is based).
+
+See the configuration section for more information about setting the model.
+
+## Configuration
+
+Configuration for the reverse engineering assistant is stored at
+`~/.config/reverse-engineering-assistant/config.yaml`. If this
+is not present on first start, a default configuration using
+OpenAI for inference and the `OPENAI_API_TOKEN` environment
+variable will be used.
+
+The most important setting is the `type` top level setting.
+This controls what inference service you use. These are the
+same as the configuration keys, for example to use Ollama,
+set type to `ollama` and configure the settings in the `ollama:`
+section.
+
+The configuration also contains the prompts used for the models.
+If you use Ollama or OpenAI these will be processed to fit the
+model specific prompt pattern (placing the system prompt in the
+correct tags, etc).
+
+For `llama-cpp` and `text-generation-webui` these may need to be
+configured for your specific model. For this reason Ollama is
+preferred for self hosting.
+
+## Workflow
+
+RevA has a two step workflow.
+1. Generate knowledge base
+2. Perform inference
+
+To generate the knowledge base, use the plugin for your disassembler and run the Assistant script.
+See [Ghidra Support](#ghidra-support) and [BinaryNinja Support](#binary-ninja-support) below.
+
+First your disassembler extracts the information required for the knowledge base and embeddings.
+This involes extracting each function, it's decompilation and some metadata. These are written to a "project". This allows
+multiple programs and data sources to be combined into one set of knowledge for the assistant. For example multiple malware
+samples, or a program and its libraries could be included along with previous RE notes.
+
+Projects are stored in `~/.cache/reverse-engineering-assistant/projects`. If you make significant changes to your
+annotations or analysis in your disassembler, you should delete and regenerate your project directory. This cache
+is a _snapshot_ of the state of your disassembler.
+
+To ask questions and run the inference a command line tool is provided. Run `revassistant --project ${NAME_OF_YOUR_PROJECT}` to begin the chat session.
+
+`revassistant` will hash the knowledge base and generate and combine the embeddings into a searchable
+index. Once this is complete the index is saved to disk and the chat session begins.
+
+## Installation
+
+To install the particular extension for your disassembler see:
+- [Ghidra Support](#ghidra-support)
+- [Binary Ninja Support](#binary-ninja-support)
+
+To install the chat component you can do the following:
+
+```sh
+python3 -m pip install ./reverse-engineering-assistant
+```
+
+The chat can be started with:
+
+```sh
+revassistant --project ${NAME_OF_YOUR_PROJECT}
+```
+
+
+# Ghidra Support
+
+## Prerequisites
+- [Ghidrathon](https://github.com/mandiant/Ghidrathon) >= 2.2.0 installed into Ghidra
+
+## Usage
+
+After installation, enable the [Ghidrathon extension](https://github.com/mandiant/Ghidrathon#installing-ghidrathon)
+and the Ghidra Assistant Extension.
+
+You can generate the knowledge base by running the Ghidra Assistant analysis from the Analysis menu in the Code Browser.
+
+# Binary Ninja Support
+
+Install the ReVA BinaryNinja plugin by opening your BinaryNinja plugin directory (Plugins -> Open Plugin Folder)
+and copying or symbolic linking the [binary-ninja-assistant](./binary-ninja-assistant) directory into the plugin
+directory.
+
+Restart Binary Ninja and "ReVA Push" will be available in the Plugin menu.
+Press this to push data from BinaryNinja to ReVA, then follow the instructions in the [Workflow section](#workflow).
+The project name will be the name of the current open file.
+
+# Support
+
+Do you like my work? Want to support this project and others? Interested in how this project was designed and built?
+This project and many others are built live on my stream at https://twitch.tv/cyberkaida !

reverse-engineering-assistant-1.0.0/README.md

@@ -0,0 +1,133 @@
+# ReVA - Reverse Engineering Assistant
+
+[✨ A quick demo! ✨](https://asciinema.org/a/626197)
+
+The reverse engineering assistant (ReVA) is a project to build a disassembler agnostic AI assistant for
+reverse engineering tasks. This includes both _offline_ and online inference and a simple architecture.
+
+RevA is different to other efforts at building AI assistants for RE tasks because it uses a technique
+called [embedding](https://openai.com/blog/introducing-text-and-code-embeddings)
+to give the AI assistant a sort of "long term memory". The model also is given access to a number of tools
+that are tweaked to perform well with queries provided by the LLM. This allows the model to reason about the whole
+program, rather than just a single function. The tools are tweaked to lead the AI to examine deeper.
+
+Using this technique you can ask general questions and get relevant answers. The model prioritises
+information from the embeddings and tools, but when there is no information it can still respond to generic
+questions from its training.
+
+You can ask questions like:
+- Does this program use encryption?
+- Draw a class diagram using plantuml syntax.
+- Rename all the variables in main with descriptive names.
+- Explain the purpose of the `__mod_init` segment.
+- What does `mmap` return?
+- What does the function at address 0x80000 do?
+
+## Large Language Model Support
+
+RevA is based on [llama-index](https://github.com/jerryjliu/llama_index),
+which supports a number of models.
+
+Built in support is provided for:
+- [OpenAI](https://platform.openai.com/overview) for online inference and easy setup (Needs an OpenAI API key)
+- [Ollama](https://ollama.ai) and any model it supports for local on-device inference or connecting to a self hosted remote inference server.
+
+Limited support is provided for:
+- [llama-cpp](https://llama-cpp-python.readthedocs.io/en/latest/) and any model it supports for local on-device inference
+- [text-generation-webui](https://github.com/oobabooga/text-generation-webui) and any model it supports for self-hosted remote inference
+
+Adding additional inference servers is easy if it is supported by llama-index or langchain (on which llama-index is based).
+
+See the configuration section for more information about setting the model.
+
+## Configuration
+
+Configuration for the reverse engineering assistant is stored at
+`~/.config/reverse-engineering-assistant/config.yaml`. If this
+is not present on first start, a default configuration using
+OpenAI for inference and the `OPENAI_API_TOKEN` environment
+variable will be used.
+
+The most important setting is the `type` top level setting.
+This controls what inference service you use. These are the
+same as the configuration keys, for example to use Ollama,
+set type to `ollama` and configure the settings in the `ollama:`
+section.
+
+The configuration also contains the prompts used for the models.
+If you use Ollama or OpenAI these will be processed to fit the
+model specific prompt pattern (placing the system prompt in the
+correct tags, etc).
+
+For `llama-cpp` and `text-generation-webui` these may need to be
+configured for your specific model. For this reason Ollama is
+preferred for self hosting.
+
+## Workflow
+
+RevA has a two step workflow.
+1. Generate knowledge base
+2. Perform inference
+
+To generate the knowledge base, use the plugin for your disassembler and run the Assistant script.
+See [Ghidra Support](#ghidra-support) and [BinaryNinja Support](#binary-ninja-support) below.
+
+First your disassembler extracts the information required for the knowledge base and embeddings.
+This involes extracting each function, it's decompilation and some metadata. These are written to a "project". This allows
+multiple programs and data sources to be combined into one set of knowledge for the assistant. For example multiple malware
+samples, or a program and its libraries could be included along with previous RE notes.
+
+Projects are stored in `~/.cache/reverse-engineering-assistant/projects`. If you make significant changes to your
+annotations or analysis in your disassembler, you should delete and regenerate your project directory. This cache
+is a _snapshot_ of the state of your disassembler.
+
+To ask questions and run the inference a command line tool is provided. Run `revassistant --project ${NAME_OF_YOUR_PROJECT}` to begin the chat session.
+
+`revassistant` will hash the knowledge base and generate and combine the embeddings into a searchable
+index. Once this is complete the index is saved to disk and the chat session begins.
+
+## Installation
+
+To install the particular extension for your disassembler see:
+- [Ghidra Support](#ghidra-support)
+- [Binary Ninja Support](#binary-ninja-support)
+
+To install the chat component you can do the following:
+
+```sh
+python3 -m pip install ./reverse-engineering-assistant
+```
+
+The chat can be started with:
+
+```sh
+revassistant --project ${NAME_OF_YOUR_PROJECT}
+```
+
+
+# Ghidra Support
+
+## Prerequisites
+- [Ghidrathon](https://github.com/mandiant/Ghidrathon) >= 2.2.0 installed into Ghidra
+
+## Usage
+
+After installation, enable the [Ghidrathon extension](https://github.com/mandiant/Ghidrathon#installing-ghidrathon)
+and the Ghidra Assistant Extension.
+
+You can generate the knowledge base by running the Ghidra Assistant analysis from the Analysis menu in the Code Browser.
+
+# Binary Ninja Support
+
+Install the ReVA BinaryNinja plugin by opening your BinaryNinja plugin directory (Plugins -> Open Plugin Folder)
+and copying or symbolic linking the [binary-ninja-assistant](./binary-ninja-assistant) directory into the plugin
+directory.
+
+Restart Binary Ninja and "ReVA Push" will be available in the Plugin menu.
+Press this to push data from BinaryNinja to ReVA, then follow the instructions in the [Workflow section](#workflow).
+The project name will be the name of the current open file.
+
+# Support
+
+Do you like my work? Want to support this project and others? Interested in how this project was designed and built?
+This project and many others are built live on my stream at https://twitch.tv/cyberkaida !

reverse-engineering-assistant-1.0.0/pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["setuptools", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "reverse-engineering-assistant"
+readme = "README.md"
+version = "1.0.0"
+authors = [
+    {name="サイバーカイダ (cyberkaida)"},
+]
+description = "An AI assistant for reverse engineering tasks"
+classifiers = [
+    "License :: OSI Approved :: Apache Software License",
+    "Natural Language :: English",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python",
+]
+dependencies = [
+    "langchain",
+    "llama-cpp-python",
+    "prompt_toolkit",
+    "sentence_transformers",
+    "PyYAML",
+    "pydantic",
+    "rich",
+    "Flask",
+]
+
+[project.scripts]
+revassistant = "reverse_engineering_assistant.assistant:main"
+reva-serve = "reverse_engineering_assistant.assistant_api_server:main"
+reva-server = "reverse_engineering_assistant.assistant_api_server:main"
+

reverse-engineering-assistant-1.0.0/reverse_engineering_assistant/api_server_tools/function_tools.py

@@ -0,0 +1,122 @@
+
+from pathlib import Path
+from typing import Dict, List, Optional
+from ..assistant import AssistantProject, RevaTool, BaseLLM, register_tool
+from ..tool_protocol import RevaGetDecompilation, RevaGetDecompilationResponse, RevaGetFunctionCount, RevaGetFunctionCountResponse, RevaGetDefinedFunctionList, RevaGetDefinedFunctionListResponse
+
+from ..reva_exceptions import RevaToolException
+
+import logging
+
+
+@register_tool
+class RevaDecompilationIndex(RevaTool):
+    """
+    An index of decompiled functions available to the
+    reverse engineering assistant.
+    """
+    index_name = "decompilation"
+    description = "Used for retrieving decompiled functions"
+    logger = logging.getLogger("reverse_engineering_assistant.RevaDecompilationIndex")
+
+    def __init__(self, project: AssistantProject, llm: BaseLLM) -> None:
+        super().__init__(project, llm)
+        self.description = "Used for retrieveing decompiled functions"
+        self.tool_functions = [
+            self.get_decompilation_for_function,
+            self.get_defined_function_list_paginated,
+            self.get_defined_function_count,
+        ]
+    
+    def get_decompilation_for_function(self, function_name_or_address: str | int) -> Dict[str, str]:
+        """
+        Return the decompilation for the given function. The function can be specified by name or address.
+        Hint: It is too slow to decompile _all_ functions, so use get_defined_function_list_paginated to get a list of functions
+        and be sure to specify the function name or address exactly.
+        """
+        from ..assistant_api_server import RevaCallbackHandler, to_send_to_tool
+
+
+        # First normalise the argument
+        address: Optional[int] = None
+        name: Optional[str] = None
+        try:
+            address = int(function_name_or_address, 16)
+            if address <= 0:
+                raise RevaToolException("Address must be > 0 and in hex format")
+        except ValueError:
+            name = function_name_or_address
+        
+        if address is None and name is None:
+            raise RevaToolException("function_name_or_address must be an address or function name")
+
+        # Now we can ask the tool
+        get_decompilation_message = RevaGetDecompilation(address=address, function=name)
+        callback_handler = RevaCallbackHandler(self.project, get_decompilation_message)
+        to_send_to_tool.put(callback_handler)
+        self.logger.debug(f"Waiting for response to {get_decompilation_message.json()}")
+        response: RevaGetDecompilationResponse = callback_handler.wait()
+
+        if response.error_message:
+            raise RevaToolException(response.error_message, send_to_llm=True)
+
+        if not isinstance(response, RevaGetDecompilationResponse):
+            raise ValueError(f"Expected a RevaGetDecompilationResponse, got {response}")
+        
+        # Finally we can return the response
+        return {
+            "function": response.function,
+            "function_signature": response.function_signature,
+            "address": hex(response.address),
+            "decompilation": response.decompilation,
+            "variables": response.variables,
+        }
+
+                
+    def get_defined_function_list_paginated(self, page: int, page_size: int = 20) -> List[str]:
+        """
+        Return a paginated list of functions in the index. Use get_defined_function_count to get the total number of functions.
+        page is 1 indexed. To get the first page, set page to 1. Do not set page to 0.
+        """
+        from ..assistant_api_server import RevaCallbackHandler, to_send_to_tool
+
+        if isinstance(page, str):
+            page = int(page)
+        if isinstance(page_size, str):
+            page_size = int(page_size)
+        if page == 0:
+            raise ValueError("`page` is 1 indexed, page cannot be 0")
+        
+        get_function_list_message = RevaGetDefinedFunctionList(page=page, page_size=page_size)
+        callback_handler = RevaCallbackHandler(self.project, get_function_list_message)
+        to_send_to_tool.put(callback_handler)
+        
+        self.logger.debug(f"Waiting for response to {get_function_list_message.json()}")
+        response = callback_handler.wait()
+        if response.error_message:
+            raise RevaToolException(response.error_message, send_to_llm=True)
+
+        if not isinstance(response, RevaGetDefinedFunctionListResponse):
+            raise RevaToolException(f"Expected a RevaGetDefinedFunctionListResponse, got {response}")
+
+        return response.function_list
+    
+    def get_defined_function_count(self) -> int:
+        """
+        Return the total number of defined functions in the program.
+        """
+        from ..assistant_api_server import RevaCallbackHandler, to_send_to_tool
+
+        get_function_count_message = RevaGetFunctionCount()
+        callback_handler = RevaCallbackHandler(self.project, get_function_count_message)
+        to_send_to_tool.put(callback_handler)
+        self.logger.debug(f"Waiting for response to {get_function_count_message.json()}")
+        response = callback_handler.wait()
+
+        if response.error_message:
+            raise RevaToolException(response.error_message, send_to_llm=True)
+
+        if not isinstance(response, RevaGetFunctionCountResponse):
+            raise ValueError(f"Expected a RevaGetFunctionCountResponse, got {response}")
+
+        return response.function_count