├── .python-version ├── plugin ├── __init__.py └── lattice_server_plugin.py ├── requirements.txt ├── img └── lattice-logo.png ├── tests ├── data │ ├── test_binary │ └── test_program.c ├── conftest.py └── test_lattice_server.py ├── setup.py ├── LICENSE ├── mcp_server.py ├── README.md ├── lattice_client.py └── lib └── lattice.py /.python-version: -------------------------------------------------------------------------------- 1 | 3.13 2 | -------------------------------------------------------------------------------- /plugin/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Binary Ninja Lattice Server Plugin 3 | """ -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | mcp>=1.6.0 2 | requests>=2.32.3 3 | pytest>=7.4.0 4 | pytest-cov>=4.1.0 5 | -------------------------------------------------------------------------------- /img/lattice-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Invoke-RE/binja-lattice-mcp/HEAD/img/lattice-logo.png -------------------------------------------------------------------------------- /tests/data/test_binary: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Invoke-RE/binja-lattice-mcp/HEAD/tests/data/test_binary -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pytest 3 | 4 | @pytest.fixture(scope="session", autouse=True) 5 | def setup_test_env(): 6 | """Ensure Binary Ninja is in headless mode for testing""" 7 | os.environ["BN_DISABLE_UI"] = "1" 8 | yield -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name="binja-mcp-server", 5 | version="0.1.0", 6 | packages=find_packages(), 7 | install_requires=[ 8 | "mcp>=1.6.0", 9 | "requests>=2.32.3", 10 | ], 11 | extras_require={ 12 | "test": [ 13 | "pytest>=7.4.0", 14 | "pytest-cov>=4.1.0", 15 | ], 16 | }, 17 | ) -------------------------------------------------------------------------------- /tests/data/test_program.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | int global_var = 42; 4 | 5 | int add(int a, int b) { 6 | int result = a + b; 7 | return result; 8 | } 9 | 10 | void print_message(const char* message) { 11 | printf("%s\n", message); 12 | } 13 | 14 | int main() { 15 | int local_var = 10; 16 | int sum = add(local_var, global_var); 17 | print_message("Hello, World!"); 18 | return sum; 19 | } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2025 Invoke Reversing Inc. 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 8 | -------------------------------------------------------------------------------- /mcp_server.py: -------------------------------------------------------------------------------- 1 | from mcp.server.fastmcp import FastMCP 2 | from lib.lattice import Lattice 3 | import os, json 4 | 5 | # Initialize FastMCP server 6 | mcp = FastMCP("binja-lattice", log_level="ERROR") 7 | 8 | @mcp.tool() 9 | def get_all_function_names() -> str: 10 | """Get all function names""" 11 | response = lattice_client.get_all_function_names() 12 | if response and 'status' in response and response['status'] == 'success': 13 | return '\n'.join([f"{func['name']}" for func in response['function_names']]) 14 | return "Error: Could not retrieve function names" 15 | 16 | @mcp.tool() 17 | def get_binary_info() -> str: 18 | """Get information about the binary being analyzed""" 19 | response = lattice_client.get_binary_info() 20 | if response and 'status' in response and response['status'] == 'success': 21 | return json.dumps(response, indent=2) 22 | return "Error: Could not retrieve binary information" 23 | 24 | @mcp.tool() 25 | def update_function_name(name: str, new_name: str) -> str: 26 | """Update the name of a function""" 27 | response = lattice_client.update_function_name(name, new_name) 28 | if response and 'status' in response and response['status'] == 'success': 29 | return f"Successfully renamed function {name} to {new_name}" 30 | return f"Error: Could not update function name {name}" 31 | 32 | @mcp.tool() 33 | def add_comment_to_address(address: int, comment: str) -> str: 34 | """Add a comment to an address""" 35 | response = lattice_client.add_comment_to_address(address, comment) 36 | if response and 'status' in response and response['status'] == 'success': 37 | return f"Successfully added comment to address {address}" 38 | return f"Error: Could not add comment to address {address}" 39 | 40 | @mcp.tool() 41 | def add_comment_to_function(name: str, comment: str) -> str: 42 | """Add a comment to a function with specified function name""" 43 | response = lattice_client.add_comment_to_function(name, comment) 44 | if response and 'status' in response and response['status'] == 'success': 45 | return f"Successfully added comment to function {name}" 46 | return f"Error: Could not add comment to function {name}" 47 | 48 | @mcp.tool() 49 | def get_function_disassembly(name: str) -> str: 50 | """Get disassembly for the function""" 51 | response = lattice_client.get_function_disassembly(name) 52 | if response and 'status' in response and response['status'] == 'success': 53 | return '\n'.join([f"{block['address']}: {block['text']}" for block in response['disassembly']]) 54 | return f"Error: Could not retrieve function disassembly for function {name}" 55 | 56 | @mcp.tool() 57 | def get_function_pseudocode(name: str) -> str: 58 | """Get pseudocode for the function""" 59 | response = lattice_client.get_function_pseudocode(name) 60 | if response and 'status' in response and response['status'] == 'success': 61 | return '\n'.join([f"{block['address']}: {block['text']}" for block in response['pseudocode']]) 62 | return f"Error: Could not retrieve function pseudocode for function {name}" 63 | 64 | @mcp.tool() 65 | def get_function_variables(name: str) -> str: 66 | """Get variables for the function""" 67 | response = lattice_client.get_function_variables(name) 68 | if response and 'status' in response and response['status'] == 'success': 69 | rstr = 'Parameters: ' + '\n'.join([f"{param['name']}: {param['type']}" for param in response['variables']['parameters']]) \ 70 | + '\nLocal Variables: ' + '\n'.join([f"{var['name']}: {var['type']}" for var in response['variables']['local_variables']]) \ 71 | + '\nGlobal Variables: ' + '\n'.join([f"{var['name']}: {var['type']}" for var in response['variables']['global_variables']]) 72 | return rstr 73 | 74 | return f"Error: Could not retrieve function variables for function {name}" 75 | 76 | @mcp.tool() 77 | def update_variable_name(function_name: str, var_name: str, new_name: str) -> str: 78 | """Update the name of a variable""" 79 | response = lattice_client.update_variable_name(function_name, var_name, new_name) 80 | if response and 'status' in response and response['status'] == 'success': 81 | return f"Successfully renamed variable {var_name} to {new_name}" 82 | return f"Error: Could not update variable name {var_name}" 83 | 84 | @mcp.tool() 85 | def get_global_variable_data(function_name: str, global_var_name: str) -> str: 86 | """Get data pointed to by a global variable name""" 87 | response = lattice_client.get_global_variable_data(function_name, global_var_name) 88 | if response and 'status' in response and response['status'] == 'success': 89 | return response['message'] 90 | return f"Error: Could not retrieve global variable data for function {function_name} and variable {global_var_name}" 91 | 92 | @mcp.tool() 93 | def get_cross_references_to_function(name: str) -> str: 94 | """Get cross references to the specified function with function name""" 95 | response = lattice_client.get_cross_references_to_function(name) 96 | if response and 'status' in response and response['status'] == 'success': 97 | return '\n'.join([f"{ref['function']}" for ref in response['cross_references']]) 98 | return f"Error: Could not retrieve cross references for function {name}" 99 | 100 | # Initialize and run the server 101 | api_key = os.getenv("BNJLAT") 102 | if not api_key: 103 | raise ValueError("BNJLAT environment variable not set") 104 | 105 | global lattice_client 106 | lattice_client = Lattice() 107 | print(f"Authenticating with {api_key}") 108 | lattice_client.authenticate("mcp-user", api_key) 109 | mcp.run(transport='stdio') 110 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![BinjaLattice Logo](img/lattice-logo.png) 2 | 3 | # BinjaLattice 4 | 5 | BinjaLattice is a secure communication protocol for Binary Ninja that enables interaction with external Model Context Protocol (MCP) servers and tools. It provides a structured way to acquire information from Binary Ninja and the ability to modify an active Binary Ninja database over HTTP with a REST API. 6 | 7 | ## Demo 8 | 9 | [![BinjaLattice Demo](https://img.youtube.com/vi/xfDRVn0VIA0/0.jpg)](https://www.youtube.com/watch?v=xfDRVn0VIA0) 10 | 11 | ## Features 12 | 13 | - **Secure Authentication**: Token-based authentication system 14 | - **Encrypted Communication**: Optional SSL/TLS encryption 15 | - **Binary Analysis Context**: Export pseudocode, disassembly, variable names, binary information etc. 16 | - **Binary Modification**: Update function names, add comments, rename variables 17 | - **Token Management**: Automatic expiration and renewal of authentication tokens 18 | 19 | ## Installation 20 | 21 | 1. Copy `lattice_server_plugin.py` to your Binary Ninja plugins directory: 22 | - Linux: `~/.binaryninja/plugins/` 23 | - macOS: `~/Library/Application Support/Binary Ninja/plugins/` 24 | - Windows: `%APPDATA%\Binary Ninja\plugins\` 25 | 26 | 2. Create a virtual environment `pip -m venv venv-test` (or your preferred dependency manager) 27 | 28 | 3. Activate your virtual environment and install required Python dependencies: 29 | - Install with: `pip install -r requirements.txt` (or your preferred method) 30 | 31 | ## Usage 32 | 33 | ### Starting the Server in Binary Ninja 34 | 35 | 1. Open Binary Ninja and load a binary file 36 | 2. Go to `Plugins > Start Lattice Protocol Server` 37 | 3. The server will start and display the API key in the log console 38 | 4. Set the API key as the `BNJLAT` environment variable in your MCP configuration 39 | 40 | Example MCP configuration (`mcp.json`) from Cursor: 41 | ```json 42 | { 43 | "mcpServers": { 44 | "binja-lattice-mcp": { 45 | "command": "/path/to/venv/bin/python", 46 | "args": ["/path/to/mcp_server.py"], 47 | "env": { 48 | "BNJLAT": "your_api_key_here" 49 | } 50 | } 51 | } 52 | } 53 | ``` 54 | 55 | ### Available MCP Tools 56 | 57 | The following tools are available through the MCP server: 58 | 59 | - `get_all_function_names`: Get a list of all function names in the binary 60 | - `get_binary_info`: Get information about the binary being analyzed 61 | - `update_function_name`: Update the name of a function 62 | - `update_variable_name`: Change variable name within function to specified name 63 | - `get_global_variable_data`: Get data from global variable 64 | - `add_comment_to_address`: Add a comment to a specific address 65 | - `add_comment_to_function`: Add a comment to a function 66 | - `get_function_disassembly`: Get disassembly for a function 67 | - `get_function_pseudocode`: Get pseudocode for a function 68 | - `get_function_variables`: Get variables and parameters for a function 69 | - `get_cross_references_to_function`: Get cross references to a function 70 | 71 | ### Client Library Usage 72 | 73 | The `Lattice` client library provides a Python interface for interacting with the BinjaLattice server: 74 | 75 | ```python 76 | from lib.lattice import Lattice 77 | 78 | # Initialize client 79 | client = Lattice(host='localhost', port=9000, use_ssl=False) 80 | 81 | # Authenticate with API key 82 | client.authenticate("username", "API_KEY") 83 | 84 | # Example: Get binary information 85 | binary_info = client.get_binary_info() 86 | 87 | # Example: Update function name 88 | client.update_function_name("old_name", "new_name") 89 | 90 | # Example: Add comment to function 91 | client.add_comment_to_function("function_name", "This function handles authentication") 92 | ``` 93 | 94 | ### Command Line Interface 95 | 96 | The project includes `lattice_client.py`, which provides an interactive command-line interface for testing and debugging the BinjaLattice server: 97 | 98 | ```bash 99 | python lattice_client.py --host localhost --port 9000 [--ssl] --username user --password YOUR_API_KEY 100 | ``` 101 | 102 | #### Command Line Options 103 | 104 | - `--host`: Server host (default: localhost) 105 | - `--port`: Server port (default: 9000) 106 | - `--ssl`: Enable SSL/TLS encryption 107 | - `--interactive`, `-i`: Run in interactive mode 108 | - `--username`: Username for authentication 109 | - `--password`: Password/API key for authentication 110 | - `--token`: Authentication token (if you have one from previous authentication) 111 | 112 | #### Interactive Mode 113 | 114 | The interactive mode provides a menu-driven interface with the following options: 115 | 116 | 1. Get Binary Information 117 | 2. Get Function Context by Address 118 | 3. Get Function Context by Name 119 | 4. Update Function Name 120 | 5. Update Variable Name 121 | 6. Add Comment to Function 122 | 7. Add Comment to Address 123 | 8. Reconnect to Server 124 | 9. Get All Function Names 125 | 10. Get Function Disassembly 126 | 11. Get Function Pseudocode 127 | 12. Get Function Variables 128 | 13. Get Cross References to Function 129 | 14. Exit 130 | 131 | Example usage with interactive mode: 132 | 133 | ```bash 134 | python lattice_client.py -i --ssl --username user --password YOUR_API_KEY 135 | ``` 136 | 137 | #### Non-Interactive Commands 138 | 139 | You can also use the client to execute single commands: 140 | 141 | ```bash 142 | # Get binary information 143 | python lattice_client.py --username user --password YOUR_API_KEY --get-binary-info 144 | 145 | # Get function disassembly 146 | python lattice_client.py --username user --password YOUR_API_KEY --get-function-disassembly "main" 147 | 148 | # Add comment to a function 149 | python lattice_client.py --username user --password YOUR_API_KEY --add-comment-to-function "main" "Entry point of the program" 150 | ``` 151 | 152 | ### Security Notes 153 | 154 | - The API key is generated randomly on server start and shown in the Binary Ninja log 155 | - Tokens expire after 8 hours by default 156 | - SSL/TLS requires a certificate and key be provided by the user (disabled by default) 157 | - All requests require authentication via API key or token 158 | - The server runs locally by default on port 9000 159 | 160 | ## Development 161 | 162 | - The main server implementation is in `plugin/lattice_server_plugin.py` 163 | - MCP server implementation is in `mcp_server.py` 164 | - Client library is in `lib/lattice.py` 165 | 166 | ### Adding New Features 167 | 168 | To add new functionality: 169 | 170 | 1. Add new endpoint handlers in `LatticeRequestHandler` class in `lattice_server_plugin.py` 171 | 2. Add corresponding client methods in `Lattice` class in `lib/lattice.py` 172 | 3. Add new MCP tools in `mcp_server.py` 173 | 174 | ### Running Tests 175 | 176 | 1. Create a Python virtual environment and install the `requirements.txt` 177 | 2. Install the Binary Ninja Python API with the `install_api.py` provided in your Binary Ninja installation directory 178 | 3. Run the tests with `pytest tests/ -v` 179 | 180 | ## License 181 | 182 | [MIT License](LICENSE) 183 | -------------------------------------------------------------------------------- /tests/test_lattice_server.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import time 4 | import requests 5 | import pytest 6 | import binaryninja 7 | 8 | # Add the project root directory to Python path 9 | project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 10 | sys.path.insert(0, project_root) 11 | 12 | from plugin.lattice_server_plugin import BinjaLattice 13 | 14 | @pytest.fixture(scope="module") 15 | def bv(): 16 | """Create a headless BinaryView from our test binary""" 17 | bv = binaryninja.load("tests/data/test_binary") 18 | return bv 19 | 20 | @pytest.fixture(scope="module") 21 | def lattice_server(bv): 22 | """Start the Lattice server and yield its configuration""" 23 | protocol = BinjaLattice(bv, host="127.0.0.1", port=0) 24 | protocol.start_server() 25 | # Give it a moment to bind 26 | time.sleep(0.1) 27 | 28 | host, port = protocol.server.server_address 29 | base = f"http://{host}:{port}" 30 | api_key = protocol.auth_manager.api_key 31 | 32 | yield { 33 | "protocol": protocol, 34 | "base": base, 35 | "api_key": api_key, 36 | "bv": bv 37 | } 38 | 39 | # Teardown 40 | protocol.stop_server() 41 | # TODO figure out if this is actually possible 42 | # bv.close() 43 | 44 | def test_auth_fails_without_credentials(lattice_server): 45 | """Test that authentication fails without proper credentials""" 46 | r = requests.post(f"{lattice_server['base']}/auth", json={}) 47 | assert r.status_code == 401 48 | assert r.json()["status"] == "error" 49 | 50 | def test_auth_succeeds_with_valid_credentials(lattice_server): 51 | """Test that authentication succeeds with valid API key""" 52 | base = lattice_server["base"] 53 | api_key = lattice_server["api_key"] 54 | 55 | r = requests.post( 56 | f"{base}/auth", 57 | json={"username": "test_user", "password": api_key} 58 | ) 59 | assert r.status_code == 200 60 | j = r.json() 61 | assert j["status"] == "success" 62 | assert "token" in j 63 | return j["token"] # Return token for other tests 64 | 65 | def test_binary_info_endpoint(lattice_server): 66 | """Test the binary info endpoint""" 67 | base = lattice_server["base"] 68 | bv = lattice_server["bv"] 69 | token = test_auth_succeeds_with_valid_credentials(lattice_server) 70 | 71 | r = requests.get( 72 | f"{base}/binary/info", 73 | headers={"Authorization": f"Bearer {token}"} 74 | ) 75 | assert r.status_code == 200 76 | j = r.json() 77 | assert j["status"] == "success" 78 | info = j["binary_info"] 79 | 80 | # Assert against actual BinaryView properties 81 | assert info["filename"].endswith("test_binary") 82 | assert info["arch"] == bv.arch.name 83 | assert info["platform"] == bv.platform.name 84 | assert info["entry_point"] == bv.entry_point 85 | assert info["start"] == bv.start 86 | assert info["end"] == bv.end 87 | 88 | def test_function_operations(lattice_server): 89 | """Test function-related endpoints""" 90 | base = lattice_server["base"] 91 | bv = lattice_server["bv"] 92 | token = test_auth_succeeds_with_valid_credentials(lattice_server) 93 | headers = {"Authorization": f"Bearer {token}"} 94 | 95 | # Test getting all functions 96 | r = requests.get(f"{base}/functions", headers=headers) 97 | assert r.status_code == 200 98 | functions = r.json()["function_names"] 99 | assert len(functions) > 0 100 | 101 | # Test getting function context for main function 102 | main_func = next(f for f in functions if f["name"] == "_main") 103 | r = requests.get( 104 | f"{base}/functions/{main_func['address']}", 105 | headers=headers 106 | ) 107 | assert r.status_code == 200 108 | func_info = r.json()["function"] 109 | 110 | # Assert function properties match BinaryView 111 | actual_func = bv.get_function_at(main_func["address"]) 112 | assert func_info["name"] == actual_func.name 113 | assert func_info["start"] == actual_func.start 114 | assert func_info["end"] == actual_func.address_ranges[0].end 115 | 116 | def test_comment_operations(lattice_server): 117 | """Test comment-related endpoints""" 118 | base = lattice_server["base"] 119 | bv = lattice_server["bv"] 120 | token = test_auth_succeeds_with_valid_credentials(lattice_server) 121 | headers = {"Authorization": f"Bearer {token}"} 122 | 123 | # Test adding a comment to main function 124 | main_func = next(f for f in bv.functions if f.name == "_main") 125 | test_comment = "Test comment for main function" 126 | 127 | r = requests.post( 128 | f"{base}/functions/{main_func.name}/comments", 129 | headers=headers, 130 | json={"comment": test_comment} 131 | ) 132 | assert r.status_code == 200 133 | 134 | # Verify comment was added 135 | assert bv.get_comment_at(main_func.start) == test_comment 136 | 137 | def test_variable_operations(lattice_server): 138 | """Test variable-related endpoints""" 139 | base = lattice_server["base"] 140 | bv = lattice_server["bv"] 141 | token = test_auth_succeeds_with_valid_credentials(lattice_server) 142 | headers = {"Authorization": f"Bearer {token}"} 143 | 144 | # Test getting variables for main function 145 | main_func = next(f for f in bv.functions if f.name == "_main") 146 | r = requests.get( 147 | f"{base}/functions/{main_func.name}/variables", 148 | headers=headers 149 | ) 150 | assert r.status_code == 200 151 | vars_info = r.json()["variables"] 152 | 153 | # Assert variables match BinaryView 154 | assert len(vars_info["parameters"]) == len(main_func.parameter_vars) 155 | assert len(vars_info["local_variables"]) == len(main_func.vars) 156 | # This should be dynamic in case binja changes their pointer ref implementation etc. 157 | # Leaving it for now because getting function globals is a pain. 158 | assert len(vars_info["global_variables"]) == 4 159 | 160 | def test_variable_name_update(lattice_server): 161 | """Test variable name update endpoint""" 162 | base = lattice_server["base"] 163 | bv = lattice_server["bv"] 164 | token = test_auth_succeeds_with_valid_credentials(lattice_server) 165 | headers = {"Authorization": f"Bearer {token}"} 166 | 167 | # Test getting variables for main function 168 | main_func = next(f for f in bv.functions if f.name == "_main") 169 | r = requests.get( 170 | f"{base}/functions/{main_func.name}/variables", 171 | headers=headers 172 | ) 173 | assert r.status_code == 200 174 | vars_info = r.json()["variables"] 175 | var_name = vars_info["local_variables"][0]["name"] 176 | # Test updating variable name 177 | r = requests.put( 178 | f"{base}/variables/{main_func.name}/{var_name}/name", 179 | headers=headers, 180 | json={"name": "new_var_name"} 181 | ) 182 | assert r.status_code == 200 183 | r = requests.get( 184 | f"{base}/functions/{main_func.name}/variables", 185 | headers=headers 186 | ) 187 | assert r.status_code == 200 188 | vars_info = r.json()["variables"] 189 | print(vars_info) 190 | # This assumes that the variables maintain order after being renamed 191 | assert vars_info["local_variables"][0]["name"] == "new_var_name" 192 | 193 | def test_get_global_variable_data(lattice_server): 194 | """Test getting data for a global variable""" 195 | base = lattice_server["base"] 196 | bv = lattice_server["bv"] 197 | token = test_auth_succeeds_with_valid_credentials(lattice_server) 198 | headers = {"Authorization": f"Bearer {token}"} 199 | 200 | # This is a global variable in the main function 201 | # as per our returned convention because it is unnamed 202 | r = requests.get( 203 | f"{base}/global_variable_data/_main/data_100003f90", 204 | headers=headers 205 | ) 206 | assert r.status_code == 200 207 | data = r.json() 208 | assert data["status"] == "success" 209 | assert "Hello, World!" in data["message"] 210 | 211 | def test_cross_references(lattice_server): 212 | """Test cross-reference endpoints""" 213 | base = lattice_server["base"] 214 | bv = lattice_server["bv"] 215 | token = test_auth_succeeds_with_valid_credentials(lattice_server) 216 | headers = {"Authorization": f"Bearer {token}"} 217 | 218 | # Test getting cross references to print_message function 219 | print_func = next(f for f in bv.functions if f.name == "_print_message") 220 | r = requests.get( 221 | f"{base}/cross-references/{print_func.name}", 222 | headers=headers 223 | ) 224 | assert r.status_code == 200 225 | refs = r.json()["cross_references"] 226 | 227 | # Should have at least one reference from main 228 | assert len(refs) > 0 229 | assert any(ref["function"] == "_main" for ref in refs) 230 | 231 | def test_protected_endpoints_require_auth(lattice_server): 232 | """Test that protected endpoints require authentication""" 233 | base = lattice_server["base"] 234 | endpoints = [ 235 | "/binary/info", 236 | "/functions", 237 | "/functions/main", 238 | "/functions/main/variables" 239 | ] 240 | 241 | for endpoint in endpoints: 242 | r = requests.get(f"{base}{endpoint}") 243 | assert r.status_code == 401 244 | assert r.json()["status"] == "error" -------------------------------------------------------------------------------- /lattice_client.py: -------------------------------------------------------------------------------- 1 | from lib.lattice import Lattice 2 | import argparse, sys, json, logging 3 | 4 | logger = logging.getLogger(__name__) 5 | 6 | def print_menu(): 7 | """Print the interactive menu""" 8 | print("\nBinjaLattice Client Menu:") 9 | print("1. Get Binary Information") 10 | print("2. Get Function Context by Address") 11 | print("3. Get Function Context by Name") 12 | print("4. Update Function Name") 13 | print("5. Update Variable Name") 14 | print("6. Add Comment to Function") 15 | print("7. Add Comment to Address") 16 | print("8. Reconnect to Server") 17 | print("9. Get All Function Names") 18 | print("10. Get Function Disassembly") 19 | print("11. Get Function Pseudocode") 20 | print("12. Get Function Variables") 21 | print("13. Get Cross References to Function") 22 | print("14. Get Global Variable Data") 23 | print("15. Exit") 24 | print() 25 | 26 | def interactive_mode(client: Lattice): 27 | """Run the interactive REPL mode""" 28 | while True: 29 | print_menu() 30 | try: 31 | choice = input("Enter your choice (1-15): ").strip() 32 | 33 | if choice == '1': 34 | result = client.get_binary_info() 35 | print(json.dumps(result, indent=2)) 36 | 37 | elif choice == '2': 38 | addr = input("Enter function address (hex or decimal): ").strip() 39 | try: 40 | address = int(addr, 0) 41 | result = client.get_function_context(address) 42 | print(json.dumps(result, indent=2)) 43 | except ValueError: 44 | logger.error("Invalid address format") 45 | elif choice == '3': 46 | name = input("Enter function name: ").strip() 47 | try: 48 | result = client.get_function_context_by_name(name) 49 | print(json.dumps(result, indent=2)) 50 | except ValueError: 51 | logger.error("Invalid function name") 52 | 53 | elif choice == '4': 54 | name = input("Enter function name: ").strip() 55 | new_name = input("Enter new function name: ").strip() 56 | try: 57 | result = client.update_function_name(name, new_name) 58 | print(json.dumps(result, indent=2)) 59 | except ValueError: 60 | logger.error("Invalid function name") 61 | 62 | elif choice == '5': 63 | func_name = input("Enter function name: ").strip() 64 | var_name = input("Enter variable name: ").strip() 65 | new_name = input("Enter new variable name: ").strip() 66 | try: 67 | result = client.update_variable_name(func_name, var_name, new_name) 68 | print(json.dumps(result, indent=2)) 69 | except ValueError: 70 | logger.error("Invalid input format") 71 | 72 | elif choice == '6': 73 | name = input("Enter function name: ").strip() 74 | comment = input("Enter comment: ").strip() 75 | try: 76 | result = client.add_comment_to_function(name, comment) 77 | print(json.dumps(result, indent=2)) 78 | except ValueError: 79 | logger.error("Invalid function name") 80 | elif choice == '7': 81 | address = input("Enter address (hex or decimal): ").strip() 82 | comment = input("Enter comment: ").strip() 83 | try: 84 | result = client.add_comment_to_address(address, comment) 85 | logger.error(json.dumps(result, indent=2)) 86 | except ValueError: 87 | print("Invalid address format") 88 | elif choice == '8': 89 | client.close() 90 | if client.connect(): 91 | print("Reconnected successfully") 92 | if client.auth_token: 93 | print("Previous authentication token is still valid") 94 | else: 95 | logger.error("Failed to reconnect") 96 | elif choice == '9': 97 | result = client.get_all_function_names() 98 | print(json.dumps(result, indent=2)) 99 | elif choice == '10': 100 | name = input("Enter function name: ").strip() 101 | try: 102 | result = client.get_function_disassembly(name) 103 | print(json.dumps(result, indent=2)) 104 | except ValueError: 105 | logger.error("Invalid function name") 106 | elif choice == '11': 107 | name = input("Enter function name: ").strip() 108 | try: 109 | result = client.get_function_pseudocode(name) 110 | print(json.dumps(result, indent=2)) 111 | except ValueError: 112 | logger.error("Invalid function name") 113 | elif choice == '12': 114 | name = input("Enter function name: ").strip() 115 | try: 116 | result = client.get_function_variables(name) 117 | print(json.dumps(result, indent=2)) 118 | except ValueError: 119 | logger.error("Invalid function name") 120 | elif choice == '13': 121 | name = input("Enter function name: ").strip() 122 | try: 123 | result = client.get_cross_references_to_function(name) 124 | print(json.dumps(result, indent=2)) 125 | except ValueError: 126 | logger.error("Invalid function name") 127 | elif choice == '14': 128 | func_name = input("Enter function name: ").strip() 129 | var_name = input("Enter variable name: ").strip() 130 | try: 131 | result = client.get_global_variable_data(func_name, var_name) 132 | print(json.dumps(result, indent=2)) 133 | except ValueError: 134 | logger.error("Invalid function name") 135 | elif choice == '15': 136 | print("Goodbye!") 137 | break 138 | else: 139 | print("Invalid choice. Please try again.") 140 | 141 | except KeyboardInterrupt: 142 | print("\nGoodbye!") 143 | break 144 | except Exception as e: 145 | logger.error(f"Error: {e}") 146 | print("Try reconnecting to the server (option 8)") 147 | 148 | def main(): 149 | parser = argparse.ArgumentParser(description='BinjaLattice Client - Communicate with Binary Ninja Lattice Protocol Server') 150 | parser.add_argument('--host', default='localhost', help='Server host (default: localhost)') 151 | parser.add_argument('--port', type=int, default=9000, help='Server port (default: 9000)') 152 | parser.add_argument('--ssl', action='store_true', help='Use SSL/TLS encryption') 153 | parser.add_argument('--interactive', '-i', action='store_true', help='Run in interactive mode') 154 | 155 | # Authentication options 156 | auth_group = parser.add_argument_group('Authentication') 157 | auth_group.add_argument('--username', help='Username for authentication') 158 | auth_group.add_argument('--password', help='Password/API key for authentication') 159 | auth_group.add_argument('--token', help='Authentication token') 160 | 161 | # Command options (only used in non-interactive mode) 162 | command_group = parser.add_argument_group('Commands') 163 | command_group.add_argument('--get-binary-info', action='store_true', help='Get binary information') 164 | command_group.add_argument('--get-function-context', type=lambda x: int(x, 0), help='Get function context at address (hex or decimal)') 165 | command_group.add_argument('--get-basic-block-context', type=lambda x: int(x, 0), help='Get basic block context at address (hex or decimal)') 166 | command_group.add_argument('--update-function-name', nargs=2, help='Update function name:
') 167 | command_group.add_argument('--update-variable-name', nargs=3, help='Update variable name: ') 168 | command_group.add_argument('--add-comment-to-address', nargs=2, help='Add comment to address:
') 169 | command_group.add_argument('--add-comment-to-function', nargs=2, help='Add comment to function: ') 170 | command_group.add_argument('--get-function-disassembly', type=str, help='Get function disassembly for function name') 171 | command_group.add_argument('--get-function-pseudocode', type=str, help='Get function pseudocode for function name') 172 | command_group.add_argument('--get-function-variables', type=str, help='Get function variables for function name') 173 | command_group.add_argument('--get-cross-references-to-function', type=str, help='Get cross references to function name') 174 | args = parser.parse_args() 175 | 176 | # Create client 177 | client = Lattice(host=args.host, port=args.port, use_ssl=args.ssl) 178 | 179 | # Authenticate 180 | if args.token: 181 | if not client.authenticate_with_token(args.token): 182 | print("Authentication failed with token") 183 | client.close() 184 | sys.exit(1) 185 | elif args.username and args.password: 186 | if not client.authenticate(args.username, args.password): 187 | print("Authentication failed with username/password") 188 | client.close() 189 | sys.exit(1) 190 | else: 191 | print("Authentication credentials required (--token or --username/--password)") 192 | client.close() 193 | sys.exit(1) 194 | 195 | try: 196 | if args.interactive: 197 | interactive_mode(client) 198 | else: 199 | # Execute requested command 200 | if args.get_binary_info: 201 | result = client.get_binary_info() 202 | print(json.dumps(result, indent=2)) 203 | 204 | elif args.get_function_context: 205 | result = client.get_function_context(args.get_function_context) 206 | print(json.dumps(result, indent=2)) 207 | elif args.update_function_name: 208 | address = int(args.update_function_name[0], 0) 209 | new_name = args.update_function_name[1] 210 | result = client.update_function_name(address, new_name) 211 | print(json.dumps(result, indent=2)) 212 | 213 | elif args.update_variable_name: 214 | func_addr = int(args.update_variable_name[0], 0) 215 | var_id = int(args.update_variable_name[1]) 216 | new_name = args.update_variable_name[2] 217 | result = client.update_variable_name(func_addr, var_id, new_name) 218 | print(json.dumps(result, indent=2)) 219 | 220 | elif args.add_comment_to_address: 221 | address = int(args.add_comment_to_address[0], 0) 222 | comment = args.add_comment_to_address[1] 223 | result = client.add_comment_to_address(address, comment) 224 | print(json.dumps(result, indent=2)) 225 | 226 | elif args.add_comment_to_function: 227 | name = args.add_comment_to_function[0] 228 | comment = args.add_comment_to_function[1] 229 | result = client.add_comment_to_function(name, comment) 230 | 231 | print(json.dumps(result, indent=2)) 232 | elif args.get_function_disassembly: 233 | result = client.get_function_disassembly(args.get_function_disassembly) 234 | print(json.dumps(result, indent=2)) 235 | 236 | elif args.get_function_pseudocode: 237 | result = client.get_function_pseudocode(args.get_function_pseudocode) 238 | print(json.dumps(result, indent=2)) 239 | 240 | elif args.get_function_variables: 241 | result = client.get_function_variables(args.get_function_variables) 242 | print(json.dumps(result, indent=2)) 243 | 244 | elif args.get_cross_references_to_function: 245 | result = client.get_cross_references_to_function(args.get_cross_references_to_function) 246 | print(json.dumps(result, indent=2)) 247 | else: 248 | print("No command specified. Use --help to see available commands.") 249 | 250 | finally: 251 | client.close() 252 | 253 | if __name__ == "__main__": 254 | main() 255 | -------------------------------------------------------------------------------- /lib/lattice.py: -------------------------------------------------------------------------------- 1 | import requests 2 | import json 3 | import logging 4 | import sys 5 | from typing import Optional, Dict, Any, List, Tuple, Union 6 | from urllib.parse import urljoin 7 | 8 | logging.basicConfig(level=logging.INFO) 9 | logger = logging.getLogger(__name__) 10 | 11 | class Lattice: 12 | """Client for communicating with a BinjaLattice server""" 13 | 14 | def __init__(self, host: str = "localhost", port: int = 9000, use_ssl: bool = False): 15 | """ 16 | Initialize the client. 17 | 18 | Args: 19 | host: Host address of the server 20 | port: Port number of the server 21 | use_ssl: Whether to use SSL/TLS encryption 22 | """ 23 | self.host = host 24 | self.port = port 25 | self.use_ssl = use_ssl 26 | self.auth_token = None 27 | self.base_url = f"{'https' if use_ssl else 'http'}://{host}:{port}" 28 | self.session = requests.Session() 29 | if not use_ssl: 30 | self.session.verify = False # Disable SSL verification for non-SSL connections 31 | 32 | def connect(self) -> bool: 33 | """Connect to the server""" 34 | #try: 35 | response = self.session.get(urljoin(self.base_url, '/binary/info')) 36 | if response.status_code == 200: 37 | logger.info(f"Connected to {self.host}:{self.port}") 38 | return True 39 | elif response.status_code == 401: 40 | logger.error(f"Authentication failed with status code: {response.status_code}") 41 | logger.error(f"Response: {response.text}") 42 | return False 43 | else: 44 | logger.error(f"Failed to connect: {response.status_code}") 45 | return False 46 | #except Exception as e: 47 | # logger.error(f"Failed to connect: {e}") 48 | # return False 49 | 50 | def authenticate(self, username: str, password: str) -> bool: 51 | """ 52 | Authenticate with the server using username/password 53 | 54 | Args: 55 | username: Username for authentication 56 | password: Password (API key) for authentication 57 | 58 | Returns: 59 | True if authentication successful, False otherwise 60 | """ 61 | response = self.session.post( 62 | urljoin(self.base_url, '/auth'), 63 | json={ 64 | 'username': username, 65 | 'password': password 66 | } 67 | ) 68 | 69 | if response.status_code == 200: 70 | print(response.content) 71 | data = json.loads(response.content) 72 | if data.get('status') == 'success': 73 | self.auth_token = data.get('token') 74 | self.session.headers.update({'Authorization': f'Bearer {self.auth_token}'}) 75 | logger.info("Authentication successful") 76 | return True 77 | else: 78 | logger.error(f"Authentication failed: {data.get('message')}") 79 | else: 80 | logger.error(f"Authentication failed with status code: {response.status_code}") 81 | 82 | return False 83 | 84 | def authenticate_with_token(self, token: str) -> bool: 85 | """ 86 | Authenticate with the server using a token 87 | 88 | Args: 89 | token: Authentication token 90 | 91 | Returns: 92 | True if authentication successful, False otherwise 93 | """ 94 | try: 95 | response = self.session.post( 96 | urljoin(self.base_url, '/auth'), 97 | json={'token': token} 98 | ) 99 | 100 | if response.status_code == 200: 101 | data = response.json() 102 | if data.get('status') == 'success': 103 | self.auth_token = token 104 | self.session.headers.update({'Authorization': f'Bearer {self.auth_token}'}) 105 | logger.info("Token authentication successful") 106 | return True 107 | else: 108 | logger.error(f"Token authentication failed: {data.get('message')}") 109 | else: 110 | logger.error(f"Token authentication failed with status code: {response.status_code}") 111 | 112 | return False 113 | 114 | except Exception as e: 115 | logger.error(f"Token authentication error: {e}") 116 | return False 117 | 118 | def get_binary_info(self) -> Optional[Dict[str, Any]]: 119 | """Get information about the binary""" 120 | try: 121 | response = self.session.get(urljoin(self.base_url, '/binary/info')) 122 | if response.status_code == 200: 123 | return response.json() 124 | return None 125 | except Exception as e: 126 | logger.error(f"Error getting binary info: {e}") 127 | return None 128 | 129 | def get_function_context(self, address: int) -> Optional[Dict[str, Any]]: 130 | """ 131 | Get context for a function at the specified address 132 | 133 | Args: 134 | address: Address of the function 135 | 136 | Returns: 137 | Dictionary containing function context 138 | """ 139 | try: 140 | response = self.session.get(urljoin(self.base_url, f'/functions/{address}')) 141 | if response.status_code == 200: 142 | return response.json() 143 | return {'status': 'error', 'message': 'Failed to get function context'} 144 | except Exception as e: 145 | logger.error(f"Error getting function context: {e}") 146 | return {'status': 'error', 'message': str(e)} 147 | 148 | def get_function_context_by_name(self, name: str) -> Optional[Dict[str, Any]]: 149 | """ 150 | Get context for a function by name 151 | 152 | Args: 153 | name: Name of the function 154 | 155 | Returns: 156 | Dictionary containing function context 157 | """ 158 | try: 159 | response = self.session.get(urljoin(self.base_url, f'/functions/name/{name}')) 160 | if response.status_code == 200: 161 | return response.json() 162 | return {'status': 'error', 'message': 'Failed to get function context by name'} 163 | except Exception as e: 164 | logger.error(f"Error getting function context by name: {e}") 165 | return {'status': 'error', 'message': str(e)} 166 | 167 | def get_all_function_names(self) -> Optional[Dict[str, Any]]: 168 | """ 169 | Get all function names 170 | """ 171 | try: 172 | response = self.session.get(urljoin(self.base_url, '/functions')) 173 | if response.status_code == 200: 174 | return response.json() 175 | return {'status': 'error', 'message': 'Failed to get all function names'} 176 | except Exception as e: 177 | logger.error(f"Error getting all function names: {e}") 178 | return {'status': 'error', 'message': str(e)} 179 | 180 | def update_function_name(self, name: str, new_name: str) -> Optional[Dict[str, Any]]: 181 | """ 182 | Update the name of a function 183 | 184 | Args: 185 | name: Current name of the function 186 | new_name: New name for the function 187 | 188 | Returns: 189 | Dictionary containing the result of the operation 190 | """ 191 | try: 192 | response = self.session.put( 193 | urljoin(self.base_url, f'/functions/{name}/name'), 194 | json={'name': new_name} 195 | ) 196 | if response.status_code == 200: 197 | return response.json() 198 | return {'status': 'error', 'message': 'Failed to update function name'} 199 | except Exception as e: 200 | logger.error(f"Error updating function name: {e}") 201 | return {'status': 'error', 'message': str(e)} 202 | 203 | def update_variable_name(self, function_name: str, var_name: str, new_name: str) -> Optional[Dict[str, Any]]: 204 | """ 205 | Update the name of a variable in a function 206 | 207 | Args: 208 | function_name: Name of the function containing the variable 209 | var_name: Name of the variable to rename 210 | new_name: New name for the variable 211 | 212 | Returns: 213 | Dictionary containing the result of the operation 214 | """ 215 | try: 216 | response = self.session.put( 217 | urljoin(self.base_url, f'/variables/{function_name}/{var_name}/name'), 218 | json={'name': new_name} 219 | ) 220 | if response.status_code == 200: 221 | return response.json() 222 | return {'status': 'error', 'message': 'Failed to update variable name'} 223 | except Exception as e: 224 | logger.error(f"Error updating variable name: {e}") 225 | return {'status': 'error', 'message': str(e)} 226 | 227 | def get_global_variable_data(self, function_name: str, global_var_name: str) -> Optional[Dict[str, Any]]: 228 | """ 229 | Get data for a global variable 230 | """ 231 | try: 232 | response = self.session.get(urljoin(self.base_url, f'/global_variable_data/{function_name}/{global_var_name}')) 233 | if response.status_code == 200: 234 | return response.json() 235 | return {'status': 'error', 'message': 'Failed to get global variable data'} 236 | except Exception as e: 237 | logger.error(f"Error getting global variable data: {e}") 238 | return {'status': 'error', 'message': str(e)} 239 | 240 | def add_comment_to_address(self, address: int, comment: str) -> Optional[Dict[str, Any]]: 241 | """ 242 | Add a comment at the specified address 243 | 244 | Args: 245 | address: Address to add the comment at 246 | comment: Comment text to add 247 | 248 | Returns: 249 | Dictionary containing the result of the operation 250 | """ 251 | try: 252 | response = self.session.post( 253 | urljoin(self.base_url, f'/comments/{address}'), 254 | json={'comment': comment} 255 | ) 256 | if response.status_code == 200: 257 | return response.json() 258 | return {'status': 'error', 'message': 'Failed to add comment'} 259 | except Exception as e: 260 | logger.error(f"Error adding comment: {e}") 261 | return {'status': 'error', 'message': str(e)} 262 | 263 | def add_comment_to_function(self, name: str, comment: str) -> Optional[Dict[str, Any]]: 264 | """ 265 | Add a comment to a function with specified function name 266 | """ 267 | try: 268 | response = self.session.post( 269 | urljoin(self.base_url, f'/functions/{name}/comments'), 270 | json={'comment': comment} 271 | ) 272 | if response.status_code == 200: 273 | return response.json() 274 | return {'status': 'error', 'message': 'Failed to add comment'} 275 | except Exception as e: 276 | logger.error(f"Error adding comment to function: {e}") 277 | return {'status': 'error', 'message': str(e)} 278 | 279 | def get_function_disassembly(self, name: str) -> Optional[Dict[str, Any]]: 280 | """ 281 | Get disassembly for a function with specified function name 282 | 283 | Args: 284 | name: Address of the function 285 | 286 | Returns: 287 | Dictionary containing function disassembly 288 | """ 289 | try: 290 | response = self.session.get(urljoin(self.base_url, f'/functions/{name}/disassembly')) 291 | if response.status_code == 200: 292 | return response.json() 293 | return {'status': 'error', 'message': 'Failed to get function disassembly'} 294 | except Exception as e: 295 | logger.error(f"Error getting function disassembly: {e}") 296 | return {'status': 'error', 'message': str(e)} 297 | 298 | def get_cross_references_to_function(self, name: str) -> Optional[Dict[str, Any]]: 299 | """ 300 | Get cross references to a function 301 | """ 302 | try: 303 | response = self.session.get(urljoin(self.base_url, f'/cross-references/{name}')) 304 | if response.status_code == 200: 305 | return response.json() 306 | return {'status': 'error', 'message': 'Failed to get cross references to function'} 307 | except Exception as e: 308 | logger.error(f"Error getting cross references to function: {e}") 309 | return {'status': 'error', 'message': str(e)} 310 | 311 | def get_function_pseudocode(self, name: str) -> Optional[Dict[str, Any]]: 312 | """ 313 | Get pseudocode for a function with specified function name 314 | 315 | Args: 316 | name: Name of the function 317 | 318 | Returns: 319 | Dictionary containing function pseudocode 320 | """ 321 | try: 322 | response = self.session.get(urljoin(self.base_url, f'/functions/{name}/pseudocode')) 323 | if response.status_code == 200: 324 | return response.json() 325 | return {'status': 'error', 'message': 'Failed to get function pseudocode'} 326 | except Exception as e: 327 | logger.error(f"Error getting function pseudocode: {e}") 328 | return {'status': 'error', 'message': str(e)} 329 | 330 | def get_function_variables(self, name: str) -> Optional[Dict[str, Any]]: 331 | """ 332 | Get variables for a function at the specified address 333 | 334 | Args: 335 | name: Name of function 336 | 337 | Returns: 338 | Dictionary containing function variables 339 | """ 340 | try: 341 | response = self.session.get(urljoin(self.base_url, f'/functions/{name}/variables')) 342 | if response.status_code == 200: 343 | return response.json() 344 | return {'status': 'error', 'message': 'Failed to get function variables'} 345 | except Exception as e: 346 | logger.error(f"Error getting function variables: {e}") 347 | return {'status': 'error', 'message': str(e)} 348 | 349 | def close(self): 350 | """Close the connection to the server""" 351 | self.session.close() 352 | logger.info("Connection closed") 353 | -------------------------------------------------------------------------------- /plugin/lattice_server_plugin.py: -------------------------------------------------------------------------------- 1 | from binaryninja import * 2 | from binaryninja.binaryview import BinaryView 3 | from binaryninja.enums import DisassemblyOption 4 | from binaryninja.function import DisassemblySettings, Function 5 | from binaryninja.lineardisassembly import LinearViewCursor, LinearViewObject 6 | from binaryninja.plugin import PluginCommand 7 | from binaryninja.log import Logger 8 | from typing import Optional, Dict, Any, List, Tuple 9 | from http.server import HTTPServer, BaseHTTPRequestHandler 10 | from urllib.parse import urlparse 11 | import json 12 | import os 13 | import secrets 14 | import time 15 | import ssl 16 | import re 17 | import traceback 18 | import threading 19 | 20 | logger = Logger(session_id=0, logger_name=__name__) 21 | 22 | class AuthManager: 23 | """Manages authentication for the Lattice Protocol""" 24 | def __init__(self, token_expiry_seconds=28800): 25 | """ 26 | Initialize the authentication manager 27 | 28 | Args: 29 | token_expiry_seconds: How long tokens are valid (default: 1 hour) 30 | """ 31 | self.token_expiry_seconds = token_expiry_seconds 32 | self.tokens = {} # Map of token -> (expiry_time, client_info) 33 | 34 | # Generate a secure API key on startup 35 | self.api_key = secrets.token_hex(16) 36 | logger.log_info(f"API key: {self.api_key}") 37 | 38 | def generate_token(self, client_info: Dict[str, Any]) -> str: 39 | """ 40 | Generate a new authentication token 41 | 42 | Args: 43 | client_info: Information about the client requesting the token 44 | 45 | Returns: 46 | A new authentication token 47 | """ 48 | token = secrets.token_hex(16) 49 | expiry = time.time() + self.token_expiry_seconds 50 | self.tokens[token] = (expiry, client_info) 51 | 52 | # Cleanup expired tokens 53 | self._cleanup_expired_tokens() 54 | 55 | return token 56 | 57 | def validate_token(self, token: str) -> Tuple[bool, Optional[Dict[str, Any]]]: 58 | """ 59 | Validate an authentication token 60 | 61 | Args: 62 | token: The token to validate 63 | 64 | Returns: 65 | Tuple of (is_valid, client_info) 66 | """ 67 | logger.log_info(f"Validating token: {token}") 68 | if token not in self.tokens: 69 | return False, None 70 | 71 | expiry, client_info = self.tokens[token] 72 | 73 | if time.time() > expiry: 74 | # Token has expired 75 | del self.tokens[token] 76 | return False, None 77 | 78 | return True, client_info 79 | 80 | def revoke_token(self, token: str) -> bool: 81 | """ 82 | Revoke a token 83 | 84 | Args: 85 | token: The token to revoke 86 | 87 | Returns: 88 | True if the token was revoked, False if it didn't exist 89 | """ 90 | if token in self.tokens: 91 | del self.tokens[token] 92 | return True 93 | return False 94 | 95 | def _cleanup_expired_tokens(self): 96 | """Remove expired tokens from the tokens dictionary""" 97 | current_time = time.time() 98 | expired_tokens = [ 99 | token for token, (expiry, _) in self.tokens.items() 100 | if current_time > expiry 101 | ] 102 | 103 | for token in expired_tokens: 104 | del self.tokens[token] 105 | 106 | def verify_credentials(self, password: str) -> bool: 107 | """ 108 | Verify a username and password against stored credentials. 109 | For simplicity, this just verifies against the API key. 110 | In a real implementation, this would check against a secure credential store. 111 | 112 | Args: 113 | username: The username to tie to session token 114 | password: The password to verify 115 | 116 | Returns: 117 | True if the credentials are valid, False otherwise 118 | """ 119 | # For simplicity, we're using the API key as the "password" 120 | # In a real implementation, this would use secure password hashing 121 | return password == self.api_key 122 | 123 | class LatticeRequestHandler(BaseHTTPRequestHandler): 124 | """HTTP request handler for the Lattice Protocol""" 125 | 126 | def __init__(self, *args, **kwargs): 127 | self.protocol = kwargs.pop('protocol') 128 | super().__init__(*args, **kwargs) 129 | 130 | def _send_response(self, data: Dict[str, Any], status: int = 200): 131 | """Send JSON response""" 132 | self.send_response(status) 133 | self.send_header('Content-type', 'application/json') 134 | self.end_headers() 135 | self.wfile.write(json.dumps(data).encode()) 136 | 137 | def _require_auth(self, handler): 138 | """Decorator to require authentication""" 139 | def decorated(*args, **kwargs): 140 | auth_header = self.headers.get('Authorization') 141 | if not auth_header: 142 | self._send_response({'status': 'error', 'message': 'No token provided'}, 401) 143 | return 144 | 145 | # Remove 'Bearer ' prefix if present 146 | token = auth_header[7:] if auth_header.startswith('Bearer ') else auth_header 147 | 148 | is_valid, client_info = self.protocol.auth_manager.validate_token(token) 149 | if not is_valid: 150 | self._send_response({'status': 'error', 'message': 'Invalid token'}, 401) 151 | return 152 | 153 | return handler(*args, **kwargs) 154 | return decorated 155 | 156 | def do_POST(self): 157 | """Handle POST requests""" 158 | parsed_path = urlparse(self.path) 159 | path = parsed_path.path 160 | 161 | try: 162 | content_length = int(self.headers.get('Content-Length', 0)) 163 | body = self.rfile.read(content_length) 164 | data = json.loads(body.decode()) 165 | except Exception as e: 166 | self._send_response({'status': 'error', 'message': str(e)}, 400) 167 | return 168 | 169 | if path == '/auth': 170 | self._handle_auth(data) 171 | elif path.startswith('/comments/'): 172 | self._require_auth(self._handle_add_comment_to_address)(data) 173 | elif path.startswith('/functions/'): 174 | logger.log_info(f"Handling add comment to function request: {data}") 175 | self._require_auth(self._handle_add_comment_to_function)(data) 176 | else: 177 | self._send_response({'status': 'error', 'message': 'Invalid endpoint'}, 404) 178 | 179 | def do_PUT(self): 180 | """Handle PUT requests""" 181 | parsed_path = urlparse(self.path) 182 | path = parsed_path.path 183 | 184 | try: 185 | content_length = int(self.headers.get('Content-Length', 0)) 186 | body = self.rfile.read(content_length) 187 | data = json.loads(body.decode()) 188 | except Exception as e: 189 | self._send_response({'status': 'error', 'message': str(e)}, 400) 190 | return 191 | 192 | if path.startswith('/functions/') and path.endswith('/name'): 193 | self._require_auth(self._handle_update_function_name)(data) 194 | elif path.startswith('/variables/') and path.endswith('/name'): 195 | self._require_auth(self._handle_update_variable_name)(data) 196 | else: 197 | self._send_response({'status': 'error', 'message': 'Invalid endpoint'}, 404) 198 | 199 | def do_GET(self): 200 | """Handle GET requests""" 201 | parsed_path = urlparse(self.path) 202 | path = parsed_path.path 203 | 204 | if path == '/binary/info': 205 | self._require_auth(self._handle_get_binary_info)() 206 | elif path == '/functions': 207 | self._require_auth(self._handle_get_all_function_names)() 208 | elif path.startswith('/functions/'): 209 | if path.startswith('/functions/name/'): 210 | self._require_auth(self._handle_get_function_context_by_name)() 211 | elif path.endswith('/disassembly'): 212 | self._require_auth(self._handle_get_function_disassembly)() 213 | elif path.endswith('/pseudocode'): 214 | self._require_auth(self._handle_get_function_pseudocode)() 215 | elif path.endswith('/variables'): 216 | self._require_auth(self._handle_get_function_variables)() 217 | else: 218 | self._require_auth(self._handle_get_function_context_by_address)() 219 | elif path.startswith('/global_variable_data'): 220 | self._require_auth(self._handle_get_global_variable_data)() 221 | elif path.startswith('/cross-references/'): 222 | self._require_auth(self._handle_get_cross_references_to_function)() 223 | else: 224 | self._send_response({'status': 'error', 'message': 'Invalid endpoint'}, 404) 225 | 226 | def _handle_auth(self, data): 227 | """Handle authentication requests""" 228 | username = data.get('username') 229 | password = data.get('password') 230 | token = data.get('token') 231 | 232 | if token: 233 | is_valid, client_info = self.protocol.auth_manager.validate_token(token) 234 | if is_valid: 235 | self._send_response({ 236 | 'status': 'success', 237 | 'message': 'Authentication successful', 238 | 'token': token 239 | }) 240 | return 241 | 242 | if password: 243 | if self.protocol.auth_manager.verify_credentials(password): 244 | client_info = {'username': username, 'address': self.client_address[0]} 245 | new_token = self.protocol.auth_manager.generate_token(client_info) 246 | self._send_response({ 247 | 'status': 'success', 248 | 'message': 'Authentication successful', 249 | 'token': new_token 250 | }) 251 | return 252 | 253 | self._send_response({'status': 'error', 'message': 'Authentication failed'}, 401) 254 | 255 | def _handle_get_binary_info(self): 256 | """Handle requests for binary information""" 257 | try: 258 | binary_info = { 259 | 'filename': self.protocol.bv.file.filename, 260 | 'file_size': self.protocol.bv.end, 261 | 'start': self.protocol.bv.start, 262 | 'end': self.protocol.bv.end, 263 | 'entry_point': self.protocol.bv.entry_point, 264 | 'arch': self.protocol.bv.arch.name, 265 | 'platform': self.protocol.bv.platform.name, 266 | 'segments': self.protocol._get_segments_info(), 267 | 'sections': self.protocol._get_sections_info(), 268 | 'functions_count': len(self.protocol.bv.functions), 269 | 'symbols_count': len(self.protocol.bv.symbols) 270 | } 271 | 272 | self._send_response({ 273 | 'status': 'success', 274 | 'binary_info': binary_info 275 | }) 276 | 277 | except Exception as e: 278 | logger.log_error(f"Error getting binary info: {e}") 279 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 280 | self._send_response({'status': 'error', 'message': str(e)}, 500) 281 | 282 | def _get_function_context(self, address: int) -> Dict[str, Any]: 283 | res = self.protocol.bv.get_functions_containing(address) 284 | func = None 285 | if len(res) > 0: 286 | func = res[0] 287 | else: 288 | return None 289 | 290 | function_info = { 291 | 'name': func.name, 292 | 'start': func.address_ranges[0].start, 293 | 'end': func.address_ranges[0].end, 294 | 'pseudo_c': self._get_pseudo_c_text(self.protocol.bv, func), 295 | 'call_sites': self._get_call_sites(func), 296 | 'basic_blocks': self._get_basic_blocks_info(func), 297 | 'parameters': self._get_parameters(func), 298 | 'variables': self._get_variables(func), 299 | 'global_variables': self._get_global_variables(), 300 | 'disassembly': self._get_disassembly(func), 301 | 'incoming_calls': self._get_incoming_calls(func) 302 | } 303 | return function_info 304 | 305 | def _handle_get_function_context_by_address(self): 306 | """Handle requests for function context""" 307 | try: 308 | address = int(self.path.split('/')[-1], 0) 309 | function_info = self._get_function_context(address) 310 | if function_info is None: 311 | self._send_response({'status': 'error', 'message': f'No function found at address 0x{address:x}'}, 404) 312 | return 313 | 314 | self._send_response({ 315 | 'status': 'success', 316 | 'function': function_info 317 | }) 318 | 319 | except Exception as e: 320 | logger.log_error(f"Error getting function context: {e}") 321 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 322 | self._send_response({'status': 'error', 'message': str(e)}, 500) 323 | 324 | def _handle_get_function_context_by_name(self): 325 | """Handle requests for function context by name""" 326 | try: 327 | name = self.path.split('/')[-1] 328 | res = self.protocol.bv.get_functions_by_name(name) 329 | func = None 330 | if len(res) > 0: 331 | func = res[0] 332 | else: 333 | self._send_response({'status': 'error', 'message': f'No function found with name: {name}'}, 404) 334 | return 335 | 336 | function_info = self._get_function_context(func.start) 337 | if function_info is None: 338 | self._send_response({'status': 'error', 'message': f'No function found with name: {name}'}, 404) 339 | return 340 | self._send_response({ 341 | 'status': 'success', 342 | 'function': function_info 343 | }) 344 | except Exception as e: 345 | logger.log_error(f"Error getting function context by name: {e}") 346 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 347 | self._send_response({'status': 'error', 'message': str(e)}, 500) 348 | 349 | def _handle_get_all_function_names(self): 350 | """Handle requests for all function names""" 351 | try: 352 | function_names = [{'name': func.name, 'address': func.start} for func in self.protocol.bv.functions] 353 | self._send_response({ 354 | 'status': 'success', 355 | 'function_names': function_names 356 | }) 357 | except Exception as e: 358 | logger.log_error(f"Error getting all function names: {e}") 359 | self._send_response({'status': 'error', 'message': str(e)}, 500) 360 | 361 | def _handle_update_function_name(self, data): 362 | """Handle requests to update function name""" 363 | try: 364 | if not data or 'name' not in data: 365 | self._send_response({'status': 'error', 'message': 'New name is required'}, 400) 366 | return 367 | 368 | new_name = data['name'] 369 | name = self.path.split('/')[-2] 370 | func = self._get_function_by_name(name) 371 | if not func: 372 | self._send_response({'status': 'error', 'message': f'No function found with name {name}'}, 404) 373 | return 374 | 375 | old_name = func.name 376 | func.name = new_name 377 | 378 | self._send_response({ 379 | 'status': 'success', 380 | 'message': f'Function name updated from "{old_name}" to "{new_name}"' 381 | }) 382 | 383 | except Exception as e: 384 | logger.log_error(f"Error updating function name: {e}") 385 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 386 | self._send_response({'status': 'error', 'message': str(e)}, 500) 387 | 388 | def _handle_update_variable_name(self, data): 389 | """Handle requests to update variable name""" 390 | try: 391 | if not data or 'name' not in data: 392 | self._send_response({'status': 'error', 'message': 'New name is required'}, 400) 393 | return 394 | 395 | new_name = data['name'] 396 | func_name = self.path.split('/')[-3] 397 | func = self._get_function_by_name(func_name) 398 | if not func: 399 | self._send_response({'status': 'error', 'message': f'No function found at address {func_name}'}, 404) 400 | return 401 | 402 | # Find the variable by name 403 | for var in func.vars: 404 | if var.name == self.path.split('/')[-2]: 405 | old_name = var.name 406 | var.name = new_name 407 | self._send_response({ 408 | 'status': 'success', 409 | 'message': f'Variable name updated from "{old_name}" to "{new_name}"' 410 | }) 411 | return 412 | """ 413 | We need to handle the case where the LLM is trying to change 414 | the name of a global variable. We need to find the global and 415 | rename it. 416 | """ 417 | for var in self._get_globals_from_func(func): 418 | current_var_name = self.path.split('/')[-2] 419 | if var['name'] == current_var_name: 420 | for addr, gvar in self.protocol.bv.data_vars.items(): 421 | if addr == var['location']: 422 | gvar.name = new_name 423 | self._send_response({ 424 | 'status': 'success', 425 | 'message': f'Variable name updated from "{current_var_name}" to "{new_name}"' 426 | }) 427 | return 428 | 429 | self._send_response({'status': 'error', 'message': f'No variable with name {self.path.split("/")[-1]} found in function'}, 404) 430 | 431 | except Exception as e: 432 | logger.log_error(f"Error updating variable name: {e}") 433 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 434 | self._send_response({'status': 'error', 'message': str(e)}, 500) 435 | 436 | def _handle_get_global_variable_data(self): 437 | """Handle requests access data from a global address""" 438 | try: 439 | func_name = self.path.split('/')[-2] 440 | func = self._get_function_by_name(func_name) 441 | if not func: 442 | self._send_response({'status': 'error', 'message': f'No function found at address {func_name}'}, 404) 443 | return 444 | # Find the variable by name 445 | global_name = self.path.split('/')[-1] 446 | """ 447 | We need to handle the case where the LLM is trying to change 448 | the name of a global variable. We need to find the global and 449 | rename it. 450 | """ 451 | for var in self._get_globals_from_func(func): 452 | if var['name'] == global_name: 453 | for addr, gvar in self.protocol.bv.data_vars.items(): 454 | if addr == var['location']: 455 | read_address = None 456 | rbytes = None 457 | if gvar.value: 458 | target_val = gvar.value 459 | # Getting the .value for a value found with heuristics 460 | # will actually return this value. If it's an int 461 | # then it's likely a pointer for us to follow. 462 | if isinstance(target_val, bytes): 463 | rbytes = target_val 464 | elif isinstance(target_val, int): 465 | read_address = target_val 466 | else: 467 | read_address = addr 468 | 469 | # If there is not a defined value at address, then read 470 | # an arbitrary amount of data as a last ditch effort. 471 | if read_address and not rbytes: 472 | rbytes = self.protocol.bv.read(read_address, 256) 473 | self._send_response({ 474 | 'status': 'success', 475 | 'message': f'Byte slice from global: {rbytes}' 476 | }) 477 | return 478 | except Exception as e: 479 | logger.log_error(f"Error updating variable name: {e}") 480 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 481 | self._send_response({'status': 'error', 'message': str(e)}, 500) 482 | 483 | def _handle_add_comment_to_address(self, data): 484 | """Handle requests to add a comment to an address""" 485 | try: 486 | if not data or 'comment' not in data: 487 | self._send_response({'status': 'error', 'message': 'Comment text is required'}, 400) 488 | return 489 | 490 | comment = data['comment'] 491 | self.protocol.bv.set_comment_at(int(self.path.split('/')[-1], 0), comment) 492 | 493 | self._send_response({ 494 | 'status': 'success', 495 | 'message': f'Comment added at address 0x{int(self.path.split("/")[-1], 0):x}' 496 | }) 497 | 498 | except Exception as e: 499 | logger.log_error(f"Error adding comment: {e}") 500 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 501 | self._send_response({'status': 'error', 'message': str(e)}, 500) 502 | 503 | def _handle_add_comment_to_function(self, data): 504 | """Handle requests to add a comment to a function""" 505 | try: 506 | if not data or 'comment' not in data: 507 | self._send_response({'status': 'error', 'message': 'Comment text is required'}, 400) 508 | return 509 | 510 | comment = data['comment'] 511 | name = self.path.split('/')[-2] 512 | func = self._get_function_by_name(name) 513 | if not func: 514 | self._send_response({'status': 'error', 'message': f'No function found with name: {name}'}, 404) 515 | return 516 | self.protocol.bv.set_comment_at(func.start, comment) 517 | 518 | self._send_response({ 519 | 'status': 'success', 520 | 'message': f'Comment added to function {name}' 521 | }) 522 | 523 | except Exception as e: 524 | logger.log_error(f"Error adding comment: {e}") 525 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 526 | self._send_response({'status': 'error', 'message': str(e)}, 500) 527 | 528 | def _get_function_by_name(self, name): 529 | """Acquire function by name instead of address""" 530 | logger.log_info(f"Getting function by name: {name}") 531 | res = self.protocol.bv.get_functions_by_name(name) 532 | # TODO: is there a scenario where there's more than one with the same name? 533 | if len(res) > 0: 534 | return res[0] 535 | else: 536 | return None 537 | 538 | def _get_function_by_address(self, address): 539 | """Acquire function by address instead of name""" 540 | res = self.protocol.bv.get_functions_containing(address) 541 | if res: 542 | return res[0] 543 | else: 544 | return None 545 | 546 | def _handle_get_function_disassembly(self): 547 | """Handle requests for function disassembly with function name""" 548 | try: 549 | name = self.path.split('/')[-2] 550 | func = self._get_function_by_name(name) 551 | if not func: 552 | self._send_response({'status': 'error', 'message': f'No function found with name: {name}'}, 404) 553 | return 554 | else: 555 | disassembly = self._get_disassembly(func) 556 | self._send_response({ 557 | 'status': 'success', 558 | 'disassembly': disassembly 559 | }) 560 | except Exception as e: 561 | logger.log_error(f"Error getting function disassembly: {e}") 562 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 563 | self._send_response({'status': 'error', 'message': str(e)}, 500) 564 | 565 | def _handle_get_function_pseudocode(self): 566 | """Handle requests for function pseudocode with function name""" 567 | try: 568 | name = self.path.split('/')[-2] 569 | func = self._get_function_by_name(name) 570 | if not func: 571 | self._send_response({'status': 'error', 'message': f'No function found with name: {name}'}, 404) 572 | return 573 | 574 | pseudocode = self._get_pseudo_c_text(self.protocol.bv, func) 575 | 576 | self._send_response({ 577 | 'status': 'success', 578 | 'pseudocode': pseudocode 579 | }) 580 | 581 | except Exception as e: 582 | logger.log_error(f"Error getting function pseudocode: {e}") 583 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 584 | self._send_response({'status': 'error', 'message': str(e)}, 500) 585 | 586 | def _is_global_ptr(self, obj): 587 | """Callback to look for a HighLevelILConstPtr in instruction line""" 588 | if(isinstance(obj, HighLevelILConstPtr)): 589 | return obj 590 | 591 | def _get_globals_from_func(self, func: binaryninja.function.Function) -> List[Dict[str, Any]]: 592 | """Get global variables in a given HLIL function""" 593 | res = [] 594 | gvar_results = [] 595 | """ 596 | We enumerate all instructions in basic blocks to find 597 | pointers to global variables. We recursively enumerate 598 | each instruction line for HighLevelILConstPtr to do this. 599 | """ 600 | for bb in func.hlil: 601 | for instr in bb: 602 | res += (list(instr.traverse(self._is_global_ptr))) 603 | 604 | """ 605 | Once we find a pointer, we get the pointer's address value 606 | and find the data variable that this corresponds to in 607 | order to find the variable's name. Unnamed variables 608 | in the format of data_[address] return None for their name 609 | so we need to format this ourselves to match the pseudocode 610 | output. 611 | """ 612 | for r in res: 613 | address = r.constant 614 | for gaddr, gvar in self.protocol.bv.data_vars.items(): 615 | if address == gaddr: 616 | var_name = None 617 | if not gvar.name: 618 | var_name = f"data_{address:2x}" 619 | else: 620 | var_name = gvar.name 621 | gvar_results.append({ 622 | 'name': var_name, 623 | 'type': str(gvar.type), 624 | 'location': gaddr 625 | }) 626 | return gvar_results 627 | 628 | def _handle_get_function_variables(self): 629 | """Handle requests for function variables""" 630 | try: 631 | name = self.path.split('/')[-2] 632 | func = self._get_function_by_name(name) 633 | if not func: 634 | self._send_response({'status': 'error', 'message': f'No function found with name {name}'}, 404) 635 | return 636 | 637 | variables = { 638 | 'parameters': self._get_parameters(func), 639 | 'local_variables': self._get_variables(func), 640 | 'global_variables': self._get_globals_from_func(func) 641 | } 642 | 643 | self._send_response({ 644 | 'status': 'success', 645 | 'variables': variables 646 | }) 647 | 648 | except Exception as e: 649 | logger.log_error(f"Error getting function variables: {e}") 650 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 651 | self._send_response({'status': 'error', 'message': str(e)}, 500) 652 | 653 | def _handle_get_cross_references_to_function(self): 654 | """Handle requests for cross references to a function by address or name""" 655 | try: 656 | val = self.path.split('/')[-1] 657 | logger.log_info(f"Getting cross references to function: {val}") 658 | if val.startswith('0x'): 659 | val = int(val, 0) 660 | func = self._get_function_by_address(val) 661 | else: 662 | func = self._get_function_by_name(val) 663 | if func is None: 664 | self._send_response({'status': 'error', 'message': f'No function found with name {val}'}, 404) 665 | return 666 | cross_references = self._get_cross_references_to_function(func.name) 667 | if len(cross_references) == 0: 668 | self._send_response({'status': 'error', 'message': f'No cross references found for function {name}'}, 404) 669 | self._send_response({ 670 | 'status': 'success', 671 | 'cross_references': cross_references 672 | }) 673 | except Exception as e: 674 | logger.log_error(f"Error getting cross references to function: {e}") 675 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 676 | self._send_response({'status': 'error', 'message': str(e)}, 500) 677 | 678 | def _get_llil_text(self, func: binaryninja.function.Function) -> List[str]: 679 | """Get LLIL text for a function""" 680 | result = [] 681 | for block in func.llil: 682 | for instruction in block: 683 | result.append({'address': instruction.address, 'text': str(instruction)}) 684 | return result 685 | 686 | def _get_mlil_text(self, func: binaryninja.function.Function) -> List[str]: 687 | """Get MLIL text for a function""" 688 | result = [] 689 | for block in func.mlil: 690 | for instruction in block: 691 | result.append({'address': instruction.address, 'text': str(instruction)}) 692 | return result 693 | 694 | def _get_hlil_text(self, func: binaryninja.function.Function) -> List[str]: 695 | """Get HLIL text for a function""" 696 | result = [] 697 | for block in func.hlil: 698 | for instruction in block: 699 | result.append({'address': instruction.address, 'text': str(instruction)}) 700 | return result 701 | 702 | def _get_pseudo_c_text(self, bv: BinaryView, function: Function) -> List[str]: 703 | """ 704 | Get pseudo-c text for a function, big thanks to Asher Devila L. 705 | for help with this https://github.com/AsherDLL/PCDump-bn/blob/main/__init__.py 706 | """ 707 | lines = [] 708 | settings = DisassemblySettings() 709 | settings.set_option(DisassemblyOption.ShowAddress, True) 710 | settings.set_option(DisassemblyOption.WaitForIL, True) 711 | obj = LinearViewObject.language_representation(bv, settings) 712 | cursor_end = LinearViewCursor(obj) 713 | cursor_end.seek_to_address(function.highest_address) 714 | body = bv.get_next_linear_disassembly_lines(cursor_end) 715 | cursor_end.seek_to_address(function.highest_address) 716 | header = bv.get_previous_linear_disassembly_lines(cursor_end) 717 | for line in header: 718 | lines.append(f'{str(line)}\n') 719 | for line in body: 720 | lines.append(f'{str(line)}\n') 721 | with_addr = self._get_addr_pseudo_c_from_text(lines) 722 | return with_addr 723 | 724 | def _get_addr_pseudo_c_from_text(self, lines: list) -> List[str]: 725 | """Get addresses and pseudo-c from pseudo-c text output""" 726 | if lines is None: 727 | return [] 728 | else: 729 | result = [] 730 | for l in lines: 731 | lr = re.findall("(^[0-9A-Fa-f]+)(.*)$", l) 732 | if lr: 733 | # Converting binja address format of 0x[Address] 734 | addr = int("0x" + lr[0][0], 0) 735 | pseudo_c = lr[0][1] 736 | result.append({'address': addr, 'text': pseudo_c}) 737 | return result 738 | 739 | def _get_call_sites(self, func: binaryninja.function.Function) -> List[Dict[str, Any]]: 740 | """Get call sites within a function""" 741 | result = [] 742 | for ref in func.call_sites: 743 | called_func = self.protocol.bv.get_function_at(ref.address) 744 | called_name = called_func.name if called_func else "unknown" 745 | result.append({ 746 | 'address': ref.address, 747 | 'target': called_name 748 | }) 749 | return result 750 | 751 | def _get_cross_references_to_function(self, name: str) -> List[Dict[str, Any]]: 752 | """ 753 | Get cross references to a function by name. 754 | This returns functions containing cross-reference locations, 755 | instead of the actual cross-reference locations. 756 | """ 757 | result = [] 758 | func = self._get_function_by_name(name) 759 | if not func: 760 | return [] 761 | for ref in self.protocol.bv.get_code_refs(func.start): 762 | called_func = self.protocol.bv.get_functions_containing(ref.address)[0] 763 | result.append({ 764 | 'address': ref.address, 765 | 'function': called_func.name 766 | }) 767 | return result 768 | 769 | def _get_basic_blocks_info(self, func: binaryninja.function.Function) -> List[Dict[str, Any]]: 770 | """Get information about basic blocks in a function""" 771 | result = [] 772 | for block in func.basic_blocks: 773 | result.append({ 774 | 'start': block.start, 775 | 'end': block.end, 776 | 'incoming_edges': [edge.source.start for edge in block.incoming_edges], 777 | 'outgoing_edges': [edge.target.start for edge in block.outgoing_edges] 778 | }) 779 | return result 780 | 781 | def _get_parameters(self, func: binaryninja.function.Function) -> List[Dict[str, Any]]: 782 | """Get information about function parameters""" 783 | result = [] 784 | for param in func.parameter_vars: 785 | result.append({ 786 | 'name': param.name, 787 | 'type': str(param.type), 788 | 'location': str(param.storage) 789 | }) 790 | return result 791 | 792 | def _get_variables(self, func: binaryninja.function.Function) -> List[Dict[str, Any]]: 793 | """Get information about function variables""" 794 | result = [] 795 | for var in func.vars: 796 | result.append({ 797 | 'name': var.name, 798 | 'type': str(var.type), 799 | 'location': str(var.storage), 800 | 'id': var.identifier 801 | }) 802 | return result 803 | 804 | def _get_global_variables(self) -> List[Dict[str, Any]]: 805 | """Get information about global variables""" 806 | result = [] 807 | for address, var in self.protocol.bv.data_vars.items(): 808 | result.append({ 809 | 'name': var.name, 810 | 'type': str(var.type), 811 | 'location': address 812 | }) 813 | return result 814 | 815 | def _get_disassembly(self, func: binaryninja.function.Function) -> List[Dict[str, Any]]: 816 | """Get disassembly for a function""" 817 | result = [] 818 | for block in func: 819 | all_dis = block.get_disassembly_text() 820 | for i, instruction in enumerate(all_dis): 821 | if i == len(all_dis)-1: 822 | instr_len = block.end-instruction.address 823 | else: 824 | instr_len = all_dis[i+1].address-all_dis[i].address 825 | result.append({ 826 | 'address': instruction.address, 827 | 'text': str(instruction) 828 | }) 829 | return result 830 | 831 | def _get_incoming_calls(self, func: binaryninja.function.Function) -> List[Dict[str, Any]]: 832 | """Get incoming calls to a function""" 833 | result = [] 834 | for ref in self.protocol.bv.get_code_refs(func.start): 835 | caller = self.protocol.bv.get_function_at(ref.address) 836 | if caller: 837 | result.append({ 838 | 'address': ref.address, 839 | 'function': caller.name 840 | }) 841 | return result 842 | 843 | def _get_block_disassembly(self, block) -> List[Dict[str, Any]]: 844 | """Get disassembly for a basic block""" 845 | result = [] 846 | for instruction in block: 847 | result.append({ 848 | 'address': instruction.address, 849 | 'text': instruction.get_disassembly_text(), 850 | 'bytes': [b for b in instruction.bytes], 851 | 'length': instruction.length 852 | }) 853 | return result 854 | 855 | def _get_block_llil(self, block) -> List[str]: 856 | """Get LLIL text for a basic block""" 857 | result = [] 858 | func = block.function 859 | llil_block = func.get_low_level_il_at(block.start).ssa_form 860 | if llil_block: 861 | for instruction in llil_block: 862 | result.append(f"0x{instruction.address:x}: {instruction}") 863 | return result 864 | 865 | def _get_block_mlil(self, block) -> List[str]: 866 | """Get MLIL text for a basic block""" 867 | result = [] 868 | func = block.function 869 | mlil_block = func.get_medium_level_il_at(block.start).ssa_form 870 | if mlil_block: 871 | for instruction in mlil_block: 872 | result.append(f"0x{instruction.address:x}: {instruction}") 873 | return result 874 | 875 | def _get_block_hlil(self, block) -> List[str]: 876 | """Get HLIL text for a basic block""" 877 | result = [] 878 | func = block.function 879 | hlil_block = func.get_high_level_il_at(block.start).ssa_form 880 | if hlil_block: 881 | for instruction in hlil_block: 882 | result.append(f"0x{instruction.address:x}: {instruction}") 883 | return result 884 | 885 | class BinjaLattice: 886 | """ 887 | Protocol for communicating between Binary Ninja an external MCP Server or tools. 888 | This protocol handles sending context from Binary Ninja to MCP Server and receiving 889 | responses to integrate back into the Binary Ninja UI. 890 | """ 891 | 892 | def __init__(self, bv: BinaryView, port: int = 9000, host: str = "localhost", use_ssl: bool = False): 893 | """ 894 | Initialize the model context protocol. 895 | 896 | Args: 897 | bv: BinaryView object representing the currently analyzed binary 898 | port: Port number for communication 899 | host: Host address for the server 900 | use_ssl: Whether to use SSL/TLS encryption 901 | """ 902 | self.bv = bv 903 | self.port = port 904 | self.host = host 905 | self.use_ssl = use_ssl 906 | self.auth_manager = AuthManager() 907 | self.server = None 908 | 909 | def start_server(self): 910 | """Start the HTTP server""" 911 | try: 912 | if self.use_ssl: 913 | logger.log_info("Starting server with SSL") 914 | cert_file = os.path.join(os.path.dirname(__file__), "server.crt") 915 | key_file = os.path.join(os.path.dirname(__file__), "server.key") 916 | 917 | self.server = HTTPServer((self.host, self.port), 918 | lambda *args, **kwargs: LatticeRequestHandler(*args, protocol=self, **kwargs)) 919 | self.server.socket = ssl.wrap_socket(self.server.socket, 920 | server_side=True, 921 | certfile=cert_file, 922 | keyfile=key_file) 923 | else: 924 | self.server = HTTPServer((self.host, self.port), 925 | lambda *args, **kwargs: LatticeRequestHandler(*args, protocol=self, **kwargs)) 926 | 927 | # Run server in a separate thread 928 | server_thread = threading.Thread(target=self.server.serve_forever) 929 | server_thread.daemon = True 930 | server_thread.start() 931 | 932 | logger.log_info(f"Server started on {self.host}:{self.port}") 933 | logger.log_info(f"Authentication API key: {self.auth_manager.api_key}") 934 | logger.log_info(f"Use this key to authenticate clients") 935 | 936 | except Exception as e: 937 | logger.log_error(f"Failed to start server: {e}") 938 | logger.log_error("Stack trace: %s" % traceback.format_exc()) 939 | self.stop_server() 940 | 941 | def stop_server(self): 942 | """Stop the server""" 943 | if self.server: 944 | self.server.shutdown() 945 | self.server.server_close() 946 | logger.log_info("Server stopped") 947 | 948 | def _get_segments_info(self) -> List[Dict[str, Any]]: 949 | """Get information about binary segments""" 950 | result = [] 951 | for segment in self.bv.segments: 952 | result.append({ 953 | 'start': segment.start, 954 | 'end': segment.end, 955 | 'length': segment.length, 956 | 'permissions': { 957 | 'read': segment.readable, 958 | 'write': segment.writable, 959 | 'execute': segment.executable 960 | } 961 | }) 962 | return result 963 | 964 | def _get_sections_info(self) -> List[Dict[str, Any]]: 965 | """Get information about binary sections""" 966 | result = [] 967 | for section in self.bv.sections.values(): 968 | result.append({ 969 | 'name': section.name, 970 | 'start': section.start, 971 | 'end': section.end, 972 | 'length': section.length, 973 | 'semantics': str(section.semantics) 974 | }) 975 | return result 976 | 977 | protocol_instances = {} 978 | 979 | def register_plugin_command(view): 980 | protocol = BinjaLattice(view, use_ssl=False) 981 | protocol.start_server() 982 | protocol_instances[view] = protocol 983 | return protocol 984 | 985 | def stop_lattice_protocol_server(view): 986 | protocol = protocol_instances.get(view) 987 | if protocol: 988 | protocol.stop_server() 989 | del protocol_instances[view] 990 | 991 | PluginCommand.register( 992 | "Start Lattice Protocol Server", 993 | "Start server for Binary Ninja Lattice protocol with authentication", 994 | register_plugin_command 995 | ) 996 | 997 | PluginCommand.register( 998 | "Stop Lattice Protocol Server", 999 | "Stop server for Binary Ninja Lattice protocol", 1000 | stop_lattice_protocol_server 1001 | ) 1002 | --------------------------------------------------------------------------------