ExportFunction.scala

//ExportFunction.scala
// 
// This script exports one or all functions to C/H files using the CppExporter in Ghidra
// This script is designed to run in headless mode only
//
// @category: Export
//

// Use with https://github.com/edmcman/ghidra-scala-loader

import ghidra.app.script.GhidraScript
import ghidra.app.util.exporter.CppExporter
import ghidra.app.decompiler.DecompileOptions
import ghidra.util.task.TaskMonitor
import ghidra.program.model.listing.{Function, FunctionManager}

import java.io.File
import java.nio.file.{Files, Paths}
import scala.util.{Try, Success, Failure}

// Example usage in headless mode:
// For specific function by name:
// analyzeHeadless /path/to/project_dir MyProject -process binary_name -postScript ExportFunction.scala output_dir function_name
// For specific function by address:
// analyzeHeadless /path/to/project_dir MyProject -process binary_name -postScript ExportFunction.scala output_dir 0x1234abcd
// For all functions:
// analyzeHeadless /path/to/project_dir MyProject -process binary_name -postScript ExportFunction.scala output_dir

class ExportFunction extends GhidraScript:

  override def run(): Unit =
    if currentProgram == null then
      println("No program is loaded. Make sure to use -process or -import in analyzeHeadless")
      return
    
    // Get arguments - need output directory and optionally a function name or address
    val args = getScriptArgs
    
    if args.length < 1 then
      println("Error: Script requires at least one argument: <output_directory> [function_name_or_address]")
      println("Example for specific function by name: analyzeHeadless /path/to/project_dir MyProject -process binary_name -postScript ExportFunction.scala output_dir function_name")
      println("Example for specific function by address: analyzeHeadless /path/to/project_dir MyProject -process binary_name -postScript ExportFunction.scala output_dir 0x1234abcd")
      println("Example for all functions: analyzeHeadless /path/to/project_dir MyProject -process binary_name -postScript ExportFunction.scala output_dir")
      return
    
    val outputDirStr = args(0)
    
    // If function name/address provided, export just that function, otherwise export all
    if args.length >= 2 then
      val funcIdentifier = args(1)
      // Check if the identifier is a hex address
      if funcIdentifier.toLowerCase.startsWith("0x") then
        try
          val address = currentProgram.getAddressFactory.getAddress(funcIdentifier.substring(2))
          exportFunctionByAddress(outputDirStr, address)
        catch 
          case e: Exception => 
            println(s"Error parsing address: ${funcIdentifier}")
            println(e.getMessage)
      else
        // Treat as function name
        exportSpecificFunction(outputDirStr, funcIdentifier)
    else
      exportAllFunctions(outputDirStr)

  private def exportSpecificFunction(outputDirStr: String, funcName: String): Unit =
    println(s"Looking for function: $funcName")
    
    // Find the function
    val fm = currentProgram.getFunctionManager
    
    // Scala version of finding the function - fixed to avoid type mismatch
    val functionIter = fm.getFunctions(true)
    val functionOpt = Iterator.continually {
      if (functionIter.hasNext) Some(functionIter.next) else None
    }
    .takeWhile(_.isDefined)
    .map(_.get)
    .find(_.getName == funcName)
      
    if functionOpt.isEmpty then
      println(s"Error: Could not find function '$funcName' in program")
      return
    
    val function = functionOpt.get
    exportFunction(outputDirStr, function)
  
  private def exportFunctionByAddress(outputDirStr: String, address: ghidra.program.model.address.Address): Unit =
    println(s"Looking for function at address: ${address}")
    
    // Find the function at the given address
    val fm = currentProgram.getFunctionManager
    val function = fm.getFunctionAt(address)
    
    if function == null then
      println(s"Error: Could not find function at address '${address}' in program")
      return
    
    println(s"Found function: ${function.getName()}")
    exportFunction(outputDirStr, function)

  private def exportAllFunctions(outputDirStr: String): Unit =
    println("Exporting all functions...")
    
    // Create output directory if it doesn't exist
    val outputDir = new File(outputDirStr)
    if !outputDir.exists then
      outputDir.mkdirs
    
    // Create decompiler options
    val decompileOptions = new DecompileOptions
    
    // Initialize the CPP exporter with the correct parameters
    val exporter = new CppExporter(decompileOptions, true, true, true, false, null)
    
    // Create output file with the program name as base
    val progName = currentProgram.getName
    val cFilePath = new File(outputDir, s"$progName.c")
    
    println(s"Exporting all functions to ${cFilePath.getAbsolutePath}")
    
    // Export the entire program at once without specifying an address set
    val success = exporter.`export`(cFilePath, currentProgram, null, monitor)
    
    if success then
      println("Successfully exported all functions")
      val hFilePath = new File(outputDir, s"$progName.h")
      if hFilePath.exists then
        println(s"Header file created: ${hFilePath.getAbsolutePath}")
    else
      println("Failed to export functions")
  
  private def exportFunction(outputDirStr: String, function: Function): Unit =
    val funcName = function.getName
    println(s"Exporting function: $funcName at ${function.getEntryPoint}")
    
    // Create output directory if it doesn't exist
    val outputDir = new File(outputDirStr)
    if !outputDir.exists then
      outputDir.mkdirs
    
    // Create decompiler options
    val decompileOptions = new DecompileOptions
    
    // Step 1: Export the specific function's source code (C file)
    // Initialize the CPP exporter with createHeader and createFile set to true
    val sourceExporter = new CppExporter(decompileOptions, true, true, true, false, null)
    
    // Get function body
    val functionBody = function.getBody
    
    // Export just this function's source
    val cFilePath = new File(outputDir, s"$funcName.c")
    
    println(s"Exporting function $funcName to ${cFilePath.getAbsolutePath}")
    
    // Export only the selected function's source code
    val sourceSuccess = sourceExporter.`export`(cFilePath, currentProgram, functionBody, monitor)
    
    if !sourceSuccess then
      println(s"Failed to export function $funcName source")
      throw new RuntimeException(s"Export failed for function $funcName source")
    
    println(s"Successfully exported function $funcName source")
    
    // Step 2: Export the header file with all functions
    // Initialize a new CPP exporter with only createHeader set to true
    val headerExporter = new CppExporter(decompileOptions, true, false, true, false, null)
    
    // Create header file path with the function name instead of program name
    val hFilePath = new File(outputDir, s"$funcName.h")
    
    println(s"Exporting header file with all functions to ${hFilePath.getAbsolutePath}")
    
    // Export the header file for all functions (no specific address set)
    val headerSuccess = headerExporter.`export`(hFilePath, currentProgram, null, monitor)
    
    if headerSuccess then
      println(s"Successfully exported header file for all functions: ${hFilePath.getAbsolutePath}")
    else
      println("Failed to export header file")
      throw new RuntimeException("Export failed for header file")