test3

backend/src/api/nftables_api.py

@@ -1,13 +1,17 @@
-# fastapi_nft_stateless_comment_enums.py
+# fastapi_nft_stateless_comment_enums_resilient_full.py
 """
-Stateless FastAPI nftables router using pyroute2.nftables.
+Stateless FastAPI nftables router with enums and resilient pyroute2 binding.
 
 - Stateless: no in-process or on-disk rule store.
 - Rules may include an optional 'comment' field that will be written into nft's comment.
-- Enums introduced for Action.type, ip_proto (common names) and family.
-- Endpoints:
-  - GET  /nft/rules   -> reconstruct rules from kernel (returns comment if present)
-  - PUT  /nft/rules   -> replace entire ordered rule set (clients supply optional comment per rule)
+- Uses pyroute2 binding if available and can be instantiated synchronously, otherwise falls back to calling the `nft` CLI via subprocess.
+- Exposes two endpoints under /nft:
+    GET  /nft/rules  -> list rules reconstructed from kernel state (returns comment if present)
+    PUT  /nft/rules  -> replace entire ordered ruleset (clients supply optional comment per rule)
+
+Notes:
+- The service must run with privileges to modify nftables (root / CAP_NET_ADMIN) when applying rules.
+- If the subprocess fallback is used, ensure the `nft` binary is present.
 """
 from fastapi import APIRouter, HTTPException, Header, Request
 from pydantic import BaseModel, Field
@@ -16,13 +20,24 @@ import logging
 import json
 import uuid
 from enum import Enum
+import asyncio
+import subprocess
 
-# pyroute2 nftables
-from pyroute2.nftables.main import NFTables
+# Try to import NFTables binding (various pyroute2 layouts)
+try:
+    # common location
+    from pyroute2.nftables.main import NFTables as NFTablesBinding  # type: ignore
+except Exception:
+    try:
+        # alternate location
+        from pyroute2.nftables import NFTables as NFTablesBinding  # type: ignore
+    except Exception:
+        NFTablesBinding = None  # will fall back to subprocess wrapper
 
+# Router + logger
 router = APIRouter()
 logger = logging.getLogger("nftables")
-logger.debug("nftables stateless-comment-enums router loaded")
+logger.debug("nftables router module loaded")
 
 # Defaults and version token
 DEFAULT_TABLE = "mitm_tbl"
@@ -58,14 +73,13 @@ class MatchModel(BaseModel):
     iif: Optional[str] = None
     oif: Optional[str] = None
     meta_length: Optional[Any] = None  # int or range string
-    # ip_proto may be Protocol enum, integer or arbitrary string (name)
     ip_proto: Optional[Union[int, Protocol, str]] = None
     tcp_dport: Optional[int] = None
     udp_dport: Optional[int] = None
 
 
 class ActionModel(BaseModel):
-    type: ActionType  # ActionType enum
+    type: ActionType
     queue_num: Optional[int] = None
     redirect_port: Optional[int] = None
 
@@ -86,38 +100,33 @@ class ReplaceResult(BaseModel):
     rules_count: int
 
 
-# ---------------------- nft wrapper ----------------------
-class NFT:
+# ---------------------- resilient NFT wrapper selection ----------------------
+class NFTSubprocessWrapper:
+    """
+    Fallback wrapper calling the `nft` CLI via subprocess.
+    """
+
     def __init__(self):
