Blame for runtime/python/dibaguard/llm.py

"""LLM bridge: DibaGuard as a tool for any language model - MCP (stdio), OpenAI / Anthropic tool schemas, HTTP.

Tools (the same five everywhere):
  guard_status        engine state: counts, thresholds and where they come from, drift, pressure
  guard_score_flows   score a time-ordered batch of flows (FLOW_FIELDS columns) -> verdict per flow + summary
  guard_score_http    score canonical HTTP request texts (dibaguard.http.canon) -> verdict per request
  guard_calibrate     feed this site's own benign flows / requests (no attack labels) -> site thresholds
  guard_report        incident report for a time window: JSON (dibaguard.report.REPORT_SCHEMA) + text en / fa
  guard_fast_state    what is enforced now (the proxy fast path polls this)      | with --policy (shield layer):
  guard_posture       host exposure audit + remediation                        | guard_policy (show / set),
  guard_playbook      recent-attack playbook and what matched here             | guard_* act through the policy
  guard_db / guard_db_calibrate / guard_db_status   the database tier (query / audit logs)

    python -m dibaguard.llm mcp  --models models            # MCP server over stdio (Claude, Cursor, any MCP client)
    python -m dibaguard.llm http --models models --port 8787 # POST /tools/<name> JSON, GET /tools (schemas)
                                                             # GET /console: the operator console (dibaguard/console.py)
  guard_rpc           JSON-RPC calls of a chain node (method abuse, floods, mempool spam)
Over HTTP, calls that change what the layer enforces (is_control: policy set, verdicts, rule installs, port closing) need
the console token (DIBAGUARD_CONSOLE_TOKEN, or the one printed at startup) and the header X-DibaGuard-Console: 1. Reads
stay open on loopback; traffic feeds that carry sources need the feed token (DIBAGUARD_FEED_TOKEN) or the console token.
No dependency beyond the standard library; the model never sees raw tokens or private data it did not send.
"""
from __future__ import annotations

import json
import sys

import numpy as np

from dibaguard import report as rep_mod
from dibaguard.engine import FLOW_FIELDS, Guard

_FLOWS_SCHEMA = {"type": "object", "description": "columns of equal length, time-ordered",
                 "properties": {k: {"type": "array"} for k in FLOW_FIELDS}, "required": ["t", "src", "dst"]}
