ida-plugin/server_func.py

import argparse
import json
import sys
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from threading import Lock

PROJECT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT))

import capstone
import laya

from laya_x64 import normalize, Analyzer
from build_labeling import (ALGO_OPTIONS, CATEGORY_OPTIONS, BUG_OPTIONS, ALGO_Q, CAT_Q, VULN_Q, BUG_Q)

DEFAULT_MODEL = str(PROJECT / "models" / "laya-x64-func")
DISASM = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)

# Base v2 questions. If a later model (v3) is loaded, the extra questions are added at startup
# only if the model was fine-tuned with them (detected by trying a probe predict).
BASE_QUESTIONS = {
    "algorithm": {"type": "choice", "instructions": ALGO_Q, "criteria": {o: None for o in ALGO_OPTIONS}},
    "category": {"type": "choice", "instructions": CAT_Q, "criteria": {o: None for o in CATEGORY_OPTIONS}},
    "vulnerable": {"type": "noul", "instructions": VULN_Q},
    "bug_type": {"type": "choice", "instructions": BUG_Q, "criteria": {o: None for o in BUG_OPTIONS}},
}

try:
    from build_libs import LIBRARY_OPTIONS, PURPOSE_OPTIONS
    from build_labeling_v3 import LIB_Q, PURPOSE_Q, COMPILER_Q, COMPILER_OPTIONS
    V3_QUESTIONS = {
        "library": {"type": "choice", "instructions": LIB_Q, "criteria": {o: None for o in LIBRARY_OPTIONS}},
        "purpose": {"type": "choice", "instructions": PURPOSE_Q, "criteria": {o: None for o in PURPOSE_OPTIONS}},
        "compiler": {"type": "choice", "instructions": COMPILER_Q, "criteria": {o: None for o in COMPILER_OPTIONS}},
    }
except Exception:
    V3_QUESTIONS = {}

try:
    import api_knowledge
    import api_docs
    V4_QUESTIONS = {f"tag_{t}": {"type": "noul", "instructions": api_knowledge.TAG_QUESTIONS[t]}
                    for t in api_knowledge.TAGS}
except Exception:
    api_knowledge = None
    api_docs = None
    V4_QUESTIONS = {}


def disasm_function(insn_hexes):
    lines = []
    for h in insn_hexes:
        for i in DISASM.disasm(bytes.fromhex(h), 0x1000):
            lines.append(normalize(i.mnemonic, i.op_str))
    return "\n".join(lines)


def summarize(answers):
    out = {}
    for qid, a in answers.items():
        if a["type"] == "noul":
            out[qid] = {"value": a["noul"] >= 0.5, "conf": round(max(a["noul"], 1 - a["noul"]), 3)}
        else:
            out[qid] = {"value": a["choice"], "conf": round(a["probabilities"][a["choice"]], 3)}
    return out


def disasm_lines(insn_hexes):
    """Return [(asm_line, None)] for each provided instruction's bytes."""
    lines = []
    for h in insn_hexes:
        for i in DISASM.disasm(bytes.fromhex(h), 0x1000):
            lines.append(normalize(i.mnemonic, i.op_str))
    return lines


