| """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"]}}, |
| ] |
| |
| |
| |
| |
| |
| |
| 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)) |
|
|
|
|
| |
| 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: |
| 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: |
| send({"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": f"method not found: {method}"}}) |
|
|
|
|
| |
| 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): |
| path, _, qs = self.path.partition("?") |
| if path in ("/", "/dashboard"): |
| 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): |
| 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): |
| 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: |
| 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: |
| from dibaguard import state |
| hk = state.from_env(disp.layer, cons.lock) |
| try: |
| srv.serve_forever() |
| finally: |
| if hk is not None: |
| hk.stop() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