TOOLS = [
    {"name": "guard_status", "description": "DibaGuard engine status: traffic counts, active thresholds (site-calibrated or shipped), drift flag, attack pressure.",
     "input_schema": {"type": "object", "properties": {}}},
    {"name": "guard_score_flows", "description": "Score network flows (NetFlow-like records, time-ordered). Returns per-flow attack probability, alarm, family, 0-100 score, tier that decided, per-source risk, and a summary.",
     "input_schema": {"type": "object", "properties": {"flows": _FLOWS_SCHEMA, "max_rows": {"type": "integer", "default": 200}}, "required": ["flows"]}},
    {"name": "guard_score_http", "description": "Score HTTP requests given as canonical text ('method path | q:k=v | h:k=v | b:...'). Optional client ips, hosts, statuses, timestamps make them part of the entity stream.",
     "input_schema": {"type": "object", "properties": {"texts": {"type": "array", "items": {"type": "string"}}, "deepest": {"type": "array", "items": {"type": "integer"}},
                                                       "client": {"type": "array", "items": {"type": "string"}}, "host": {"type": "array", "items": {"type": "string"}},
                                                       "status": {"type": "array", "items": {"type": "integer"}}, "t": {"type": "array", "items": {"type": "number"}}},
                      "required": ["texts"]}},
    {"name": "guard_calibrate", "description": "Give the engine this site's own normal traffic (no attack labels): benign flows and/or benign HTTP request texts. Site thresholds and the HTTP site profile are derived from them.",
     "input_schema": {"type": "object", "properties": {"flows": _FLOWS_SCHEMA, "texts": {"type": "array", "items": {"type": "string"}}}}},
    {"name": "guard_report", "description": "Incident report for a time window (unix seconds; omit for everything logged): level, pressure, top risky entities with evidence and recommended action, recommendations, text in English and Persian.",
     "input_schema": {"type": "object", "properties": {"since": {"type": "number"}, "until": {"type": "number"}, "top": {"type": "integer", "default": 10}}}},
    {"name": "guard_fast_state", "description": "What is enforced right now (blocked / challenged / rate-limited sources with deadlines): the map the proxy fast path applies.",
     "input_schema": {"type": "object", "properties": {}}},
    {"name": "guard_posture", "description": "Audit this host's exposure (listening ports vs policy, services that must never be public, chain RPC exposure); closes what the policy allows, reports the rest.",
     "input_schema": {"type": "object", "properties": {}}},
    {"name": "guard_policy", "description": "Show the active policy (mode, actions per family, allowlist, posture rules) or change fields of it (mode, actions, allow, min_action_score, posture).",
     "input_schema": {"type": "object", "properties": {"set": {"type": "object"}}}},
    {"name": "guard_playbook", "description": "The recent-attack playbook (2024-2026): patterns, how each is seen, family and action; with counts of what was matched here.",
     "input_schema": {"type": "object", "properties": {}}},
    {"name": "guard_feedback_pending", "description": "Alarms awaiting an operator's verdict, grouped per source (newest, highest score first).",
     "input_schema": {"type": "object", "properties": {"limit": {"type": "integer", "default": 50}, "since": {"type": "number"}}}},
    {"name": "guard_feedback", "description": "The operator's verdict on a source's alarms: confirm (take the action now), reject (lift it, the source's traffic becomes benign reference), or benign_source (allowlist it). Clearing a source with hard evidence (a named signature, a dangerous database command, campaign membership) or a burst of clears returns needs_confirmation instead of applying; repeat with force=true only after the human explicitly confirms.",
     "input_schema": {"type": "object", "properties": {"entity": {"type": "string"}, "verdict": {"type": "string", "enum": ["confirm", "reject", "benign_source"]}, "note": {"type": "string"},
                                                       "force": {"type": "boolean", "description": "the human explicitly confirmed clearing hard evidence"}}, "required": ["entity", "verdict"]}},
    {"name": "guard_feedback_stats", "description": "How many alarms the operator has judged, and the alarm precision as they see it.",
     "input_schema": {"type": "object", "properties": {}}},
    {"name": "guard_rules", "description": "Show the operator's own rules, or add natural-language rules (fa/en) that the layer enforces from the first second (interpreted to a precise form, with the interpretation returned for confirmation).",
     "input_schema": {"type": "object", "properties": {"add": {"type": "array", "items": {"type": "string"}}, "interpret_only": {"type": "boolean"}}}},
    {"name": "guard_db", "description": "Score database query / audit events (t, engine, client, db, user, query, status ok|auth_fail, rows, bytes) or raw log lines of one engine (postgres, mysql, mongodb, cassandra, redis, elasticsearch). Per event: alarm, families, severity, evidence, whether the client is a known application server, the remediation to run on the database; plus the actions taken.",
     "input_schema": {"type": "object", "properties": {"events": {"type": "array", "items": {"type": "object"}}, "engine": {"type": "string"},
                                                       "lines": {"type": "array", "items": {"type": "string"}}, "max_rows": {"type": "integer", "default": 200}}}},
    {"name": "guard_db_calibrate", "description": "Give the database tier the applications' normal workload (events or raw log lines of one engine, no attack labels): which clients talk to which database, the query shapes they run, their volumes.",
     "input_schema": {"type": "object", "properties": {"events": {"type": "array", "items": {"type": "object"}}, "engine": {"type": "string"},
                                                       "lines": {"type": "array", "items": {"type": "string"}}}}},
    {"name": "guard_db_status", "description": "Database tier state: known databases with their application clients and query shapes, the learned tier's reference and threshold per database, counts, and the recent database alarms with their remediation.",
     "input_schema": {"type": "object", "properties": {"recent": {"type": "integer", "default": 20}}}},
    {"name": "guard_rpc", "description": "Score JSON-RPC calls a chain node received (bodies as JSON text or objects, one client IP per body, optional unix times / error flags): RPC method abuse, expensive-method floods, mempool spam. Returns per call alarm, family, score, playbook ids, and the actions taken.",
     "input_schema": {"type": "object", "properties": {"bodies": {"type": "array", "items": {}}, "client": {"type": "array", "items": {"type": "string"}},
                                                       "t": {"type": "array", "items": {"type": "number"}}, "node": {"type": "string", "default": "node"},
                                                       "errors": {"type": "array", "items": {"type": "boolean"}}, "max_rows": {"type": "integer", "default": 200}},
                      "required": ["bodies", "client"]}},
]
# tools that feed traffic to the engine: they change its state and the layer acts on what they carry - whoever can post
# to them can make the layer block an address of their choice. Over HTTP a feed that carries sources (client / t / flows /
# events) needs the feed token (DIBAGUARD_FEED_TOKEN, header X-DibaGuard-Feed-Token; the Go proxy sends it) or the console
# token. A score-only call (request texts without client or time) changes nothing and stays open. DIBAGUARD_FEEDS_OPEN=1
# reopens every feed (an isolated loopback deployment only); DIBAGUARD_FEEDS_NEED_TOKEN=1 demands the console token and
# the console header. Re-calibration is a control call (it rewrites what the site's normal is).
FEED_TOOLS = {"guard_score_flows", "guard_score_http", "guard_calibrate", "guard_db", "guard_db_calibrate", "guard_rpc"}
CALIBRATE_TOOLS = {"guard_calibrate", "guard_db_calibrate"}


