Scripting Ghidra to auto rename functions and variables with an LLM

# -*- coding: utf-8 -*-
# Ghidra Function Auto-Namer using Claude Code CLI
# @author Matt Webb + Claude Code
# @category Analysis
# @keybinding
# @menupath Tools.Claude Function Namer
# @toolbar

import os
import re
import json
import time
import java.io.BufferedReader as BufferedReader
import java.io.InputStreamReader as InputStreamReader
import java.lang.ProcessBuilder as ProcessBuilder
import java.util.concurrent.TimeUnit as TimeUnit
import java.util.ArrayList as ArrayList
import java.util.concurrent.Executors as Executors
import java.util.concurrent.Callable as Callable
import java.util.concurrent.TimeoutException as TimeoutException
import java.util.concurrent.ExecutionException as ExecutionException
from ghidra.app.decompiler import DecompInterface, DecompileOptions
from ghidra.program.model.symbol import SourceType
from ghidra.program.model.pcode import HighFunctionDBUtil
from ghidra.util.exception import DuplicateNameException

# --- CONFIGURATION ---
CLAUDE_PATH = "/opt/homebrew/bin/claude"
CLAUDE_OAUTH_TOKEN = os.environ["CLAUDE_CODE_OAUTH_TOKEN"]
CONCURRENCY = 30  # Number of parallel Claude API calls
PROMPT_SUFFIX = "<additional info to add to the prompt about the binary>"
# --- END CONFIGURATION ---

def get_decompiled_code(function, decompiler):
    """Decompile a function and return its C code and decompile results."""
    try:
        results = decompiler.decompileFunction(function, 60, monitor)
        if not results or not results.decompileCompleted():
            return None, None

        decomp = results.getDecompiledFunction()
        if not decomp:
            return None, None

        c_code = decomp.getC()
        if not c_code or len(c_code.strip()) == 0:
            return None, None

        return c_code, results
    except Exception as e:
        print("Decompilation exception: {}".format(e))
        return None, None

def is_valid_function_to_rename(function):
    """Check if a function is worth renaming."""
    if function.isThunk() or function.isExternal():
        return False
    return True

def calls_unnamed_functions(function):
    """Check if a function calls any other FUN_ functions (excluding itself for recursion)."""
    try:
        called_functions = function.getCalledFunctions(monitor)
        function_addr = function.getEntryPoint()

        for called_func in called_functions:
            if called_func.getEntryPoint().equals(function_addr):
                continue

            if called_func.getName().startswith("FUN_"):
                return True

        return False
    except Exception as e:
        print("  -> WARNING: Error checking dependencies for {}: {}".format(function.getName(), e))
        return False

def get_unnamed_functions(func_manager):
    """Get all functions that start with FUN_ and are valid for renaming."""
    unnamed_funcs = []
    for function in func_manager.getFunctions(True):
        if function.getName().startswith("FUN_"):
            if is_valid_function_to_rename(function):
                unnamed_funcs.append(function)
    return unnamed_funcs

