package re3lib; import java.io.File; import java.sql.Connection; import java.sql.DriverManager; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; import org.sqlite.JDBC; import ghidra.app.script.GhidraScript; import ghidra.program.model.address.Address; import ghidra.program.model.listing.Function; public class FunctionDatabase implements AutoCloseable { public static enum Type { Auto(0), Fix(1), Stub(2), Ref(3); private final int value; Type(int value) { this.value = value; } public int getValue() { return value; } public static Type fromValue(int value) { for (Type type : Type.values()) { if (type.value == value) { return type; } } throw new IllegalArgumentException("Unknown type value: " + value); } } public static enum CallingConvention { Cdecl(0), Stdcall(1), Fastcall(2); private final int value; CallingConvention(int value) { this.value = value; } public int getValue() { return value; } public static CallingConvention fromValue(int value) { for (CallingConvention conv : CallingConvention.values()) { if (conv.value == value) { return conv; } } throw new IllegalArgumentException("Unknown calling convention value: " + value); } public static CallingConvention fromString(String convStr) { if (convStr == null) { return Cdecl; // Default } String lower = convStr.toLowerCase(); if (lower.contains("fastcall")) { return Fastcall; } else if (lower.contains("stdcall")) { return Stdcall; } else { return Cdecl; // Default } } @Override public String toString() { switch (this) { case Fastcall: return "fastcall"; case Stdcall: return "stdcall"; case Cdecl: default: return "cdecl"; } } } public static class FunctionEntry { public Address address; public String name; public File file; public Type type; public CallingConvention callingConvention; public String parameterNames; // Semicolon-separated parameter names public String parameterTypes; // Semicolon-separated parameter types public String returnType; // Function return type public FunctionEntry(Address address, String name, File file, Type type, CallingConvention callingConvention, String parameterNames, String parameterTypes, String returnType) { this.address = address; this.name = name; this.file = file; this.type = type; this.callingConvention = callingConvention; this.parameterNames = parameterNames != null ? parameterNames : ""; this.parameterTypes = parameterTypes != null ? parameterTypes : ""; this.returnType = returnType != null ? returnType : "void"; } // Helper methods to work with semicolon-separated parameter lists public String[] getParameterNamesArray() { if (parameterNames == null || parameterNames.trim().isEmpty()) { return new String[0]; } return parameterNames.split(";"); } public String[] getParameterTypesArray() { if (parameterTypes == null || parameterTypes.trim().isEmpty()) { return new String[0]; } return parameterTypes.split(";"); } public void setParameterNamesArray(String[] names) { if (names == null || names.length == 0) { this.parameterNames = ""; } else { this.parameterNames = String.join(";", names); } } public void setParameterTypesArray(String[] types) { if (types == null || types.length == 0) { this.parameterTypes = ""; } else { this.parameterTypes = String.join(";", types); } } } public static class GlobalEntry { public Address address; public String name; public File file; public GlobalEntry(Address address, String name, File file) { this.address = address; this.name = name; this.file = file; } } private File dbFile; private transient GhidraScript script; private Connection connection; // Prepared statements for better performance private PreparedStatement findByNameFunctions; private PreparedStatement findByNameImports; private PreparedStatement findByAddressFunctions; private PreparedStatement findByAddressImports; private PreparedStatement insertOrReplaceFunctions; private PreparedStatement insertOrReplaceImports; private PreparedStatement deleteByFilepathFunctions; private PreparedStatement deleteByFilepathImports; private PreparedStatement loadAllFunctions; private PreparedStatement loadAllImports; // Add these prepared statements after the existing ones private PreparedStatement findByNameGlobals; private PreparedStatement findByAddressGlobals; private PreparedStatement insertOrReplaceGlobals; private PreparedStatement deleteByFilepathGlobals; private PreparedStatement loadAllGlobals; public FunctionDatabase(GhidraScript script) { this.script = script; dbFile = RemanConfig.INSTANCE.databasePath; try { java.sql.DriverManager.registerDriver(new JDBC()); } catch (SQLException e) { script.printerr("Error registering JDBC driver: " + e.getMessage()); } } public void connect() throws Exception { if (connection != null && !connection.isClosed()) { return; // Already connected } if (!dbFile.exists()) { script.println("Database file not found: " + dbFile); // Create parent directories if they don't exist dbFile.getParentFile().mkdirs(); } try { connection = DriverManager.getConnection("jdbc:sqlite:" + dbFile.getAbsolutePath()); createTablesIfNotExist(); prepareCachedStatements(); script.println("Connected to database: " + dbFile); } catch (SQLException e) { script.println("Error connecting to database: " + e.getMessage()); // throw new Exception("Failed to connect to database", e); throw e; } } public void disconnect() throws Exception { if (connection != null && !connection.isClosed()) { try { // Close prepared statements closePreparedStatements(); connection.close(); script.println("Disconnected from database"); } catch (SQLException e) { script.println("Error disconnecting from database: " + e.getMessage()); throw new Exception("Failed to disconnect from database", e); } } } private void ensureConnection() throws Exception { if (connection == null || connection.isClosed()) { connect(); } } private void prepareCachedStatements() throws SQLException { // Find by name findByNameFunctions = connection.prepareStatement( "SELECT filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type FROM Functions WHERE name = ?"); findByNameImports = connection.prepareStatement( "SELECT filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type FROM Imports WHERE name = ?"); // Find by address findByAddressFunctions = connection.prepareStatement( "SELECT filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type FROM Functions WHERE address = ?"); findByAddressImports = connection.prepareStatement( "SELECT filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type FROM Imports WHERE address = ?"); // Insert or replace insertOrReplaceFunctions = connection.prepareStatement( "INSERT OR REPLACE INTO Functions (filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"); insertOrReplaceImports = connection.prepareStatement( "INSERT OR REPLACE INTO Imports (filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"); // Delete by filepath deleteByFilepathFunctions = connection.prepareStatement( "DELETE FROM Functions WHERE filepath = ?"); deleteByFilepathImports = connection.prepareStatement( "DELETE FROM Imports WHERE filepath = ?"); // Load all entries loadAllFunctions = connection.prepareStatement( "SELECT filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type FROM Functions"); loadAllImports = connection.prepareStatement( "SELECT filepath, name, address, type, calling_convention, parameter_names, parameter_types, return_type FROM Imports"); // Global statements findByNameGlobals = connection.prepareStatement( "SELECT filepath, name, address FROM Globals WHERE name = ?"); findByAddressGlobals = connection.prepareStatement( "SELECT filepath, name, address FROM Globals WHERE address = ?"); insertOrReplaceGlobals = connection.prepareStatement( "INSERT OR REPLACE INTO Globals (filepath, name, address) VALUES (?, ?, ?)"); deleteByFilepathGlobals = connection.prepareStatement( "DELETE FROM Globals WHERE filepath = ?"); loadAllGlobals = connection.prepareStatement( "SELECT filepath, name, address FROM Globals"); } private void closePreparedStatements() throws SQLException { if (findByNameFunctions != null) findByNameFunctions.close(); if (findByNameImports != null) findByNameImports.close(); if (findByAddressFunctions != null) findByAddressFunctions.close(); if (findByAddressImports != null) findByAddressImports.close(); if (insertOrReplaceFunctions != null) insertOrReplaceFunctions.close(); if (insertOrReplaceImports != null) insertOrReplaceImports.close(); if (deleteByFilepathFunctions != null) deleteByFilepathFunctions.close(); if (deleteByFilepathImports != null) deleteByFilepathImports.close(); if (loadAllFunctions != null) loadAllFunctions.close(); if (loadAllImports != null) loadAllImports.close(); if (findByNameGlobals != null) findByNameGlobals.close(); if (findByAddressGlobals != null) findByAddressGlobals.close(); if (insertOrReplaceGlobals != null) insertOrReplaceGlobals.close(); if (deleteByFilepathGlobals != null) deleteByFilepathGlobals.close(); if (loadAllGlobals != null) loadAllGlobals.close(); } public List loadAllEntries() throws Exception { ensureConnection(); List entries = new ArrayList<>(); try { // Load from Functions table try (ResultSet rs = loadAllFunctions.executeQuery()) { while (rs.next()) { FunctionEntry entry = createEntryFromResultSet(rs); if (entry != null) { entries.add(entry); } } } script.println("Loaded " + entries.size() + " function entries from database"); return entries; } catch (SQLException e) { script.println("Error loading entries: " + e.getMessage()); throw new Exception("Failed to load entries", e); } } private FunctionEntry createEntryFromResultSet(ResultSet rs) throws SQLException { String filepath = rs.getString("filepath"); String name = rs.getString("name"); String addressStr = rs.getString("address"); int typeValue = rs.getInt("type"); int callingConventionValue = rs.getInt("calling_convention"); String parameterNames = rs.getString("parameter_names"); String parameterTypes = rs.getString("parameter_types"); String returnType = rs.getString("return_type"); if (addressStr != null && !addressStr.isEmpty()) { Address address = script.getCurrentProgram().getAddressFactory().getAddress(addressStr); File file = new File(RemanConfig.INSTANCE.outputDir, filepath); Type type = Type.fromValue(typeValue); CallingConvention callingConvention = CallingConvention.fromValue(callingConventionValue); return new FunctionEntry(address, name, file, type, callingConvention, parameterNames, parameterTypes, returnType); } return null; } private void createTablesIfNotExist() throws SQLException { String createFunctions = """ CREATE TABLE IF NOT EXISTS Functions ( filepath TEXT, name TEXT, address TEXT, type INTEGER DEFAULT 0, calling_convention INTEGER DEFAULT 0, parameter_names TEXT DEFAULT '', parameter_types TEXT DEFAULT '', return_type TEXT DEFAULT '', PRIMARY KEY (name, filepath) )"""; String createImports = """ CREATE TABLE IF NOT EXISTS Imports ( filepath TEXT, name TEXT, address TEXT, type INTEGER DEFAULT 0, calling_convention INTEGER DEFAULT 0, parameter_names TEXT DEFAULT '', parameter_types TEXT DEFAULT '', return_type TEXT DEFAULT '', PRIMARY KEY (name, filepath) )"""; String createGlobals = """ CREATE TABLE IF NOT EXISTS Globals ( filepath TEXT, name TEXT, address TEXT, PRIMARY KEY (name, filepath) )"""; connection.prepareStatement(createFunctions).executeUpdate(); connection.prepareStatement(createImports).executeUpdate(); connection.prepareStatement(createGlobals).executeUpdate(); } // Helper method to find entries by name public List findEntriesByName(String name) throws Exception { ensureConnection(); List results = new ArrayList<>(); try { // Search Functions table findByNameFunctions.setString(1, name); try (ResultSet rs = findByNameFunctions.executeQuery()) { while (rs.next()) { FunctionEntry entry = createEntryFromResultSet(rs); if (entry != null) { results.add(entry); } } } // Search Imports table findByNameImports.setString(1, name); try (ResultSet rs = findByNameImports.executeQuery()) { while (rs.next()) { FunctionEntry entry = createEntryFromResultSet(rs); if (entry != null) { results.add(entry); } } } return results; } catch (SQLException e) { script.println("Error finding entries by name: " + e.getMessage()); throw new Exception("Failed to find entries by name", e); } } // Helper method to find entries by address public List findEntriesByAddress(Address address) throws Exception { ensureConnection(); List results = new ArrayList<>(); String addressStr = address.toString(); try { // Search Functions table findByAddressFunctions.setString(1, addressStr); try (ResultSet rs = findByAddressFunctions.executeQuery()) { while (rs.next()) { FunctionEntry entry = createEntryFromResultSet(rs); if (entry != null) { results.add(entry); } } } // Search Imports table findByAddressImports.setString(1, addressStr); try (ResultSet rs = findByAddressImports.executeQuery()) { while (rs.next()) { FunctionEntry entry = createEntryFromResultSet(rs); if (entry != null) { results.add(entry); } } } return results; } catch (SQLException e) { script.println("Error finding entries by address: " + e.getMessage()); throw new Exception("Failed to find entries by address", e); } } // Helper method to add/update entry (insert or replace based on filename) public void addEntryAt(FunctionEntry entry) throws Exception { ensureConnection(); String relativePath = new File(RemanConfig.INSTANCE.outputDir).toPath() .relativize(entry.file.toPath()).toString().replace('\\', '/'); try { insertOrReplaceFunctions.setString(1, relativePath); insertOrReplaceFunctions.setString(2, entry.name); insertOrReplaceFunctions.setString(3, entry.address.toString()); insertOrReplaceFunctions.setInt(4, entry.type.getValue()); insertOrReplaceFunctions.setInt(5, entry.callingConvention.getValue()); insertOrReplaceFunctions.setString(6, entry.parameterNames); insertOrReplaceFunctions.setString(7, entry.parameterTypes); insertOrReplaceFunctions.setString(8, entry.returnType); insertOrReplaceFunctions.executeUpdate(); script.println("Added/updated entry: " + entry.name + " at " + entry.address + " in " + relativePath + " (calling convention: " + entry.callingConvention + ", return type: " + entry.returnType + ")"); } catch (SQLException e) { script.println("Error adding entry: " + e.getMessage()); throw new Exception("Failed to add entry", e); } } // Helper method to add import entry public void addImportEntry(FunctionEntry entry) throws Exception { ensureConnection(); String relativePath = new File(RemanConfig.INSTANCE.outputDir).toPath() .relativize(entry.file.toPath()).toString().replace('\\', '/'); try { insertOrReplaceImports.setString(1, relativePath); insertOrReplaceImports.setString(2, entry.name); insertOrReplaceImports.setString(3, entry.address.toString()); insertOrReplaceImports.setInt(4, entry.type.getValue()); insertOrReplaceImports.setInt(5, entry.callingConvention.getValue()); insertOrReplaceImports.setString(6, entry.parameterNames); insertOrReplaceImports.setString(7, entry.parameterTypes); insertOrReplaceImports.setString(8, entry.returnType); insertOrReplaceImports.executeUpdate(); script.println("Added/updated import entry: " + entry.name + " at " + entry.address + " in " + relativePath + " (calling convention: " + entry.callingConvention + ", return type: " + entry.returnType + ")"); } catch (SQLException e) { script.println("Error adding import entry: " + e.getMessage()); throw new Exception("Failed to add import entry", e); } } // Helper method to remove entry by file path public void removeEntryAt(String filePath) throws Exception { ensureConnection(); String relativePath = new File(RemanConfig.INSTANCE.outputDir).toPath() .relativize(new File(filePath).toPath()).toString().replace('\\', '/'); try { deleteByFilepathFunctions.setString(1, relativePath); int deletedCount = deleteByFilepathFunctions.executeUpdate(); deleteByFilepathImports.setString(1, relativePath); deletedCount += deleteByFilepathImports.executeUpdate(); script.println("Removed " + deletedCount + " entries for file: " + relativePath); } catch (SQLException e) { script.println("Error removing entries: " + e.getMessage()); throw new Exception("Failed to remove entries", e); } } public void add(FunctionEntry entry) throws Exception { // Add entry directly to database addEntryAt(entry); } public void applyDefaultFilters(boolean rebuildAllGlobals) throws Exception { GlobalDumper globalDumper = new GlobalDumper(script, this); FunctionDumper dumper = new FunctionDumper(script, this, globalDumper); if (rebuildAllGlobals) { globalDumper.removeGlobalManifest(); } boolean madeAnyChanges = false; // Load all entries from database List entries = loadAllEntries(); // Create a hash map to store symbol names Map symbolNames = new HashMap<>(); Map exportedFunctionNames = new HashMap<>(); for (FunctionEntry entry : entries) { Function function = script.getFunctionAt(entry.address); if (function != null) { boolean isAuto = entry.type == Type.Auto; boolean isFix = entry.type == Type.Fix; // Get the actual symbol name and store it in the hash map String symbolName = function.getName(); symbolNames.put(entry.address, symbolName); if (isAuto && !exportedFunctionNames.containsKey(entry.name)) { exportedFunctionNames.put(entry.name, entry.file); } else if (isFix) { exportedFunctionNames.replace(entry.name, entry.file); } } } // Print the number of symbol names collected script.println("Collected " + symbolNames.size() + " symbol names"); boolean dryMode = false; HashSet functionsToRegenerate = new HashSet<>(); Iterator iterator = entries.iterator(); while (iterator.hasNext()) { FunctionEntry entry = iterator.next(); Function function = script.getFunctionAt(entry.address); boolean pendingDelete = false; boolean pendingRegenerate = false; if (rebuildAllGlobals) { pendingRegenerate = true; } // Remove CRT and other blacklisted functions if (function == null || !dumper.isValidFunction(function)) { // Remove the file if (entry.file != null && entry.file.exists()) { script.println("Removed file: " + entry.file.getAbsolutePath()); pendingDelete = true; } // Remove entry from the list script.println("Removed invalid function entry: " + entry.name + " at " + entry.address); function = null; } // Check if symbol name matches the symbol name parsed from the file if (function != null) { String actualSymbolName = symbolNames.get(entry.address); if (actualSymbolName == null) { throw new Exception( "Symbol name not found for function at " + entry.address + " in file " + entry.file.getAbsolutePath()); } if (actualSymbolName != null && !actualSymbolName.equals(entry.name)) { File fnExportedFile = exportedFunctionNames.get(entry.name); if (fnExportedFile != null && fnExportedFile != entry.file) { // Already exists elsewhere, so remove this file script.println("Removing duplicate function: " + entry.name + " at " + entry.address + " overridden by " + fnExportedFile); pendingDelete = true; } else { // Regeneral this function script.println("Symbol name mismatch for function at " + entry.address + ": " + "File name: " + entry.name + ", Actual symbol: " + actualSymbolName); entry.name = actualSymbolName; // Update the entry name to match the actual symbol pendingRegenerate = true; } } entry.name = actualSymbolName; // Update the entry name to match the actual symbol madeAnyChanges = true; } if (pendingDelete) { iterator.remove(); if (!dryMode) { entry.file.delete(); // Remove from database removeEntryAt(entry.file.getAbsolutePath()); madeAnyChanges = true; } } else if (pendingRegenerate && entry.type != Type.Stub) { if (!dryMode) { functionsToRegenerate.add(function); // Update entry in database with corrected name addEntryAt(entry); madeAnyChanges = true; } } } for (Function function : functionsToRegenerate) { script.println("Regenerating function: " + function.getName() + " at " + function.getEntryPoint()); dumper.dump(function); } if (madeAnyChanges) { // Update CMake timestamp RemanConfig.INSTANCE.touchCMakeTimestamp(); globalDumper.dumpGlobals(); globalDumper.saveGlobalManifest(); TypeDumper typeDumper = new TypeDumper(script); typeDumper.run(); } } // Global-specific methods public List loadAllGlobals() throws Exception { ensureConnection(); List globals = new ArrayList<>(); try { try (ResultSet rs = loadAllGlobals.executeQuery()) { while (rs.next()) { GlobalEntry entry = createGlobalEntryFromResultSet(rs); if (entry != null) { globals.add(entry); } } } script.println("Loaded " + globals.size() + " global entries from database"); return globals; } catch (SQLException e) { script.println("Error loading globals: " + e.getMessage()); throw new Exception("Failed to load globals", e); } } private GlobalEntry createGlobalEntryFromResultSet(ResultSet rs) throws SQLException { String filepath = rs.getString("filepath"); String name = rs.getString("name"); String addressStr = rs.getString("address"); if (addressStr != null && !addressStr.isEmpty()) { Address address = script.getCurrentProgram().getAddressFactory().getAddress(addressStr); File file = new File(RemanConfig.INSTANCE.outputDir, filepath); return new GlobalEntry(address, name, file); } return null; } public List findGlobalsByName(String name) throws Exception { ensureConnection(); List results = new ArrayList<>(); try { findByNameGlobals.setString(1, name); try (ResultSet rs = findByNameGlobals.executeQuery()) { while (rs.next()) { GlobalEntry entry = createGlobalEntryFromResultSet(rs); if (entry != null) { results.add(entry); } } } return results; } catch (SQLException e) { script.println("Error finding globals by name: " + e.getMessage()); throw new Exception("Failed to find globals by name", e); } } public List findGlobalsByAddress(Address address) throws Exception { ensureConnection(); List results = new ArrayList<>(); String addressStr = address.toString(); try { findByAddressGlobals.setString(1, addressStr); try (ResultSet rs = findByAddressGlobals.executeQuery()) { while (rs.next()) { GlobalEntry entry = createGlobalEntryFromResultSet(rs); if (entry != null) { results.add(entry); } } } return results; } catch (SQLException e) { script.println("Error finding globals by address: " + e.getMessage()); throw new Exception("Failed to find globals by address", e); } } public void addGlobal(Address address, String name) throws Exception { ensureConnection(); String filepath = RemanConfig.GLOBAL_H_FILE; // Default filepath for globals String addressStr = address.toString(); try { insertOrReplaceGlobals.setString(1, filepath); insertOrReplaceGlobals.setString(2, name); insertOrReplaceGlobals.setString(3, addressStr); insertOrReplaceGlobals.executeUpdate(); script.println("Added/updated global: " + name + " at " + address); } catch (SQLException e) { script.println("Error adding global: " + e.getMessage()); throw new Exception("Failed to add global", e); } } public void removeGlobalsByFilepath(String filePath) throws Exception { ensureConnection(); String relativePath; // Check if filePath is already a relative path or just a filename File inputFile = new File(filePath); if (inputFile.isAbsolute()) { // Convert absolute path to relative try { relativePath = new File(RemanConfig.INSTANCE.outputDir).toPath() .relativize(inputFile.toPath()).toString().replace('\\', '/'); } catch (IllegalArgumentException e) { // Fallback if paths can't be relativized (different drives, etc.) script.println("Warning: Could not relativize path: " + filePath + ", using as-is"); relativePath = filePath.replace('\\', '/'); } } else { // Already relative or just a filename, use as-is relativePath = filePath.replace('\\', '/'); } try { deleteByFilepathGlobals.setString(1, relativePath); int deletedCount = deleteByFilepathGlobals.executeUpdate(); script.println("Removed " + deletedCount + " global entries for file: " + relativePath); } catch (SQLException e) { script.println("Error removing global entries: " + e.getMessage()); throw new Exception("Failed to remove global entries", e); } } @Override public void close() throws Exception { this.disconnect(); } }