nft improve

backend/src/api/nft_manager.py

@@ -4,6 +4,7 @@ from fastapi import FastAPI, APIRouter, HTTPException, status
 from pydantic import BaseModel, Field
 import logging
 import json
+import re
 
 # libnftables (we call textual commands through its .cmd() method)
 from nftables import Nftables  # type: ignore
@@ -54,6 +55,45 @@ class NftManager:
             raise NftError(f"nft list ruleset failed: {res['stderr']}")
         return res["stdout"]
 
+    def list_rules_json(self) -> Dict[str, Any]:
+        """
+        Try to obtain nft -j list ruleset (JSON). Returns parsed JSON dict on success.
+        Raises NftError on failure or when output cannot be parsed as JSON.
+        """
+        # Try passing -j explicitly; depending on python-nftables wrapper, `set_json_output(True)` might already do it.
+        # If -j is not accepted by the wrapper, we still attempt a plain "list ruleset" which can be JSON depending on set_json_output.
+        cmd_variants = ["list ruleset -j", "list ruleset"]
+        last_err = None
+        for c in cmd_variants:
+            res = self.cmd(c)
+            if res["rc"] != 0:
+                last_err = res["stderr"]
+                continue
+            out = res["stdout"]
+            if not out:
+                # empty output is treated as error here
+                last_err = "empty output"
+                continue
+            try:
+                parsed = json.loads(out)
+                return parsed
+            except json.JSONDecodeError as e:
+                # Not JSON for this invocation; continue to next variant
+                last_err = f"json decode error: {e}"
+                continue
+        raise NftError(f"unable to get JSON ruleset: {last_err}")
+
+    def list_chain_text(self, family: str, table: str, chain: str) -> str:
+        """
+        Return textual output of `nft list chain <family> <table> <chain>`.
+        This output includes the chain header and rule lines. Caller should parse it.
+        """
+        cmd = f"list chain {family} {table} {chain}"
+        res = self.cmd(cmd)
+        if res["rc"] != 0:
+            raise NftError(f"nft {cmd} failed: {res['stderr']}")
+        return res["stdout"]
+
     def delete_rule_by_handle_text(self, family: str, table: str, chain: str, handle: int) -> None:
         """
         Delete a rule by handle using textual nft command:
@@ -89,38 +129,123 @@ class ExecResult(BaseModel):
 
 
 # ruleset may be a parsed JSON object (dict/list) OR a raw string (text) OR null
-RulesetValue = Dict
+RulesetValue = Union[Dict[str, Any], List[Any], str, None]
 
 
 class RulesetOut(BaseModel):
     ruleset: RulesetValue = Field(
-        description="Parsed JSON ruleset (dict/list).",
+        description="parsed JSON ruleset (dict/list), or the raw textual ruleset as a string.",
     )
 
 
+# ---------- Helpers to convert to desired shape ----------
+_handle_re = re.compile(r"\s+#\s*handle\s+\d+\s*$")
+
+def extract_rule_lines_from_chain_text(text: str) -> List[str]:
+    """
+    Given output of `nft list chain fam table chain`, extract the rule lines as strings.
+    Keeps the rule text as printed by nft (one rule per line), strips leading/trailing whitespace,
+    and removes trailing '# handle N' fragments.
+    """
+    lines: List[str] = []
+    if not text:
+        return lines
+    for raw in text.splitlines():
+        line = raw.rstrip()
+        # skip chain/table header lines which typically start with "table " or "chain "
+        if line.strip() == "":
+            continue
+        if line.lstrip().startswith("table "):
+            continue
+        if line.lstrip().startswith("chain "):
+            continue
+        # rule lines are indented (start with whitespace). Accept them if non-empty after stripping.
+        # Remove leading indentation:
+        stripped = line.lstrip()
+        # remove trailing " # handle N" if present
+        stripped = _handle_re.sub("", stripped)
+        if stripped:
+            lines.append(stripped)
+    return lines
+
+
+def build_custom_ruleset_from_nft_json(nft_json: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Build the desired structure:
+    { "table": [ { "name": <table>, "family": <family>, "chains": [ { "name": <chain>, "rules": [<rule strings>] } ] } ] }
+    Uses nft_json only to discover families/tables/chains, then fetches textual chain listing for exact rule strings.
+    """
+    result = {"table": []}
+    # nft_json is expected to be the parsed output of `nft -j list ruleset` which contains "nftables": [ ... ]
+    items = nft_json.get("nftables", [])
+    # discover tables and associated family/name
+    tables: Dict[tuple, Dict[str, Any]] = {}
+    # items can contain separate objects for table/chain/rule entries
+    for item in items:
+        if "table" in item:
+            t = item["table"]
+            fam = t.get("family")
+            name = t.get("name")
+            if fam and name:
+                key = (fam, name)
+                if key not in tables:
+                    tables[key] = {"name": name, "family": fam, "chains": {}}
+        elif "chain" in item:
+            ch = item["chain"]
+            fam = ch.get("family") or ch.get("table", {}).get("family")  # defensive
+            table_name = ch.get("table") or ch.get("table", {}).get("name")  # defensive
+            chain_name = ch.get("name")
+            if fam and table_name and chain_name:
+                key = (fam, table_name)
+                if key not in tables:
+                    tables[key] = {"name": table_name, "family": fam, "chains": {}}
+                # register chain placeholder
+                tables[key]["chains"].setdefault(chain_name, {"name": chain_name, "rules": []})
+
+    # Now for each discovered table+chain call textual `nft list chain ...` to get actual rule lines
+    for (fam, tname), tdata in tables.items():
+        chains_out: List[Dict[str, Any]] = []
+        for cname in sorted(tdata["chains"].keys()):
+            try:
+                chain_text = mgr.list_chain_text(fam, tname, cname)
+                rules_lines = extract_rule_lines_from_chain_text(chain_text)
+            except NftError as e:
+                logger.warning("failed to list chain text for %s %s %s: %s", fam, tname, cname, e)
+                # fallback to empty rules list on error for that chain
+                rules_lines = []
+            chains_out.append({"name": cname, "rules": rules_lines})
+        result["table"].append({"name": tname, "family": fam, "chains": chains_out})
+    return result
+
+
 # ---------- Routes ----------
 
 @router.get("/rules", response_model=RulesetOut, summary="List ruleset")
 def list_rules():
     """
-    Returns the textual nft ruleset output (as a string) or native JSON if nft returned JSON.
-    Clients should handle both cases.
+    Returns the ruleset in the custom JSON shape:
+      { "table": [ { "name": ..., "family": ..., "chains": [ { "name": ..., "rules": [ "<rule text>", ... ] } ] } ] }
+
+    Implementation:
+      1. Try to get JSON ruleset via nft -j list ruleset
+      2. Use JSON to discover tables & chains
+      3. For each chain fetch textual `nft list chain fam table chain` and extract rule lines
+      4. Return the composed structure
+    If JSON isn't available or an error occurs, fall back to returning the raw textual ruleset string (existing behavior).
     """
     try:
-        text = mgr.list_rules()
-        if text is None:
-            return {"ruleset": None}
-        s = text.strip()
-        # Try to interpret as JSON. nft may produce JSON when set_json_output(True).
+        # Try to obtain JSON ruleset
         try:
-            parsed = json.loads(s)
-            # Return parsed JSON (FastAPI will serialize)
-            return {"ruleset": parsed}
-        except json.JSONDecodeError:
-            # Not JSON — return raw text
-            # Important: we do NOT treat non-json as an internal error; return raw string.
-            logger.debug("nft output is not JSON; returning raw text")
-            return {"ruleset": s}
+            nft_json = mgr.list_rules_json()
+        except NftError as e:
+            logger.debug("could not obtain nft JSON ruleset: %s", e)
+            # fallback to returning raw textual ruleset (existing behavior)
+            text = mgr.list_rules()
+            return {"ruleset": text.strip() if text is not None else None}
+
+        # Build custom structure using the JSON to find tables/chains, and textual listing to obtain rule lines
+        custom = build_custom_ruleset_from_nft_json(nft_json)
+        return {"ruleset": custom}
     except NftError as e:
         logger.exception("list_rules failed")
         raise HTTPException(status_code=500, detail=str(e))