def query_claude_for_name(decompiled_code, claude_path):
    """Call Claude Code CLI to get suggestions for function and variable names."""
    prompt = """{code}

Analyze the function above. Provide new, accurate names for the function, its local variables, and any global variables (DAT_*).
The function name must be camelCase.

{suffix}

Respond *only* with a valid JSON object in the following format:

functionName,
  "globalVariableNames": DAT_00123456
}}

Only include variable mappings you have high confidence in. If unsure, exclude it.

JSON Response:
""".format(code=decompiled_code, suffix=PROMPT_SUFFIX)

    max_retries = 3

    for attempt in range(max_retries):
        try:
            cmd_list = ArrayList()
            cmd_list.add(claude_path)
            cmd_list.add("--print")
            cmd_list.add("--dangerously-skip-permissions")
            cmd_list.add("--disallowed-tools")
            cmd_list.add("Bash,Edit,Write,FileSystem")

            builder = ProcessBuilder(cmd_list)
            if CLAUDE_OAUTH_TOKEN and len(CLAUDE_OAUTH_TOKEN.strip()) > 0:
                builder.environment().put("CLAUDE_CODE_OAUTH_TOKEN", CLAUDE_OAUTH_TOKEN)
            builder.redirectErrorStream(True)

            process = builder.start()

            output_stream = process.getOutputStream()
            output_stream.write(prompt.encode('utf-8'))
            output_stream.close()

            reader = BufferedReader(InputStreamReader(process.getInputStream(), "UTF-8"))
            output_lines = []
            line = reader.readLine()
            while line is not None:
                output_lines.append(line)
                line = reader.readLine()
            reader.close()

            finished = process.waitFor(120, TimeUnit.SECONDS)

            if not finished:
                continue

            exit_code = process.exitValue()
            raw_output = "\n".join(output_lines).strip()

            if exit_code != 0:
                continue

            json_string = re.sub(r'^```.*?\n', '', raw_output)
            json_string = re.sub(r'\n```$', '', json_string)
            json_string = json_string.strip()

            try:
                data = json.loads(json_string)

                suggested_name = data.get("functionName")
                variable_mappings = data.get("variableNames", {})
                global_mappings = data.get("globalVariableNames", {})

                if not suggested_name or not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', suggested_name):
                    continue

                return (suggested_name, variable_mappings, global_mappings)

            except ValueError:
                continue

        except Exception:
            continue

    return (None, None, None)

def rename_function(function, new_name):
    """Rename a function in Ghidra - NO transaction, relies on outer transaction."""
    try:
        old_name = function.getName()
        function.setName(new_name, SourceType.USER_DEFINED)
        print("  -> Renamed {} -> {}".format(old_name, new_name))
        return True
    except Exception as e:
        print("  -> Error renaming function: {}".format(e))
        return False

def rename_variables(function, variable_mappings, decompile_results):
    """Rename function variables - NO transaction, relies on outer transaction."""
    if not variable_mappings:
        return

    print("  -> Renaming {} local variables...".format(len(variable_mappings)))

    try:
        high_function = decompile_results.getHighFunction()
        if not high_function:
            return

        lsymMap = high_function.getLocalSymbolMap()
        if not lsymMap:
            return

        symbols = lsymMap.getSymbols()
        for sym in symbols:
            old_name = sym.getName()
            if old_name in variable_mappings:
                new_name = variable_mappings[old_name]

                if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', new_name):
                    continue

                if old_name == new_name:
                    continue

                attempted_name = new_name
                attempt_count = 0
                max_attempts = 50

                while attempt_count < max_attempts:
                    try:
                        HighFunctionDBUtil.updateDBVariable(
                            sym,
                            attempted_name,
                            None,
                            SourceType.USER_DEFINED
                        )
                        print("    {} -> {}".format(old_name, attempted_name))
                        break
                    except DuplicateNameException:
                        attempt_count += 1
                        attempted_name = "{}_{}".format(new_name, attempt_count)
                    except Exception as e:
                        print("    Error with {}: {}".format(old_name, str(e)))
                        break

    except Exception as e:
        print("  -> Error during variable rename: {}".format(e))

def rename_globals(program, global_mappings):
    """Rename global variables - NO transaction, relies on outer transaction."""
    if not global_mappings:
        return

    print("  -> Renaming {} global variables...".format(len(global_mappings)))

    symbol_table = program.getSymbolTable()

    for old_name, new_name in global_mappings.items():
        if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', new_name):
            continue

        try:
            symbols = symbol_table.getSymbols(old_name)
            if symbols and symbols.hasNext():
                symbol = symbols.next()

                attempted_name = new_name
                attempt_count = 0
                max_attempts = 50

                while attempt_count < max_attempts:
                    try:
                        symbol.setName(attempted_name, SourceType.USER_DEFINED)
                        print("    {} -> {}".format(old_name, attempted_name))
                        break
                    except DuplicateNameException:
                        attempt_count += 1
                        attempted_name = "{}_{}".format(new_name, attempt_count)
                    except Exception as e:
                        print("    Error with {}: {}".format(old_name, str(e)))
                        break

        except Exception as e:
            print("    Error with {}: {}".format(old_name, str(e)))

