inference.py

#!/usr/bin/env python3
"""Run the Grounded Sigma Compiler adapter and verify its output."""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Any

from mlx_lm import generate, load

DEFAULT_MODEL = "mlx-community/Qwen2.5-7B-Instruct-4bit"
SYSTEM_PROMPT = (
    "You are a defensive Sigma rule compiler. Convert the supplied grounded "
    "detection specification into one valid JSON object. Preserve the supplied "
    "telemetry fields, values, condition, log source and ATT&CK mappings exactly. "
    "Never invent, remove, reinterpret or replace detection information. Return "
    "JSON only and set requires_validation to true."
)
SOURCE_FIELDS = (
    "title",
    "description",
    "logsource",
    "detection",
    "attack_techniques",
    "false_positives",
    "severity",
)
OUTPUT_FIELDS = SOURCE_FIELDS + ("requires_validation",)


def parse_args() -> argparse.Namespace:
    script_dir = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(
        description="Compile a grounded detection specification into Sigma JSON."
    )
    parser.add_argument("--input", required=True, type=Path, help="Input JSON file")
    parser.add_argument(
        "--adapter",
        type=Path,
        default=script_dir / "adapter",
        help="Adapter directory (default: ./adapter beside this script)",
    )
    parser.add_argument("--model", default=DEFAULT_MODEL, help="Base MLX model")
    parser.add_argument("--output", type=Path, help="Write verified JSON here")
    parser.add_argument("--max-tokens", type=int, default=2048)
    return parser.parse_args()


def load_spec(path: Path) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError(f"Cannot read valid JSON from {path}: {exc}") from exc
    if not isinstance(value, dict):
        raise ValueError("Input must be one JSON object.")
    missing = [field for field in SOURCE_FIELDS if field not in value]
    extra = sorted(set(value) - set(SOURCE_FIELDS))
    if missing:
        raise ValueError(f"Missing required fields: {', '.join(missing)}")
    if extra:
        raise ValueError(f"Unexpected input fields: {', '.join(extra)}")
    if not isinstance(value["logsource"], dict) or not value["logsource"]:
        raise ValueError("logsource must be a non-empty object.")
    if not isinstance(value["detection"], dict) or not value["detection"]:
        raise ValueError("detection must be a non-empty object.")
    if not isinstance(value["attack_techniques"], list):
        raise ValueError("attack_techniques must be an array.")
    if not isinstance(value["false_positives"], list):
        raise ValueError("false_positives must be an array.")
    return value


def parse_generated_json(text: str) -> dict[str, Any]:
    cleaned = text.strip()
    fenced = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", cleaned, re.DOTALL)
    if fenced:
        cleaned = fenced.group(1)
    decoder = json.JSONDecoder()
    start = cleaned.find("{")
    if start < 0:
        raise ValueError("Model output does not contain a JSON object.")
    try:
        value, end = decoder.raw_decode(cleaned[start:])
    except json.JSONDecodeError as exc:
        raise ValueError(f"Model returned invalid JSON: {exc}") from exc
    if cleaned[start + end :].strip():
        raise ValueError("Model returned extra text after the JSON object.")
    if not isinstance(value, dict):
        raise ValueError("Model output must be one JSON object.")
    return value


def verify_output(spec: dict[str, Any], output: dict[str, Any]) -> None:
    missing = [field for field in OUTPUT_FIELDS if field not in output]
    extra = sorted(set(output) - set(OUTPUT_FIELDS))
    errors: list[str] = []
    if missing:
        errors.append(f"missing fields: {', '.join(missing)}")
    if extra:
        errors.append(f"unexpected fields: {', '.join(extra)}")
    for field in SOURCE_FIELDS:
        if field not in output:
            continue
        if field == "description":
            if str(output[field]).strip() != str(spec[field]).strip():
                errors.append("description changed")
        elif output[field] != spec[field]:
            errors.append(f"{field} changed")
    if output.get("requires_validation") is not True:
        errors.append("requires_validation is not true")
    if errors:
        raise ValueError("Untrusted model output rejected: " + "; ".join(errors))


def main() -> int:
    args = parse_args()
    try:
        spec = load_spec(args.input)
        model, tokenizer = load(args.model, adapter_path=str(args.adapter))
        user_prompt = (
            "Compile this grounded detection specification into Sigma JSON:\n"
            + json.dumps(spec, ensure_ascii=False, separators=(",", ":"))
        )
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ]
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        raw = generate(
            model,
            tokenizer,
            prompt=prompt,
            max_tokens=args.max_tokens,
            verbose=False,
        )
        output = parse_generated_json(raw)
        verify_output(spec, output)
        rendered = json.dumps(output, ensure_ascii=False, indent=2) + "\n"
        if args.output:
            args.output.parent.mkdir(parents=True, exist_ok=True)
            args.output.write_text(rendered, encoding="utf-8")
            print(f"Verified output written to {args.output}", file=sys.stderr)
        else:
            print(rendered, end="")
        return 0
    except (ValueError, OSError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())