-        self.nft = NFTables()
+        # log availability
+        try:
+            subprocess.run(["nft", "--version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+            logger.debug("nft CLI available for subprocess wrapper")
+        except Exception as e:
+            logger.warning("nft CLI not available or couldn't be invoked: %s", e)
 
     def run(self, cmd: str) -> Dict[str, Any]:
-        """
-        Run a single nft command string via pyroute2.NFTables.cmd.
-        Returns parsed JSON if possible, otherwise a dict with 'out' textual output.
-        Raises RuntimeError on failure.
-        """
+        # naive split is acceptable for the limited command patterns used here
+        args = ["nft"] + cmd.split()
         try:
-            rc, out, err = self.nft.cmd(cmd)
-            if isinstance(out, bytes):
-                out = out.decode(errors="ignore")
-            if isinstance(err, bytes):
-                err = err.decode(errors="ignore")
-            if rc != 0:
-                raise RuntimeError(err or f"nft cmd failed rc={rc}")
-            if out:
+            proc = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
+            out = proc.stdout or ""
             try:
                 return json.loads(out)
             except Exception:
                 return {"out": out}
-            return {}
-        except AttributeError:
-            # Try json_cmd fallback if older pyroute2 version
-            out = self.nft.json_cmd(cmd)
-            return out or {}
-        except Exception as e:
-            logger.exception("nft wrapper error for cmd=%s: %s", cmd, e)
-            raise
+        except subprocess.CalledProcessError as e:
+            logger.error("nft CLI failed cmd=%s stderr=%s", cmd, e.stderr)
+            raise RuntimeError(e.stderr or str(e))
 
     def add_table(self, family: str, table: str):
         return self.run(f"add table {family} {table}")
@@ -138,7 +147,107 @@ class NFT:
         return self.run(f"delete rule {family} {table} {chain} handle {handle}")
 
 
-NFTC = NFT()
+class NFTBindingWrapper:
+    """
+    Wrapper using pyroute2 NFTables binding, instantiated synchronously if possible.
+    If the binding requires async setup and the current process already has a running event loop,
+    this wrapper raises so the factory can fall back to subprocess wrapper.
+    """
+
+    def __init__(self, binding_cls):
+        self._binding = None
+        self._constructed = False
+
+        tried_kwargs = [
+            {"async": False},
+            {"nl_async": False},
+            {"use_async": False},
+            {"asyncio": False},
+        ]
+        last_exc = None
+        for kw in tried_kwargs:
+            try:
+                self._binding = binding_cls(**kw)
+                self._constructed = True
+                logger.debug("NFTables binding instantiated with kwargs %s", kw)
+                break
+            except TypeError as e:
+                last_exc = e
+            except Exception as e:
+                last_exc = e
+
+        if not self._constructed:
+            # try no-arg
+            try:
+                self._binding = binding_cls()
+                self._constructed = True
+                logger.debug("NFTables binding instantiated with no kwargs")
+            except Exception as e:
+                last_exc = e
+
+        if not self._constructed:
+            raise RuntimeError(f"failed to instantiate NFTables binding: {last_exc}")
+
+        # If the binding exposes an async setup coroutine, ensure we can run it synchronously.
+        setup_coro = getattr(self._binding, "setup_endpoint", None)
+        if setup_coro and asyncio.iscoroutinefunction(setup_coro):
+            # If event loop running, we cannot await here — signal unsuitability.
+            if asyncio.get_event_loop().is_running():
+                raise RuntimeError("pyroute2 NFTables requires async setup but event loop is already running")
+            # otherwise run it to complete setup
+            try:
+                asyncio.get_event_loop().run_until_complete(setup_coro())
+            except Exception as e:
+                logger.error("awaiting binding.setup_endpoint() failed: %s", e)
+                raise
+
+    def run(self, cmd: str) -> Dict[str, Any]:
+        try:
+            rc, out, err = self._binding.cmd(cmd)
+            if isinstance(out, bytes):
+                out = out.decode(errors="ignore")
+            if isinstance(err, bytes):
+                err = err.decode(errors="ignore")
+            if rc != 0:
+                raise RuntimeError(err or f"nft cmd failed rc={rc}")
+            if out:
+                try:
+                    return json.loads(out)
+                except Exception:
+                    return {"out": out}
+            return {}
+        except AttributeError:
+            out = self._binding.json_cmd(cmd)
+            return out or {}
+        except Exception as e:
+            logger.exception("nft binding cmd failed: %s", e)
+            raise
+
+
+def make_nft_wrapper():
+    """
+    Choose the best NFT wrapper: try binding if available and can be used synchronously;
+    otherwise fall back to subprocess wrapper.
+    """
+    if NFTablesBinding is not None:
+        try:
+            # If an event loop is running (uvicorn), prefer subprocess to avoid async binding setup.
+            if asyncio.get_event_loop().is_running():
+                logger.info("asyncio loop is running; skipping binding and using subprocess wrapper")
+                raise RuntimeError("event loop running")
+            w = NFTBindingWrapper(NFTablesBinding)
+            logger.info("using pyroute2 NFTables binding")
+            return w
+        except Exception as e:
+            logger.warning("pyroute2 binding unavailable/synchronous construction failed: %s; falling back to nft CLI", e)
+    # Fallback
+    logger.info("using nft CLI subprocess wrapper")
+    return NFTSubprocessWrapper()
+
+
+# instantiate wrapper at import time
+NFTC = make_nft_wrapper()
+logger.info("selected nft wrapper: %s", type(NFTC).__name__)
 
 
 # ---------------------- builders / parsers ----------------------
@@ -151,7 +260,6 @@ def build_match_frag(match: MatchModel) -> List[str]:
     if match.meta_length is not None:
         frag += ["meta", "length", str(match.meta_length)]
     if match.ip_proto is not None:
-        # ip_proto can be Protocol enum, int or str
         if isinstance(match.ip_proto, Protocol):
             frag += ["ip", "protocol", match.ip_proto.value]
         elif isinstance(match.ip_proto, int):
@@ -166,7 +274,6 @@ def build_match_frag(match: MatchModel) -> List[str]:
 
 
 def build_action_frag(action: ActionModel) -> List[str]:
-    # ActionModel.type is an ActionType enum
     if action.type == ActionType.DROP:
         return ["drop"]
     if action.type == ActionType.ACCEPT:
@@ -182,21 +289,15 @@ def build_action_frag(action: ActionModel) -> List[str]:
 
 
 def nft_rule_fragment_from_model(rule: RuleModel) -> str:
-    """
-    Build fragment after 'add rule <family> <table> <chain>'.
-    Includes comment if rule.comment provided.
-    """
     match_frag = build_match_frag(rule.match)
     action_frag = build_action_frag(rule.action)
     parts = match_frag + action_frag
     if rule.comment:
-        # include comment as-is (user-provided). Wrap in quotes.
         parts += ['comment', f'"{rule.comment}"']
     return " ".join(parts)
 
 
 def parse_comment_from_exprs(exprs: List[Dict[str, Any]]) -> Optional[str]:
-    """Scan expressions for a comment expression and return its string if present."""
     for e in exprs:
         if "comment" in e:
             cm = e["comment"]
@@ -208,11 +309,6 @@ def parse_comment_from_exprs(exprs: List[Dict[str, Any]]) -> Optional[str]:
 
 
 def reconstruct_rule_from_rule_entry(entry: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Given a rule entry returned by NFTC.list_table, return a dict that resembles RuleModel
-    (family/table/chain + match + action + comment if present). This is best-effort parsing.
-    """
-    # normalize expressions
     exprs = []
     if "rule" in entry:
         r = entry["rule"]
@@ -230,8 +326,6 @@ def reconstruct_rule_from_rule_entry(entry: Dict[str, Any]) -> Dict[str, Any]:
         exprs = [exprs]
 
     comment = parse_comment_from_exprs(exprs)
-
-    # best-effort reconstruction of match & action
     match: Dict[str, Any] = {}
     action: Dict[str, Any] = {}
 
@@ -305,16 +399,14 @@ def reconstruct_rule_from_rule_entry(entry: Dict[str, Any]) -> Dict[str, Any]:
     if "type" not in action:
         action["type"] = "accept"
 
-    # convert action.type to ActionType enum if possible
+    # try to coerce action.type to ActionType
     try:
-        action_type_val = action.get("type")
-        if isinstance(action_type_val, str):
-            action["type"] = ActionType(action_type_val)
+        if isinstance(action.get("type"), str):
+            action["type"] = ActionType(action["type"])
     except Exception:
-        # leave as-is if conversion fails
         pass
 
-    # convert family to Family enum if possible
+    # coerce family to Family enum if possible
     try:
         if isinstance(family, str):
             family = Family(family)
@@ -394,9 +486,8 @@ def add_rules_replace_all(rules: List[RuleModel], family: Union[str, Family], ta
     except Exception as e:
         logger.debug("flush chain may have returned error: %s", e)
 
-    # add rules
+    # add rules in order
     for r in rules:
-        # ensure family value converted to string
         fam_r = r.family.value if isinstance(r.family, Family) else r.family
         frag = nft_rule_fragment_from_model(r)
         try:
@@ -455,7 +546,7 @@ def put_rules(
         logger.error("failed to ensure table/chain: %s", e)
         raise HTTPException(status_code=500, detail=str(e))
 
-    # ensure rule ids for client convenience
+    # ensure rule ids for convenience
     for r in rules:
         if not r.id:
             r.id = str(uuid.uuid4())