def is_stateful_feed(name: str, args: dict) -> bool:
    """a feed whose arguments reach the entity stream or the enforcer (a score-only HTTP call does not)"""
    if name not in FEED_TOOLS:
        return False
    if name == "guard_score_http":
        return any(args.get(k) is not None for k in ("client", "t", "status"))
    return True


def is_control(name: str, args: dict, layer=None) -> bool:
    """tool calls that change what the layer enforces (verdicts, rules, policy, port closing, the calibrated normal): over
    HTTP they need the console token and the X-DibaGuard-Console header; MCP over stdio is the local process's own channel."""
    if name in CALIBRATE_TOOLS:
        return True
    if name == "guard_policy":
        return bool(args.get("set"))
    if name == "guard_feedback":
        return True
    if name == "guard_rules":
        return bool(args.get("add")) and not args.get("interpret_only")
    if name == "guard_posture":
        return layer is not None and bool(layer.policy.d["posture"].get("close_unexpected")) and layer.policy.enforcing
    return False


def _db_events(args: dict) -> list[dict]:
    ev = [dict(e) for e in (args.get("events") or [])]
    if args.get("lines"):
        from dibaguard.shield.database import PARSERS

        parse = PARSERS.get(str(args.get("engine", "")).lower())
        if parse is None:
            raise ValueError(f"to parse log lines give engine, one of {sorted(PARSERS)}")
        ev += [e for e in (parse(x) for x in args["lines"]) if e]
    for e in ev:
        if not all(k in e for k in ("t", "engine", "client")):
            raise ValueError("each database event needs t, engine and client")
        e["t"] = float(e["t"])
    return ev


def as_anthropic_tools() -> list[dict]:
    return [dict(t) for t in TOOLS]


def as_openai_tools() -> list[dict]:
    return [{"type": "function", "function": {"name": t["name"], "description": t["description"], "parameters": t["input_schema"]}} for t in TOOLS]


def _jsonable(o):
    if isinstance(o, dict):
        return {str(k): _jsonable(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)):
        return [_jsonable(v) for v in o]
    if isinstance(o, (set, frozenset)):
        return [_jsonable(v) for v in sorted(o, key=str)]
    if isinstance(o, np.ndarray):
        return _jsonable(o.tolist())
    if isinstance(o, (np.floating, np.integer, np.bool_)):
        return o.item()
    return o


