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()