def format_time(seconds):
    """Format seconds into human-readable time."""
    if seconds < 60:
        return "{}s".format(int(seconds))
    elif seconds < 3600:
        return "{}m {}s".format(int(seconds / 60), int(seconds % 60))
    else:
        hours = int(seconds / 3600)
        minutes = int((seconds % 3600) / 60)
        return "{}h {}m".format(hours, minutes)

def initialize_decompiler(program):
    """Initialize the decompiler."""
    decompiler = DecompInterface()
    options = DecompileOptions()
    decompiler.setOptions(options)
    if not decompiler.openProgram(program):
        print("ERROR: Failed to initialize decompiler!")
        return None
    return decompiler

class ClaudeCallable(Callable):
    """Callable wrapper for Claude API calls."""
    def __init__(self, func, decomp_code, decomp_results, claude_path):
        self.func = func
        self.decomp_code = decomp_code
        self.decomp_results = decomp_results
        self.claude_path = claude_path

    def call(self):
        try:
            suggested_name, var_mappings, global_mappings = query_claude_for_name(
                self.decomp_code, self.claude_path)
            return (self.func, self.decomp_results, suggested_name, var_mappings, global_mappings)
        except Exception as e:
            print("  -> Error processing {}: {}".format(self.func.getName(), e))
            return (self.func, self.decomp_results, None, None, None)

def process_batch(batch_data, claude_path):
    """Process a batch of functions in parallel (Claude calls only)."""
    executor = Executors.newFixedThreadPool(CONCURRENCY)
    futures = []

    try:
        future_to_item = {}

        for item in batch_data:
            func, decomp_code, decomp_results = item
            callable_task = ClaudeCallable(func, decomp_code, decomp_results, claude_path)
            future = executor.submit(callable_task)
            futures.append(future)
            future_to_item[future] = item

        successful_results = []
        failed_items = []

        remaining_futures = list(futures)
        batch_start_time = time.time()
        per_future_timeout = 120
        max_batch_time = 150
        check_interval = 5

        while remaining_futures and (time.time() - batch_start_time) < max_batch_time:
            completed = []
            current_time = time.time()
            elapsed_time = current_time - batch_start_time

            batch_timed_out = elapsed_time > per_future_timeout

            for future in remaining_futures:
                try:
                    result = future.get(100, TimeUnit.MILLISECONDS)
                    successful_results.append(result)
                    completed.append(future)
                except TimeoutException:
                    if batch_timed_out:
                        original_item = future_to_item[future]
                        func, decomp_code, decomp_results = original_item
                        print("  -> Timeout for {}: exceeded {}s".format(
                            func.getName(), int(elapsed_time)))
                        failed_items.append(original_item)
                        completed.append(future)
                        future.cancel(True)
                except (ExecutionException, Exception) as e:
                    original_item = future_to_item[future]
                    func, decomp_code, decomp_results = original_item
                    print("  -> Error for {}: {}".format(func.getName(), str(e)[:100]))
                    failed_items.append(original_item)
                    completed.append(future)

            for future in completed:
                remaining_futures.remove(future)

            if batch_timed_out:
                break

            if remaining_futures:
                time.sleep(check_interval)

        for future in remaining_futures:
            original_item = future_to_item[future]
            func, decomp_code, decomp_results = original_item
            print("  -> Max batch time exceeded for {}".format(func.getName()))
            failed_items.append(original_item)
            future.cancel(True)

        return successful_results, failed_items
    finally:
        executor.shutdown()
        executor.awaitTermination(10, TimeUnit.SECONDS)

