# -*- coding: utf-8 -*-
#
# Part of ChipBench (AGPLv3). SPDX-FileCopyrightText: 2026 Starting Line Productions LLC
# Query.py -- headless query interface for chat-driven reverse engineering.
# Usage:
#   analyzeHeadless <projDir> <proj> -process <prog> -noanalysis -readOnly \
#       -scriptPath <dir> -postScript Query.py <command> [args...]
#
# Commands: info, funcs, decompile, disasm, xrefs-to, xrefs-from,
#           func-xrefs, strings, segments, data, search, func, mem

from ghidra.app.decompiler import DecompInterface
import jarray

args = getScriptArgs()
if len(args) == 0:
    print("Usage: Query.py <command> [args...]")
    print("Commands: info, funcs, decompile, disasm, xrefs-to, xrefs-from,")
    print("          func-xrefs, strings, segments, data, search, func, mem")
else:
    cmd = args[0]
    rest = args[1:]

    fm = currentProgram.getFunctionManager()
    rm = currentProgram.getReferenceManager()
    mem = currentProgram.getMemory()
    af = currentProgram.getAddressFactory()

    def arg(i, default):
        if i >= len(rest):
            return default
        try:
            return int(rest[i])
        except ValueError:
            return default

    def parseAddr(s):
        try:
            return af.getAddress(s)
        except:
            return None

    def resolveFunc(s):
        a = parseAddr(s)
        if a is not None:
            f = fm.getFunctionAt(a)
            if f is None:
                f = fm.getFunctionContaining(a)
            if f is not None:
                return f
        for f in fm.getFunctions(True):
            if f.getName() == s:
                return f
        return None

    if cmd == "info":
        print("Name:       " + currentProgram.getName())
        print("Language:   " + str(currentProgram.getLanguageID()))
        print("Compiler:   " + str(currentProgram.getCompilerSpec().getCompilerSpecID()))
        print("Executable: " + str(currentProgram.getExecutablePath()))
        print("ImageBase:  " + str(currentProgram.getImageBase()))
        print("Size:       " + str(mem.getSize()) + " bytes")
        print("Functions:  " + str(fm.getFunctionCount()))
        print("Segments:")
        for b in mem.getBlocks():
            r = "R" if b.isRead() else "-"
            w = "W" if b.isWrite() else "-"
            x = "X" if b.isExecute() else "-"
            print("  " + b.getName() + "  " + str(b.getStart()) + "-" + str(b.getEnd()) + "  (" + str(b.getSize()) + "B)  " + r + w + x)

    elif cmd == "funcs":
        off = arg(0, 0)
        lim = arg(1, 200)
        flt = rest[2] if len(rest) > 2 else None
        i = 0
        shown = 0
        for f in fm.getFunctions(True):
            if flt is not None and flt.lower() not in f.getName().lower():
                continue
            if i < off:
                i += 1
                continue
            if shown >= lim:
                break
            print(f.getName() + " @ " + str(f.getEntryPoint()))
            shown += 1
            i += 1
        print("--- " + str(shown) + " shown (offset " + str(off) + ") ---")

    elif cmd == "decompile":
        if len(rest) < 1:
            print("decompile needs a name or address")
        else:
            f = resolveFunc(rest[0])
            if f is None:
                print("No function found for: " + rest[0])
            else:
                d = DecompInterface()
                d.openProgram(currentProgram)
                r = d.decompileFunction(f, 60, monitor)
                if r is not None and r.decompileCompleted():
                    for line in r.getDecompiledFunction().getC().split("\n"):
                        print(line)
                else:
                    print("Decompilation failed")
                d.dispose()

    elif cmd == "disasm":
        if len(rest) < 1:
            print("disasm needs an address")
        else:
            f = resolveFunc(rest[0])
            if f is None:
                print("No function at/containing: " + rest[0])
            else:
                end = f.getBody().getMaxAddress()
                listing = currentProgram.getListing()
                it = listing.getInstructions(f.getEntryPoint(), True)
                from ghidra.program.model.listing import CodeUnit
                while it.hasNext():
                    ins = it.next()
                    if ins.getAddress().compareTo(end) > 0:
                        break
                    cmt = listing.getComment(CodeUnit.EOL_COMMENT, ins.getAddress())
                    line = str(ins.getAddress()) + ": " + ins.toString()
                    if cmt is not None:
                        line += "  ; " + cmt
                    print(line)

    elif cmd == "xrefs-to":
        if len(rest) < 1:
            print("xrefs-to needs an address")
        else:
            lim = arg(1, 200)
            a = parseAddr(rest[0])
            if a is None:
                print("Bad address: " + rest[0])
            else:
                n = 0
                it = rm.getReferencesTo(a)
                while it.hasNext():
                    if n >= lim:
                        break
                    ref = it.next()
                    ff = fm.getFunctionContaining(ref.getFromAddress())
                    extra = " in " + ff.getName() if ff is not None else ""
                    print("From " + str(ref.getFromAddress()) + extra + " [" + ref.getReferenceType().getName() + "]")
                    n += 1
                print("--- " + str(n) + " reference(s) to " + str(a) + " ---")

    elif cmd == "xrefs-from":
        if len(rest) < 1:
            print("xrefs-from needs an address")
        else:
            lim = arg(1, 200)
            a = parseAddr(rest[0])
            if a is None:
                print("Bad address: " + rest[0])
            else:
                refs = rm.getReferencesFrom(a)
                n = 0
                for ref in refs:
                    if n >= lim:
                        break
                    tf = fm.getFunctionAt(ref.getToAddress())
                    extra = " (" + tf.getName() + ")" if tf is not None else ""
                    print("To " + str(ref.getToAddress()) + extra + " [" + ref.getReferenceType().getName() + "]")
                    n += 1
                print("--- " + str(n) + " reference(s) from " + str(a) + " ---")

    elif cmd == "func-xrefs":
        if len(rest) < 1:
            print("func-xrefs needs a function name")
        else:
            lim = arg(1, 200)
            total = 0
            for f in fm.getFunctions(True):
                if f.getName() != rest[0]:
                    continue
                it = rm.getReferencesTo(f.getEntryPoint())
                while it.hasNext():
                    if total >= lim:
                        break
                    ref = it.next()
                    ff = fm.getFunctionContaining(ref.getFromAddress())
                    extra = " in " + ff.getName() if ff is not None else ""
                    print("From " + str(ref.getFromAddress()) + extra + " [" + ref.getReferenceType().getName() + "]")
                    total += 1
            print("--- " + str(total) + " reference(s) to function '" + rest[0] + "' ---")

    elif cmd == "strings":
        flt = rest[0] if len(rest) > 0 else None
        lim = arg(1, 500)
        shown = 0
        it = currentProgram.getListing().getDefinedData(True)
        while it.hasNext() and shown < lim:
            d = it.next()
            tn = d.getDataType().getName().lower()
            if "string" not in tn and "char" not in tn and tn != "unicode":
                continue
            v = str(d.getValue()) if d.getValue() is not None else ""
            if flt is not None and flt.lower() not in v.lower():
                continue
            print(str(d.getAddress()) + ": \"" + v[:120] + "\"")
            shown += 1
        print("--- " + str(shown) + " string(s) ---")

    elif cmd == "segments":
        for b in mem.getBlocks():
            print(b.getName() + ": " + str(b.getStart()) + " - " + str(b.getEnd()) + " (" + str(b.getSize()) + " bytes)")

    elif cmd == "data":
        off = arg(0, 0)
        lim = arg(1, 200)
        i = 0
        shown = 0
        for b in mem.getBlocks():
            it = currentProgram.getListing().getDefinedData(b.getStart(), True)
            while it.hasNext():
                d = it.next()
                if not b.contains(d.getAddress()):
                    break
                if i < off:
                    i += 1
                    continue
                if shown >= lim:
                    break
                lbl = d.getLabel() if d.getLabel() is not None else "(unnamed)"
                print(str(d.getAddress()) + ": " + lbl + " = " + d.getDefaultValueRepresentation())
                shown += 1
                i += 1
        print("--- " + str(shown) + " data item(s) ---")

    elif cmd == "search":
        if len(rest) < 1:
            print("search needs a query")
        else:
            lim = arg(1, 200)
            shown = 0
            for f in fm.getFunctions(True):
                if rest[0].lower() in f.getName().lower():
                    if shown >= lim:
                        break
                    print(f.getName() + " @ " + str(f.getEntryPoint()))
                    shown += 1
            print("--- " + str(shown) + " match(es) ---")

    elif cmd == "func":
        if len(rest) < 1:
            print("func needs an address")
        else:
            f = resolveFunc(rest[0])
            if f is None:
                print("No function at/containing: " + rest[0])
            else:
                print("Name:      " + f.getName())
                print("Entry:     " + str(f.getEntryPoint()))
                print("Signature: " + f.getSignature())
                print("Body:      " + str(f.getBody().getMinAddress()) + " - " + str(f.getBody().getMaxAddress()))
                print("Size:      " + str(f.getBody().getNumAddresses()) + " bytes")

    elif cmd == "mem":
        if len(rest) < 2:
            print("mem needs <addr> <len>")
        else:
            a = parseAddr(rest[0])
            length = arg(1, 64)
            if a is None:
                print("Bad address: " + rest[0])
            else:
                buf = jarray.zeros(length, "b")
                got = mem.getBytes(a, buf)
                ubuf = [(x & 0xFF) for x in buf[:got]]
                hexstr = ""
                ascstr = ""
                for idx in range(got):
                    v = ubuf[idx]
                    hexstr += "%02x " % v
                    ascstr += chr(v) if (32 <= v < 127) else "."
                    if (idx + 1) % 16 == 0:
                        base = a.add(idx - 15)
                        print("%s  %-48s %s" % (base, hexstr, ascstr))
                        hexstr = ""
                        ascstr = ""
                if hexstr:
                    print("%s  %-48s %s" % (a, hexstr, ascstr))

    else:
        print("Unknown command: " + cmd)
        print("Commands: info, funcs, decompile, disasm, xrefs-to, xrefs-from,")
        print("          func-xrefs, strings, segments, data, search, func, mem")