def _summary(out: dict, keys=("p_attack", "alarm", "family", "score", "tier", "entity_risk")) -> dict:
    al = np.asarray(out["alarm"])
    fam = np.asarray(out["family"], dtype=object)
    return {"n": int(len(al)), "alarms": int(al.sum()), "max_score": int(np.max(out["score"])) if len(al) else 0,
            "families": {f: int((fam[al] == f).sum()) for f in set(fam[al].tolist())}, "site_pressure": float(out["site_pressure"])}


class Dispatcher:
    def __init__(self, guard, layer=None):
        """guard: a Guard; layer: a shield.Layer (then verdicts are also enforced and the shield tools work)"""
        self.layer = layer
        self.guard = layer.guard if layer is not None else guard

    def call(self, name: str, args: dict) -> dict:
        g, L = self.guard, self.layer
        if name == "guard_status":
            st = g.status()
            if L is not None:
                st["enforcement"] = L.enforcer.state()
                st["position"] = {"roles": L.position["roles"], "firewall": L.position["firewall"]}
            return st
        if name == "guard_fast_state":
            return L.fast_state() if L is not None else {"block": {}, "challenge": {}, "rate_limit": {}}
        if name == "guard_posture":
            return L.posture() if L is not None else {"error": "no shield layer"}
        if name == "guard_playbook":
            from dibaguard.shield import playbook

            return {"patterns": playbook.PLAYBOOK, "matched": L.playbook_hits if L is not None else {}}
        if name == "guard_feedback_pending":
            return {"pending": L.feedback.pending(int(args.get("limit", 50)), args.get("since"))} if L is not None else {"pending": []}
        if name == "guard_feedback":
            return L.feedback.decide(args["entity"], args["verdict"], args.get("note", ""), args.get("who", "operator"),
                                     force=bool(args.get("force"))) if L is not None else {"error": "no shield layer"}
        if name == "guard_feedback_stats":
            return L.feedback.stats() if L is not None else {}
        if name == "guard_rules":
            if L is None:
                return {"error": "no shield layer"}
            if args.get("add"):
                return L.add_rules(args["add"], interpret_only=bool(args.get("interpret_only")))
            return {"rules": [r.as_dict() for r in L.rules.rules]}
        if name == "guard_policy":
            if L is None:
                return {"error": "no shield layer"}
            if args.get("set"):
                from dibaguard.shield.policy import Policy

                d = json.loads(L.policy.dump())
                for k, v in args["set"].items():
                    if isinstance(v, dict) and isinstance(d.get(k), dict):
                        d[k].update(v)
                    else:
                        d[k] = v
                L.policy = Policy(d)
                L.enforcer.policy = L.policy
            return json.loads(L.policy.dump())
        if name == "guard_score_flows":
            flows = {k: np.asarray(v) for k, v in args["flows"].items()}
            out = L.flows(flows) if L is not None else g.flows(flows)
            m = int(args.get("max_rows", 200))
            return {"summary": _summary(out), "threshold": {"value": float(out["threshold_tier1"]), "source": out["threshold_source"]},
                    "drift": bool(out["drift"]), "rows": [{k: out[k][i] for k in ("p_attack", "alarm", "family", "score", "tier", "entity_risk")}
                                                           for i in range(min(m, len(out["alarm"])))]}
        if name == "guard_score_http":
            out = (L.http if L is not None else g.http_requests)(args["texts"], args.get("deepest"), args.get("client"), args.get("host"), args.get("status"), args.get("t"))
            return {"summary": _summary(out), "threshold": {"value": float(out["threshold"]), "source": out["threshold_source"]},
                    "rows": [{k: (out[k][i] if out[k] is not None else None) for k in ("p_attack", "alarm", "family", "score", "entity_risk")}
                             for i in range(len(out["alarm"]))]}
        if name == "guard_calibrate":
            res = {}
            if args.get("flows"):
                res.update(g.calibrate_flows({k: np.asarray(v) for k, v in args["flows"].items()}))
            if args.get("texts"):
                res.update(g.calibrate_http(args["texts"]))
            res["thresholds"] = g.status()["thresholds"]
            return res
        if name == "guard_report":
            if L is not None:
                return L.report(args.get("since"), args.get("until"), int(args.get("top", 10)))
            return rep_mod.build(g, args.get("since"), args.get("until"), int(args.get("top", 10)))
        if name in ("guard_db", "guard_db_calibrate", "guard_db_status", "guard_rpc") and L is None:
            return {"error": "no shield layer"}
        if name == "guard_rpc":
            bodies = [b if isinstance(b, (str, bytes)) else json.dumps(b) for b in args["bodies"]]
            client = [str(c) for c in args["client"]]
            if len(client) != len(bodies) or (args.get("t") is not None and len(args["t"]) != len(bodies)):
                raise ValueError("bodies, client (and t) must have the same length")
            if not bodies:
                return {"summary": {"n": 0, "alarms": 0, "max_score": 0, "families": {}}, "rows": [], "actions": []}
            out = L.rpc(bodies, client, args.get("t"), str(args.get("node", "node")), args.get("errors"))
            al = np.asarray(out["alarm"], bool)
            fam = np.asarray(out["family"], dtype=object)
            return {"summary": {"n": int(len(al)), "alarms": int(al.sum()), "max_score": int(np.max(out["score"])) if len(al) else 0,
                                "families": {f: int((fam[al] == f).sum()) for f in set(fam[al].tolist())}},
                    "rows": [{"alarm": bool(al[i]), "family": fam[i], "score": int(out["score"][i]), "playbook": out["playbook"][i]}
                             for i in range(min(int(args.get("max_rows", 200)), len(al)))],
                    "actions": out["actions"]}
        if name == "guard_db":
            out = L.db(_db_events(args))
            al = [v for v in out["verdicts"] if v["alarm"]]
            fams: dict = {}
            for v in al:
                f = v["families"][0] if v["families"] else "new_client"
                fams[f] = fams.get(f, 0) + 1
            return {"summary": {"n": len(out["verdicts"]), "alarms": len(al), "max_score": max((v["score"] for v in al), default=0), "families": fams},
                    "verdicts": out["verdicts"][: int(args.get("max_rows", 200))], "actions": out["actions"]}
        if name == "guard_db_calibrate":
            return L.calibrate_db(_db_events(args))
        if name == "guard_db_status":
            d = L.dbg
            return {"calibrated": bool(d.known_clients), "dbnet": d.net is not None, "counts": dict(d.counts), "exempt_shapes": len(d.exempt_shapes),
                    "databases": {db: {"clients": len(c), "client_sample": sorted(c)[:20], "shapes": len(d.known_shapes.get(db, ())),
                                       "dbnet_reference": len(d.net_ref.get(db, ())), "dbnet_threshold": d.net_threshold(db) if d.net is not None else None}
                                  for db, c in d.known_clients.items()},
                    "policy": L.policy.d["database"],
                    "recent_alarms": [{k: v.get(k) for k in ("t", "engine", "db", "client", "user", "families", "score", "app_client", "public", "evidence", "remediation")}
                                      for v in d.log[len(d.log) - max(0, int(args.get("recent", 20))):][::-1]]}
        raise KeyError(f"unknown tool {name}")

    def call_json(self, name: str, args: dict) -> dict:
        return _jsonable(self.call(name, args))