def checkpoint_program(program, iteration):
    """Commit current transaction and start a new one to preserve changes."""
    try:
        print("  -> Creating checkpoint {} (committing transaction)...".format(iteration))
        program.flushEvents()
        end(True)
        start()
        print("  -> Checkpoint complete (changes committed)")
        return True
    except Exception as e:
        print("  -> Error creating checkpoint: {}".format(e))
        try:
            start()
        except:
            pass
        return False

def main():
    """Main function."""
    program = getCurrentProgram()
    func_manager = program.getFunctionManager()

    print("=" * 60)
    print("Claude Function Auto-Namer (Parallel)")
    print("=" * 60)
    print("Using Claude at: {}".format(CLAUDE_PATH))
    print("Concurrency: {}".format(CONCURRENCY))

    if not os.path.exists(CLAUDE_PATH):
        print("\nERROR: Claude executable not found at: {}".format(CLAUDE_PATH))
        return

    if not CLAUDE_OAUTH_TOKEN or "YOUR_TOKEN_HERE" in CLAUDE_OAUTH_TOKEN:
        print("\nERROR: CLAUDE_OAUTH_TOKEN is not set!")
        return

    decompiler = initialize_decompiler(program)
    if not decompiler:
        return

    initial_unnamed = get_unnamed_functions(func_manager)
    total_functions = len(initial_unnamed)
    print("Found {} unnamed functions to process\n".format(total_functions))

    total_renamed = 0
    total_failed = 0
    start_time = time.time()
    function_times = []

    retry_queue = {}
    MAX_RETRIES = 3

    max_iterations = 1000
    iteration = 0

    while iteration < max_iterations:
        iteration += 1

        unnamed_funcs = get_unnamed_functions(func_manager)

        if not unnamed_funcs and not retry_queue:
            print("\nNo more unnamed functions!")
            break

        if not unnamed_funcs and retry_queue:
            all_maxed_out = all(retry_count >= MAX_RETRIES for _, _, _, retry_count in retry_queue.values())
            if all_maxed_out:
                print("\nOnly max-retry functions remain. Exiting.")
                break

        remaining = len(unnamed_funcs)
        processed = total_functions - remaining
        progress = (float(processed) / float(total_functions)) * 100 if total_functions > 0 else 0

        eta_str = "calculating..."
        if function_times:
            avg_batch_time = sum(function_times) / len(function_times)
            batches_needed = float(remaining) / float(CONCURRENCY)
            eta_seconds = avg_batch_time * batches_needed
            eta_str = format_time(eta_seconds)

        elapsed = time.time() - start_time
        print("\n[{:.1f}%] Iteration {} | {}/{} remaining | Elapsed: {} | ETA: {}".format(
            progress, iteration, remaining, total_functions, format_time(elapsed), eta_str))

        batch = []
        batch_retry_counts = {}
        batch_start_time = time.time()

        retry_limit = max(1, CONCURRENCY / 2)
        retry_addresses = list(retry_queue.keys())

        for addr_str in retry_addresses:
            if len(batch) >= retry_limit:
                break

            func, decomp_code, decomp_results, retry_count = retry_queue[addr_str]

            if not func.getName().startswith("FUN_"):
                print("Removing {} from retry queue (already renamed)".format(addr_str))
                del retry_queue[addr_str]
                continue

            if retry_count >= MAX_RETRIES:
                print("Max retries exceeded for {} (giving up)".format(func.getName()))
                del retry_queue[addr_str]
                total_failed += 1
                continue

            print("Retry #{} for: {}".format(retry_count + 1, func.getName()))
            batch.append((func, decomp_code, decomp_results))
            batch_retry_counts[addr_str] = retry_count
            del retry_queue[addr_str]

        candidates = []
        for function in unnamed_funcs:
            addr_str = str(function.getEntryPoint())
            if addr_str in retry_queue:
                continue

            if calls_unnamed_functions(function):
                continue

            candidates.append(function)

        for function in candidates:
            if len(batch) >= CONCURRENCY:
                break

            func_name = function.getName()
            print("Decompiling: {} at {}".format(func_name, function.getEntryPoint()))

            decompiled_code, decompile_results = get_decompiled_code(function, decompiler)

            if not decompiled_code:
                print("  -> Could not decompile (skipping)")
                total_failed += 1
                continue

            batch.append((function, decompiled_code, decompile_results))

        if not batch and unnamed_funcs:
            print("\n" + "=" * 60)
            print("WARNING: Circular dependencies detected!")
            print("=" * 60)
            print("All remaining {} unnamed functions call other unnamed functions.".format(len(unnamed_funcs)))
            print("Attempting to break the cycle by force-processing some functions...")
            print("=" * 60)

            force_candidates = []
            for function in unnamed_funcs:
                addr_str = str(function.getEntryPoint())
                if addr_str in retry_queue:
                    continue

                try:
                    called_funcs = function.getCalledFunctions(monitor)
                    unnamed_call_count = 0
                    for cf in called_funcs:
                        if cf.getName().startswith("FUN_") and not cf.getEntryPoint().equals(function.getEntryPoint()):
                            unnamed_call_count += 1
                    force_candidates.append((function, unnamed_call_count))
                except:
                    pass

            force_candidates.sort(key=lambda x: x[1])
            force_batch_size = min(5, len(force_candidates))

            print("Force processing {} functions with fewest dependencies:".format(force_batch_size))
            for function, dep_count in force_candidates[:force_batch_size]:
                print("  {} (calls {} unnamed functions)".format(function.getName(), dep_count))

                decompiled_code, decompile_results = get_decompiled_code(function, decompiler)
                if decompiled_code:
                    batch.append((function, decompiled_code, decompile_results))
                else:
                    print("    -> Could not decompile")
                    total_failed += 1

            if not batch:
                print("\nCould not decompile any functions. Stopping.")
                break

        print("Processing batch of {} functions in parallel...".format(len(batch)))
        results, failed_items = process_batch(batch, CLAUDE_PATH)

        for func, decomp_code, decomp_results in failed_items:
            addr_str = str(func.getEntryPoint())

            if addr_str in batch_retry_counts:
                old_retry_count = batch_retry_counts[addr_str]
                new_retry_count = old_retry_count + 1
            else:
                new_retry_count = 1

            retry_queue[addr_str] = (func, decomp_code, decomp_results, new_retry_count)
            print("  -> Added {} to retry queue (attempt {}/{})".format(
                func.getName(), new_retry_count, MAX_RETRIES))

        for func, decomp_results, suggested_name, var_mappings, global_mappings in results:
            print("Applying results for: {}".format(func.getName()))

            if suggested_name:
                if rename_function(func, suggested_name):
                    total_renamed += 1
                    if var_mappings:
                        rename_variables(func, var_mappings, decomp_results)
                    if global_mappings:
                        rename_globals(program, global_mappings)

                    addr_str = str(func.getEntryPoint())
                    if addr_str in retry_queue:
                        del retry_queue[addr_str]
                        print("  -> Removed {} from retry queue (succeeded)".format(func.getName()))
            else:
                total_failed += 1
                addr_str = str(func.getEntryPoint())
                if addr_str in retry_queue:
                    del retry_queue[addr_str]

        if retry_queue:
            print("  -> Retry queue size: {}".format(len(retry_queue)))

        checkpoint_program(program, iteration)

        batch_time = time.time() - batch_start_time
        function_times.append(batch_time)
        if len(function_times) > 20:
            function_times.pop(0)

    total_time = time.time() - start_time
    print("\n" + "=" * 60)
    print("Renaming complete!")
    print("Successfully renamed: {}".format(total_renamed))
    print("Failed/Skipped: {}".format(total_failed))
    if retry_queue:
        print("Still in retry queue: {}".format(len(retry_queue)))
    print("Total time: {}".format(format_time(total_time)))
    print("=" * 60)

    decompiler.dispose()

if __name__ == "__main__":
    main()