// Import and analyze raw ARM firmware binary
// This script imports a raw binary file with specified ARM processor and load address
// @author MCGhidra
// @category Binary.Import
// @keybinding
// @menupath
// @toolbar

import ghidra.app.script.GhidraScript;
import ghidra.app.util.bin.ByteArrayProvider;
import ghidra.app.util.importer.MessageLog;
import ghidra.app.util.opinion.BinaryLoader;
import ghidra.app.util.opinion.LoadSpec;
import ghidra.framework.model.DomainFile;
import ghidra.framework.model.DomainFolder;
import ghidra.program.model.address.Address;
import ghidra.program.model.lang.LanguageCompilerSpecPair;
import ghidra.program.model.listing.Program;
import ghidra.util.task.TaskMonitor;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.util.List;

public class ImportRawARM extends GhidraScript {

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();

        if (args.length < 2) {
            println("========================================");
            println("  Import Raw ARM Firmware Binary");
            println("========================================");
            println("");
            println("Usage: -postScript ImportRawARM.java <binary_path> <load_address> [<language_id>]");
            println("");
            println("Arguments:");
            println("  binary_path  - Path to raw binary file");
            println("  load_address - Base address in hex (e.g., 0x00000000)");
            println("  language_id  - ARM language ID (default: ARM:LE:32:v5t)");
            println("");
            println("Common ARM Language IDs:");
            println("  ARM:LE:32:v4       - ARMv4 little-endian");
            println("  ARM:LE:32:v5       - ARMv5 little-endian");
            println("  ARM:LE:32:v5t      - ARMv5T little-endian (Thumb)");
            println("  ARM:LE:32:v6       - ARMv6 little-endian");
            println("  ARM:LE:32:v7       - ARMv7 little-endian");
            println("  ARM:LE:32:Cortex   - ARM Cortex");
            println("  ARM:BE:32:v5t      - ARMv5T big-endian");
            println("");
            println("Example:");
            println("  analyzeHeadless /projects MyProject \\");
            println("    -postScript ImportRawARM.java /binaries/firmware.bin 0x00000000 ARM:LE:32:v5t");
            println("");
            return;
        }

        String binaryPath = args[0];
        String loadAddressStr = args[1];
        String languageIDStr = args.length > 2 ? args[2] : "ARM:LE:32:v5t";

        // Parse load address
        long loadAddress;
        if (loadAddressStr.toLowerCase().startsWith("0x")) {
            loadAddress = Long.parseLong(loadAddressStr.substring(2), 16);
        } else {
            loadAddress = Long.parseLong(loadAddressStr, 16);
        }

        File binaryFile = new File(binaryPath);
        if (!binaryFile.exists()) {
            printerr("ERROR: Binary file not found: " + binaryPath);
            return;
        }

        long fileSize = binaryFile.length();

        println("========================================");
        println("  Importing Raw ARM Firmware");
        println("========================================");
        println("  Binary:       " + binaryFile.getName());
        println("  Size:         " + fileSize + " bytes (0x" + Long.toHexString(fileSize) + ")");
        println("  Load Address: 0x" + String.format("%08X", loadAddress));
        println("  Language:     " + languageIDStr);
        println("========================================");
        println("");

        try {
            // Read the binary file
            byte[] bytes = Files.readAllBytes(binaryFile.toPath());
            ByteArrayProvider provider = new ByteArrayProvider(bytes);

            // Parse language string (format: "ARM:LE:32:v5t")
            String[] parts = languageIDStr.split(":");
            if (parts.length < 4) {
                printerr("ERROR: Invalid language ID format. Expected format: ARCH:ENDIAN:SIZE:VARIANT");
                printerr("       Example: ARM:LE:32:v5t");
                return;
            }

            // Create language compiler spec pair
            LanguageCompilerSpecPair lcsPair = new LanguageCompilerSpecPair(languageIDStr, "default");

            // Use BinaryLoader to import
            BinaryLoader loader = new BinaryLoader();

            // Create load spec with our language
            LoadSpec loadSpec = new LoadSpec(loader, loadAddress, lcsPair, false);

            // Get project folder
            DomainFolder rootFolder = state.getProject().getProjectData().getRootFolder();

            String programName = binaryFile.getName();

            println("Creating program: " + programName);
            println("Using language: " + languageIDStr);
            println("");

            // Import using BinaryLoader
            MessageLog log = new MessageLog();
            List<Program> programs = loader.load(
                provider,
                programName,
                rootFolder,
                loadSpec,
                List.of(),  // options
                log,
                this,       // consumer
                monitor
            );

            if (programs == null || programs.isEmpty()) {
                printerr("ERROR: Failed to load binary");
                println("");
                println("Loader messages:");
                println(log.toString());
                return;
            }

            Program program = programs.get(0);

            println("Import successful!");
            println("  Program: " + program.getName());
            println("  Base Address: " + program.getImageBase());
            println("  Memory blocks: " + program.getMemory().getBlocks().length);
            println("");

            // Set entry point at load address
            Address entryAddr = program.getAddressFactory().getDefaultAddressSpace().getAddress(loadAddress);

            int txId = program.startTransaction("Set Entry Point");
            try {
                program.getSymbolTable().addExternalEntryPoint(entryAddr);
                program.getSymbolTable().createLabel(entryAddr, "entry",
                    ghidra.program.model.symbol.SourceType.ANALYSIS);
                program.endTransaction(txId, true);

                println("Entry point set at: " + entryAddr);
                println("");
            } catch (Exception e) {
                program.endTransaction(txId, false);
                printerr("Warning: Could not set entry point: " + e.getMessage());
            }

            // Save the program
            try {
                program.save("ARM firmware import", monitor);
                println("Program saved successfully");
                println("");
            } catch (Exception e) {
                printerr("Warning: Could not save program: " + e.getMessage());
            }

            // Run analysis
            println("Starting auto-analysis...");
            println("(This may take a while for large binaries)");
            println("");

            analyzeAll(program);

            println("");
            println("========================================");
            println("  Import Complete!");
            println("========================================");
            println("  Program: " + program.getName());
            println("  Functions found: " + program.getFunctionManager().getFunctionCount());
            println("  Defined data: " + program.getListing().getNumDefinedData());
            println("========================================");
            println("");

            // Set as current program for subsequent scripts
            state.setCurrentProgram(program);

        } catch (IOException e) {
            printerr("ERROR: Failed to read binary file: " + e.getMessage());
            e.printStackTrace();
        } catch (Exception e) {
            printerr("ERROR: Import failed: " + e.getMessage());
            e.printStackTrace();
        }
    }
}