# ---------------------------------------------------------------- MCP over stdio (JSON-RPC 2.0, protocol 2024-11-05)
def serve_mcp(disp: Dispatcher, inp=None, out=None) -> None:
    inp, out = inp or sys.stdin, out or sys.stdout

    def send(obj):
        out.write(json.dumps(obj, ensure_ascii=False) + "\n")
        out.flush()

    for line in inp:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except json.JSONDecodeError:
            continue
        rid, method, params = req.get("id"), req.get("method"), req.get("params") or {}
        if method == "initialize":
            send({"jsonrpc": "2.0", "id": rid, "result": {"protocolVersion": "2024-11-05", "capabilities": {"tools": {}},
                                                           "serverInfo": {"name": "dibaguard", "version": "2"}}})
        elif method == "tools/list":
            send({"jsonrpc": "2.0", "id": rid, "result": {"tools": [{"name": t["name"], "description": t["description"], "inputSchema": t["input_schema"]} for t in TOOLS]}})
        elif method == "tools/call":
            try:
                res = disp.call_json(params.get("name", ""), params.get("arguments") or {})
                send({"jsonrpc": "2.0", "id": rid, "result": {"content": [{"type": "text", "text": json.dumps(res, ensure_ascii=False)}], "isError": False}})
            except Exception as e:  # noqa: BLE001
                send({"jsonrpc": "2.0", "id": rid, "result": {"content": [{"type": "text", "text": f"{type(e).__name__}: {e}"}], "isError": True}})
        elif method == "ping":
            send({"jsonrpc": "2.0", "id": rid, "result": {}})
        elif rid is not None:   # notifications (no id) are ignored
            send({"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": f"method not found: {method}"}})


