// Exports compact Ghidra facts for MXAccess/LMX/NMX reverse-engineering.
// The output is intentionally metadata-oriented: functions, references,
// imports, strings, and call relationships, not full proprietary decompiled
// source listings.

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Data;
import ghidra.program.model.listing.DataIterator;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.listing.InstructionIterator;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.symbol.Reference;
import ghidra.program.model.symbol.Symbol;
import ghidra.program.model.symbol.SymbolIterator;
import ghidra.program.model.symbol.SymbolTable;

import java.io.File;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class MxNmxExport extends GhidraScript {
    private final List<String> interestingStringNeedles = Arrays.asList(
        "CLMXProxyServer::Write",
        "CLMXProxyServer::Advise",
        "CLMXProxyServer::AdviseSupervisory",
        "CLMXProxyServer::Register",
        "CLMXProxyServer::AddItem",
        "Fire_OnWriteComplete",
        "Fire_OnDataChange",
        "RemoteWrite",
        "WriteSecured",
        "WriteVerified",
        "PrebindReference",
        "SupervisoryRegisterPreboundReference",
        "UserRegisterPreboundReference",
        "TransferData",
        "DataReceived",
        "StatusReceived",
        "PutRequest",
        "GetResponse",
        "Nmx",
        "Lmx",
        "MX_E_",
        "MxSecurity",
        "MxSource",
        "MxCategory",
        "IDataClient",
        "PublishWriteComplete",
        "Write2",
        "RegisterItems",
        "socket",
        "WSASend",
        "WSARecv",
        "send",
        "recv",
        "Ndr",
        "RPC"
    );

    private final List<String> interestingCallNeedles = Arrays.asList(
        "send",
        "recv",
        "WSASend",
        "WSARecv",
        "connect",
        "bind",
        "listen",
        "accept",
        "closesocket",
        "Ndr",
        "Rpc",
        "CoCreateInstance",
        "CoGetClassObject",
        "QueryInterface",
        "Variant",
        "SafeArray",
        "SysAllocString",
        "SysFreeString",
        "memcpy",
        "memmove",
        "memset"
    );

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();
        File outDir = new File(args.length > 0 ? args[0] : "analysis/ghidra/exports");
        outDir.mkdirs();

        String baseName = sanitize(currentProgram.getName());
        writeProgramMarkdown(new File(outDir, baseName + ".ghidra.md"));
        writeFunctionsTsv(new File(outDir, baseName + ".functions.tsv"));
        writeStringRefsTsv(new File(outDir, baseName + ".string-refs.tsv"));
        writeCallRefsTsv(new File(outDir, baseName + ".call-refs.tsv"));
        println("Wrote MX/NMX Ghidra export for " + currentProgram.getName() + " to " + outDir.getAbsolutePath());
    }

    private void writeProgramMarkdown(File outFile) throws Exception {
        PrintWriter out = new PrintWriter(new FileWriter(outFile));
        out.println("# " + currentProgram.getName());
        out.println();
        out.println("## Program");
        out.println();
        out.println("- Language: `" + currentProgram.getLanguageID() + "`");
        out.println("- Compiler spec: `" + currentProgram.getCompilerSpec().getCompilerSpecID() + "`");
        out.println("- Image base: `" + currentProgram.getImageBase() + "`");
        out.println("- Executable format: `" + currentProgram.getExecutableFormat() + "`");
        out.println();
        writeMemoryBlocks(out);
        writeExternalImports(out);
        writeExports(out);
        writeInterestingStringSummary(out);
        writeInterestingCallerSummary(out);
        out.close();
    }

    private void writeMemoryBlocks(PrintWriter out) {
        out.println("## Memory Blocks");
        out.println();
        out.println("| Name | Start | End | Size | R | W | X |");
        out.println("| --- | ---: | ---: | ---: | :---: | :---: | :---: |");
        for (MemoryBlock block : currentProgram.getMemory().getBlocks()) {
            out.println("| `" + escape(block.getName()) + "` | `" + block.getStart() + "` | `" + block.getEnd() + "` | " +
                block.getSize() + " | " + yn(block.isRead()) + " | " + yn(block.isWrite()) + " | " + yn(block.isExecute()) + " |");
        }
        out.println();
    }

    private void writeExternalImports(PrintWriter out) {
        out.println("## External Imports");
        out.println();
        List<String> imports = new ArrayList<String>();
        SymbolIterator it = currentProgram.getSymbolTable().getExternalSymbols();
        while (it.hasNext()) {
            imports.add(it.next().getName(true));
        }
        Collections.sort(imports);
        for (String name : imports) {
            out.println("- `" + escape(name) + "`");
        }
        out.println();
    }

    private void writeExports(PrintWriter out) {
        out.println("## Exports and Globals");
        out.println();
        out.println("| Name | Address | Function |");
        out.println("| --- | ---: | --- |");
        SymbolIterator symbols = currentProgram.getSymbolTable().getSymbolIterator(true);
        while (symbols.hasNext()) {
            Symbol symbol = symbols.next();
            if (symbol.isExternal() || !symbol.isGlobal()) {
                continue;
            }
            String name = symbol.getName();
            if (name.startsWith("FUN_") || name.startsWith("DAT_") || name.startsWith("LAB_")) {
                continue;
            }
            Function function = getFunctionContaining(symbol.getAddress());
            out.println("| `" + escape(name) + "` | `" + symbol.getAddress() + "` | `" +
                (function == null ? "" : escape(function.getName())) + "` |");
        }
        out.println();
    }

    private void writeInterestingStringSummary(PrintWriter out) {
        out.println("## Interesting Strings and Referencing Functions");
        out.println();
        out.println("| Address | String | Referencing Functions |");
        out.println("| ---: | --- | --- |");
        for (StringRecord record : collectInterestingStrings()) {
            out.println("| `" + record.address + "` | `" + escape(record.value) + "` | `" + escape(join(record.functions, "`, `")) + "` |");
        }
        out.println();
    }

    private void writeInterestingCallerSummary(PrintWriter out) {
        out.println("## Interesting API Callers");
        out.println();
        out.println("| Caller | Entry | Call Targets |");
        out.println("| --- | ---: | --- |");
        List<FunctionCalls> calls = collectInterestingFunctionCalls();
        Collections.sort(calls, new Comparator<FunctionCalls>() {
            public int compare(FunctionCalls a, FunctionCalls b) {
                return a.function.getEntryPoint().compareTo(b.function.getEntryPoint());
            }
        });
        for (FunctionCalls item : calls) {
            out.println("| `" + escape(item.function.getName()) + "` | `" + item.function.getEntryPoint() + "` | `" +
                escape(join(item.targets, "`, `")) + "` |");
        }
        out.println();
    }

    private void writeFunctionsTsv(File outFile) throws Exception {
        PrintWriter out = new PrintWriter(new FileWriter(outFile));
        out.println("entry\tname\tsignature\tbody_size\tcall_count\tinteresting_calls");
        FunctionIterator functions = currentProgram.getFunctionManager().getFunctions(true);
        while (functions.hasNext()) {
            Function function = functions.next();
            Set<String> targets = directCallTargets(function);
            Set<String> interesting = filterInterestingCalls(targets);
            out.println(function.getEntryPoint() + "\t" + tsv(function.getName()) + "\t" +
                tsv(function.getSignature().getPrototypeString()) + "\t" + function.getBody().getNumAddresses() + "\t" +
                targets.size() + "\t" + tsv(join(new ArrayList<String>(interesting), ";")));
        }
        out.close();
    }

    private void writeStringRefsTsv(File outFile) throws Exception {
        PrintWriter out = new PrintWriter(new FileWriter(outFile));
        out.println("string_address\tstring_value\tref_from\tref_function");
        for (StringRecord record : collectInterestingStrings()) {
            for (String ref : record.references) {
                out.println(record.address + "\t" + tsv(record.value) + "\t" + tsv(ref) + "\t" + tsv(record.refFunctionByAddress.get(ref)));
            }
        }
        out.close();
    }

    private void writeCallRefsTsv(File outFile) throws Exception {
        PrintWriter out = new PrintWriter(new FileWriter(outFile));
        out.println("caller_entry\tcaller_name\tcall_address\ttarget");
        FunctionIterator functions = currentProgram.getFunctionManager().getFunctions(true);
        while (functions.hasNext()) {
            Function function = functions.next();
            InstructionIterator instructions = currentProgram.getListing().getInstructions(function.getBody(), true);
            while (instructions.hasNext()) {
                Instruction instruction = instructions.next();
                for (Reference ref : instruction.getReferencesFrom()) {
                    if (!ref.getReferenceType().isCall()) {
                        continue;
                    }
                    String target = callTargetName(ref.getToAddress());
                    if (!isInterestingCall(target)) {
                        continue;
                    }
                    out.println(function.getEntryPoint() + "\t" + tsv(function.getName()) + "\t" +
                        instruction.getAddress() + "\t" + tsv(target));
                }
            }
        }
        out.close();
    }

    private List<StringRecord> collectInterestingStrings() {
        List<StringRecord> records = new ArrayList<StringRecord>();
        Set<String> seenAt = new HashSet<String>();
        DataIterator data = currentProgram.getListing().getDefinedData(true);
        while (data.hasNext()) {
            if (monitor.isCancelled()) {
                break;
            }
            Data item = data.next();
            Object valueObject = null;
            try {
                valueObject = item.getValue();
            } catch (Exception e) {
                continue;
            }
            if (valueObject == null) {
                continue;
            }
            String value = valueObject.toString();
            if (value.length() < 4 || !isInterestingString(value)) {
                continue;
            }
            String address = item.getAddress().toString();
            if (seenAt.contains(address)) {
                continue;
            }
            seenAt.add(address);
            StringRecord record = new StringRecord(address, value);
            for (Reference ref : getReferencesTo(item.getAddress())) {
                Address from = ref.getFromAddress();
                Function function = getFunctionContaining(from);
                String functionName = function == null ? "" : function.getName() + "@" + function.getEntryPoint();
                record.references.add(from.toString());
                record.refFunctionByAddress.put(from.toString(), functionName);
                if (functionName.length() > 0) {
                    record.functions.add(functionName);
                }
            }
            records.add(record);
        }
        Collections.sort(records, new Comparator<StringRecord>() {
            public int compare(StringRecord a, StringRecord b) {
                return a.address.compareTo(b.address);
            }
        });
        return records;
    }

    private List<FunctionCalls> collectInterestingFunctionCalls() {
        List<FunctionCalls> out = new ArrayList<FunctionCalls>();
        FunctionIterator functions = currentProgram.getFunctionManager().getFunctions(true);
        while (functions.hasNext()) {
            Function function = functions.next();
            Set<String> targets = filterInterestingCalls(directCallTargets(function));
            if (!targets.isEmpty()) {
                out.add(new FunctionCalls(function, new ArrayList<String>(targets)));
            }
        }
        return out;
    }

    private Set<String> directCallTargets(Function function) {
        Set<String> targets = new LinkedHashSet<String>();
        InstructionIterator instructions = currentProgram.getListing().getInstructions(function.getBody(), true);
        while (instructions.hasNext()) {
            Instruction instruction = instructions.next();
            for (Reference ref : instruction.getReferencesFrom()) {
                if (!ref.getReferenceType().isCall()) {
                    continue;
                }
                targets.add(callTargetName(ref.getToAddress()));
            }
        }
        return targets;
    }

    private String callTargetName(Address address) {
        Function function = getFunctionAt(address);
        if (function != null) {
            return function.getName();
        }
        Symbol symbol = getSymbolAt(address);
        if (symbol != null) {
            return symbol.getName(true);
        }
        SymbolTable symbols = currentProgram.getSymbolTable();
        Symbol primary = symbols.getPrimarySymbol(address);
        if (primary != null) {
            return primary.getName(true);
        }
        return address.toString();
    }

    private Set<String> filterInterestingCalls(Set<String> targets) {
        Set<String> out = new LinkedHashSet<String>();
        for (String target : targets) {
            if (isInterestingCall(target)) {
                out.add(target);
            }
        }
        return out;
    }

    private boolean isInterestingString(String value) {
        String lower = value.toLowerCase();
        for (String needle : interestingStringNeedles) {
            if (lower.contains(needle.toLowerCase())) {
                return true;
            }
        }
        return false;
    }

    private boolean isInterestingCall(String value) {
        String lower = value.toLowerCase();
        for (String needle : interestingCallNeedles) {
            if (lower.contains(needle.toLowerCase())) {
                return true;
            }
        }
        return false;
    }

    private String sanitize(String value) {
        return value.replaceAll("[^A-Za-z0-9_.-]", "_");
    }

    private String yn(boolean value) {
        return value ? "Y" : "";
    }

    private String join(List<String> items, String separator) {
        Collections.sort(items);
        StringBuilder builder = new StringBuilder();
        for (int i = 0; i < items.size(); i++) {
            if (i > 0) {
                builder.append(separator);
            }
            builder.append(items.get(i));
        }
        return builder.toString();
    }

    private String escape(String value) {
        if (value == null) {
            return "";
        }
        return value.replace("\\", "\\\\").replace("`", "\\`").replace("|", "\\|").replace("\r", " ").replace("\n", " ");
    }

    private String tsv(String value) {
        if (value == null) {
            return "";
        }
        return value.replace("\t", " ").replace("\r", " ").replace("\n", " ");
    }

    private static class StringRecord {
        String address;
        String value;
        List<String> references = new ArrayList<String>();
        Map<String, String> refFunctionByAddress = new HashMap<String, String>();
        List<String> functions = new ArrayList<String>();

        StringRecord(String address, String value) {
            this.address = address;
            this.value = value;
        }
    }

    private static class FunctionCalls {
        Function function;
        List<String> targets;

        FunctionCalls(Function function, List<String> targets) {
            this.function = function;
            this.targets = targets;
        }
    }
}