File size: 2,449 Bytes
cf58573 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | import argparse
import json
from pathlib import Path
import laya
ROOT = Path(__file__).parent
DEFAULT_MODEL = str(ROOT / "models" / "laya-x64-func")
from build_labeling import (ALGO_OPTIONS, CATEGORY_OPTIONS, BUG_OPTIONS, ALGO_Q, CAT_Q, VULN_Q, BUG_Q)
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}},
}
def read_asm(source):
"""source is a path to a text file of assembly, or raw hex bytes to disassemble with Capstone."""
p = Path(source)
if p.exists():
return p.read_text(encoding="utf-8").strip()
import capstone
from laya_x64 import normalize
md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
data = bytes.fromhex(source.replace(" ", ""))
return "\n".join(normalize(i.mnemonic, i.op_str) for i in md.disasm(data, 0x1000))
def main():
ap = argparse.ArgumentParser(description="Ask the v2 model about a whole x86-64 function.")
ap.add_argument("source", help="path to an .asm text file, or a hex byte string to disassemble")
ap.add_argument("--model", default=DEFAULT_MODEL)
ap.add_argument("--max-len", type=int, default=1024)
ap.add_argument("--json", action="store_true")
args = ap.parse_args()
asm = read_asm(args.source)
agent = laya.load(args.model, device="cuda")
result = agent.predict({"function": asm}, QUESTIONS, max_len=args.max_len)["answers"]
n_insns = sum(1 for l in asm.splitlines() if l.strip() and not l.rstrip().endswith(":"))
print(f"\nfunction: {n_insns} instructions\n" + "-" * 40)
if args.json:
print(json.dumps(result, indent=2, default=str)); return
a = result["algorithm"]
print(f"algorithm : {a['choice']:24} ({a['probabilities'][a['choice']]:.0%})")
c = result["category"]
print(f"category : {c['choice']:24} ({c['probabilities'][c['choice']]:.0%})")
v = result["vulnerable"]["noul"]
print(f"vulnerable: {'YES' if v >= 0.5 else 'no':24} ({max(v, 1 - v):.0%})")
if v >= 0.5:
b = result["bug_type"]
print(f"bug_type : {b['choice']:24} ({b['probabilities'][b['choice']]:.0%})")
if __name__ == "__main__":
main()
|