# ---------------------------------------------------------------- plain HTTP (stdlib)
def make_server(disp: Dispatcher, host: str = "127.0.0.1", port: int = 8787, token: str | None = None, provider="env", quiet: bool = False):
    """-> (server, console). The server is threaded (an LLM turn in the console never stalls the proxy's polling); every
    call into the engine / layer holds console.lock. token: the console secret (default: DIBAGUARD_CONSOLE_TOKEN, else
    generated and printed once to stderr). provider: "env" (DIBAGUARD_LLM_*), None (no LLM) or an object with
    complete(system, messages, tools) - tests pass a fake one."""
    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
    from urllib.parse import parse_qs

    from dibaguard import console as con

    cons = con.Console(disp, token=token, provider=provider)

    class H(BaseHTTPRequestHandler):
        server_version = "dibaguard"
        sys_version = ""

        def _send(self, code, obj, extra=None):
            body = json.dumps(con.jsonable(obj), ensure_ascii=False).encode()
            self.send_response(code)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.send_header("X-Content-Type-Options", "nosniff")
            for k, v in (extra or {}).items():
                self.send_header(k, v)
            self.end_headers()
            self.wfile.write(body)

        def _redirect(self, where):
            self.send_response(302)
            self.send_header("Location", where)
            self.send_header("Content-Length", "0")
            self.end_headers()

        def do_GET(self):  # noqa: N802
            path, _, qs = self.path.partition("?")
            if path in ("/", "/dashboard"):          # the old review page is part of the console now
                return self._redirect("/console")
            if path == "/console" or path.startswith("/console/"):
                if not cons.host_ok(self.headers.get("Host")):
                    return self._send(403, {"error": "unexpected Host header (set DIBAGUARD_CONSOLE_HOSTS to allow a name)"})
                if path == "/console":
                    body = con.PAGE_BYTES
                    self.send_response(200)
                    self.send_header("Content-Type", "text/html; charset=utf-8")
                    self.send_header("Content-Length", str(len(body)))
                    for k, v in con.PAGE_HEADERS.items():
                        self.send_header(k, v)
                    self.end_headers()
                    self.wfile.write(body)
                    return
                code, obj = cons.handle_get(path, {k: v[-1] for k, v in parse_qs(qs).items()})
                return self._send(code, obj)
            if path == "/tools":
                self._send(200, {"anthropic": as_anthropic_tools(), "openai": as_openai_tools(), "report_schema": rep_mod.REPORT_SCHEMA})
            elif path == "/health":
                self._send(200, {"ok": True})
            else:
                self._send(404, {"error": "not found"})

        def do_POST(self):  # noqa: N802
            path = self.path.partition("?")[0]
            n = int(self.headers.get("Content-Length") or 0)
            if path.startswith("/console/"):
                bad = cons.check_write(self.headers, n)
                if bad:
                    self.rfile.read(min(n, 1 << 20))
                    return self._send(bad[0], {"error": bad[1]})
                try:
                    body = json.loads(self.rfile.read(n) or b"{}")
                    if not isinstance(body, dict):
                        raise ValueError("body must be a JSON object")
                except ValueError as e:
                    return self._send(400, {"error": f"invalid JSON: {e}"})
                code, obj = cons.handle_post(path, body)
                return self._send(code, obj)
            if cons.cross_site(self.headers):            # a browser page of another site posting here: never
                self.rfile.read(n)
                return self._send(403, {"error": "cross-site request refused"})
            try:
                args = json.loads(self.rfile.read(n) or b"{}")
                if not isinstance(args, dict):
                    raise ValueError("arguments must be a JSON object")
                name = path.rsplit("/", 1)[-1]
                if is_control(name, args, disp.layer) or (cons.feeds_need_token and name in FEED_TOOLS):
                    if self.headers.get("X-DibaGuard-Console") != "1":
                        return self._send(403, {"error": f"{name} changes state: send the X-DibaGuard-Console: 1 header (CSRF guard) and the console token"})
                    if not cons.token_ok(self.headers):
                        return self._send(401, {"error": f"{name} changes state: send the console token (X-DibaGuard-Token or Authorization: Bearer)"})
                elif is_stateful_feed(name, args) and not cons.feeds_open and not (cons.feed_token_ok(self.headers) or cons.token_ok(self.headers)):
                    return self._send(401, {"error": f"{name} feeds traffic the layer acts on: send the feed token (X-DibaGuard-Feed-Token, "
                                                     f"DIBAGUARD_FEED_TOKEN) or the console token"})
                with cons.lock:
                    res = disp.call_json(name, args)
                self._send(200, res)
            except KeyError as e:
                self._send(404, {"error": str(e)})
            except Exception as e:  # noqa: BLE001
                self._send(400, {"error": f"{type(e).__name__}: {e}"})

        def log_message(self, *a):
            pass

    srv = ThreadingHTTPServer((host, port), H)
    srv.daemon_threads = True
    if not quiet:
        cons.announce(host, srv.server_address[1])
    return srv, cons


