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