def make_handler(agent, questions, max_len, base_agent=None, instr_analyzer=None):
    lock = Lock()

    class Handler(BaseHTTPRequestHandler):
        def _reply(self, code, payload):
            body = json.dumps(payload).encode()
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            self._reply(200 if self.path == "/health" else 404,
                        {"ok": self.path == "/health", "questions": list(questions)})

        def do_POST(self):
            try:
                req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))))
            except (ValueError, TypeError):
                return self._reply(400, {"error": "invalid JSON"})
            if self.path == "/explain":
                return self.do_explain(req)
            if self.path != "/label":
                return self._reply(404, {"error": "unknown route"})
            results = []
            for fn in req.get("functions", []):
                text = fn.get("asm") or disasm_function(fn.get("insns", []))
                if not text.strip():
                    results.append({"ea": fn.get("ea"), "error": "empty"})
                    continue
                # Cap the function text to what the model can actually tokenize.
                # ~320 lines comfortably exceeds 1024 tokens; anything longer would only
                # be truncated by the tokenizer anyway (and emit a "sequence too long"
                # warning first).
                text = "\n".join(text.splitlines()[:320])
                state = {"function": text}
                # Enrich with IAT imports + Microsoft one-liner docs for v4 behavior questions
                apis = fn.get("apis") or []
                if apis:
                    state["apis"] = apis[:12]
                    if api_docs is not None:
                        docs = {a: api_docs.short_doc(a) for a in state["apis"]
                                if api_docs.short_doc(a)}
                        if docs:
                            state["api_docs"] = docs
                calls = fn.get("calls") or []
                if calls:
                    state["calls"] = calls[:16]
                constants = fn.get("constants") or []
                if constants:
                    state["constants"] = constants[:12]
                with lock:
                    ans = agent.predict(state, questions, max_len=max_len)["answers"]
                    row = {"ea": fn.get("ea"), "name": fn.get("name"), "labels": summarize(ans)}
                    if base_agent is not None:
                        base_ans = base_agent.predict(state, questions, max_len=max_len)["answers"]
                        row["labels_base"] = summarize(base_ans)
                results.append(row)
            self._reply(200, {"results": results})

        def do_explain(self, req):
            lines = req.get("asm")
            lines = lines.splitlines() if isinstance(lines, str) else disasm_lines(req.get("insns", []))
            lines = [l for l in lines if l.strip()]
            if not lines:
                return self._reply(400, {"error": "no instructions"})
            per_insn = []
            if instr_analyzer is not None:
                with lock:
                    for line in lines:
                        try:
                            summary, _ = instr_analyzer.summary(line)
                        except Exception:
                            summary = ""
                        per_insn.append({"asm": line, "what": summary})
            else:
                per_insn = [{"asm": line, "what": ""} for line in lines]
            block = None
            if len(lines) > 1:
                with lock:
                    ans = agent.predict({"function": "\n".join(lines)}, questions, max_len=max_len)["answers"]
                    block = summarize(ans)
            self._reply(200, {"instructions": per_insn, "block": block})

        def log_message(self, *a):
            pass

    return Handler


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8766)
    ap.add_argument("--max-len", type=int, default=1024)
    ap.add_argument("--device", default=None)
    ap.add_argument("--questions", choices=["v2", "v3", "v4", "v5"], default="v2",
                    help="v2 = algorithm/category/vulnerable/bug_type; v3 adds library/purpose/compiler; "
                         "v4/v5 also add 28 WinAPI behavior tags")
    ap.add_argument("--compare", action="store_true",
                    help="also load the untrained base model and return its answers for side-by-side comparison")
    ap.add_argument("--base-model", default="convaiinnovations/laya",
                    help="the untrained model to compare against when --compare is set")
    ap.add_argument("--instr-model", default=str(PROJECT / "models" / "laya-x64"),
                    help="the per-instruction (v1) model used by /explain; set to '' to disable")
    args = ap.parse_args()

    agent = laya.load(args.model, device=args.device)
    questions = dict(BASE_QUESTIONS)
    if args.questions in ("v3", "v4", "v5") and V3_QUESTIONS:
        questions.update(V3_QUESTIONS)
    if args.questions in ("v4", "v5") and V4_QUESTIONS:
        questions.update(V4_QUESTIONS)
    base_agent = None
    if args.compare:
        print(f"loading untrained base model for comparison: {args.base_model}")
        base_agent = laya.load(args.base_model, device=args.device)
    instr_analyzer = None
    if args.instr_model and Path(args.instr_model).exists():
        print(f"loading per-instruction (v1) model for /explain: {args.instr_model}")
        instr_analyzer = Analyzer(args.instr_model, device=args.device)
    print(f"laya-func server on http://{args.host}:{args.port}  model={args.model}"
          f"{'  +compare vs ' + args.base_model if base_agent else ''}"
          f"{'  +explain' if instr_analyzer else ''}")
    print(f"questions: {list(questions)}")
    ThreadingHTTPServer((args.host, args.port),
                        make_handler(agent, questions, args.max_len, base_agent, instr_analyzer)).serve_forever()


if __name__ == "__main__":
    main()