def serve_http(disp: Dispatcher, host: str = "127.0.0.1", port: int = 8787, token: str | None = None, provider="env") -> None:
    make_server(disp, host, port, token, provider)[0].serve_forever()


def main(argv=None):
    import argparse

    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["mcp", "http", "tools"])
    ap.add_argument("--models", default="models")
    ap.add_argument("--alpha", type=float, default=1e-3)
    ap.add_argument("--policy", default=None, help="policy.json: with it the shield layer enforces (or reports) decisions")
    ap.add_argument("--inputs", default="http,flows", help="what this instance sees: http,flows,rpc,lb")
    ap.add_argument("--audit", default=None, help="append enforcement audit records here (jsonl)")
    ap.add_argument("--backend", default="auto", help="firewall backend: auto | nftables | iptables | windows | dryrun")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8787)
    a = ap.parse_args(argv)
    if a.mode == "tools":
        print(json.dumps({"anthropic": as_anthropic_tools(), "openai": as_openai_tools()}, ensure_ascii=False, indent=1))
        return
    if a.policy:
        from dibaguard.shield.layer import Layer

        disp = Dispatcher(None, Layer.load(a.models, a.policy, inputs={k: True for k in a.inputs.split(",") if k}, backend=a.backend, audit_path=a.audit))
    else:
        disp = Dispatcher(Guard.load(a.models, alpha=a.alpha))
    if a.mode == "mcp":
        serve_mcp(disp)
        return
    srv, cons = make_server(disp, a.host, a.port)
    hk = None
    if disp.layer is not None:                    # high availability (snapshots / replica sync) when DIBAGUARD_* is set
        from dibaguard import state
        hk = state.from_env(disp.layer, cons.lock)
    try:
        srv.serve_forever()
    finally:
        if hk is not None:
            hk.stop()                             # a final snapshot


if __name__ == "__main__":
    main()