fix: Address all production audit findings (critical through low severity)

Full end-to-end audit of every source file found 47 issues.
This commit fixes all critical, medium, AND low severity bugs.

Critical fixes:
- Fix sessions.delete() → sessions.remove() (AttributeError crash)
- Fix scope enforcement bypass when _auth_token is None on session
- Fix agent ID regex rejecting manager agent "0" and short IDs "1","12"
- Fix ACTIVE_CONNECTIONS counter leak in /sse on early errors
- Fix WazuhClient.close() swallowing indexer cleanup on aclose() failure
- Fix verify_bearer_token crash on None input
- Fix block_ip without agent_id causing Wazuh API 400 error

Medium fixes:
- Add truncation warnings for vulnerability search tools
- Fix risk_assessment defaulting to "medium" with zero risk factors
- Add IP validation to firewall_allow and host_allow
- Fix OAuth state param not URL-encoded in redirect
- Fix indexer _execute_agg_search missing timeout handling
- Fix indexer close() not nulling client reference

Low severity fixes (previously accepted, now fixed):
- Fix module-level AuthManager/SecurityManager crash on malformed env vars
  (TOKEN_LIFETIME_HOURS, RATE_LIMIT_REQUESTS, RATE_LIMIT_WINDOW)
- Remove dead code _dict_contains_text (replaced by ES query passthrough)
- Fix CPU metric always returning 0.0 (reuse psutil.Process instance)
- Fix circuit breaker HALF_OPEN allowing multiple concurrent trial
  requests (add _half_open_trial_in_progress flag)

Tests: 84 passing (removed 6 dead-code tests, added 13 regression tests)

Author: alokemajumder

src/wazuh_mcp_server/api/wazuh_client.py

@@ -20,26 +20,6 @@
 _TIME_RANGE_HOURS = {"1h": 1, "6h": 6, "12h": 12, "1d": 24, "24h": 24, "7d": 168, "30d": 720}
 
 
-def _dict_contains_text(d: Any, text: str) -> bool:
-    """Recursively check if any string value in a dict/list contains the given text.
-    Much faster than json.dumps() + string search for large nested dicts."""
-    if isinstance(d, dict):
-        for v in d.values():
-            if _dict_contains_text(v, text):
-                return True
-    elif isinstance(d, (list, tuple)):
-        for item in d:
-            if _dict_contains_text(item, text):
-                return True
-    elif isinstance(d, str):
-        if text in d.lower():
-            return True
-    elif isinstance(d, (int, float)):
-        if text in str(d):
-            return True
-    return False
-
-
 class WazuhClient:
     """Simplified Wazuh API client with rate limiting, circuit breaker, and retry logic."""
 
@@ -278,16 +258,21 @@ async def execute_active_response(self, data: Dict[str, Any]) -> Dict[str, Any]:
         # Ensure data dict doesn't contain deprecated 'custom' parameter
         if "custom" in data:
             data = {k: v for k, v in data.items() if k != "custom"}
-        # Wazuh 4.x API: agent_list must be passed as query param 'agents_list', not in body
-        # Omit agents_list entirely to target all agents (Wazuh rejects "all" as non-numeric)
+        # Wazuh 4.x API: agent_list must be passed as query param 'agents_list'
         agents_list = data.pop("agent_list", None)
         params = {}
         if agents_list:
-            # Filter out "all" — Wazuh 4.x requires numeric agent IDs only
-            numeric_agents = [str(a) for a in (agents_list if isinstance(agents_list, list) else [agents_list]) if str(a) != "all"]
-            if numeric_agents:
-                params["agents_list"] = ",".join(numeric_agents)
-            # If only "all" was specified, omit agents_list to target all agents
+            agent_items = agents_list if isinstance(agents_list, list) else [agents_list]
+            # Check if targeting all agents
+            if any(str(a).lower() == "all" for a in agent_items):
+                # Wazuh 4.x API requires a valid agents_list; use "all" as a special keyword
+                # that the API accepts for targeting all agents
+                params["agents_list"] = "all"
+            else:
+                # Filter to numeric agent IDs only
+                numeric_agents = [str(a) for a in agent_items if str(a).isdigit()]
+                if numeric_agents:
+                    params["agents_list"] = ",".join(numeric_agents)
         result = await self._request("PUT", "/active-response", json=data, params=params)
 
         # Check for partial/total failures in the response body
@@ -807,8 +792,10 @@ async def perform_risk_assessment(self, agent_id: str = None) -> Dict[str, Any]:
             risk_level = "critical"
         elif any(f["severity"] == "high" for f in risk_factors):
             risk_level = "high"
-        else:
+        elif risk_factors:
             risk_level = "medium"
+        else:
+            risk_level = "low"
         return {
             "data": {
                 "total_agents": len(items),
@@ -1216,6 +1203,7 @@ async def restore_file(self, agent_id: str, file_path: str) -> Dict[str, Any]:
 
     async def firewall_allow(self, agent_id: str, src_ip: str) -> Dict[str, Any]:
         """Remove firewall drop rule via active response."""
+        self._validate_ip(src_ip)
         src_ip = self._sanitize_ar_argument(src_ip, "src_ip")
         data = {
             "command": "!firewall-drop",
@@ -1226,6 +1214,7 @@ async def firewall_allow(self, agent_id: str, src_ip: str) -> Dict[str, Any]:
 
     async def host_allow(self, agent_id: str, src_ip: str) -> Dict[str, Any]:
         """Remove hosts.deny entry via active response."""
+        self._validate_ip(src_ip)
         src_ip = self._sanitize_ar_argument(src_ip, "src_ip")
         data = {
             "command": "!host-deny",
@@ -1236,10 +1225,17 @@ async def host_allow(self, agent_id: str, src_ip: str) -> Dict[str, Any]:
 
     async def close(self):
         """Close the HTTP client and indexer client, releasing all connections."""
-        if self.client:
-            await self.client.aclose()
+        try:
+            if self.client:
+                await self.client.aclose()
+        except Exception:
+            pass  # Best-effort close; connection may already be broken
+        finally:
             self.client = None
-        if self._indexer_client:
-            await self._indexer_client.close()
+        try:
+            if self._indexer_client:
+                await self._indexer_client.close()
+        except Exception:
+            pass
         self.token = None
         self._cache.clear()

src/wazuh_mcp_server/api/wazuh_indexer.py

@@ -108,8 +108,13 @@ async def initialize(self):
 
     async def close(self):
         """Close the HTTP client."""
-        if self.client:
-            await self.client.aclose()
+        try:
+            if self.client:
+                await self.client.aclose()
+        except Exception:
+            pass  # Best-effort close
+        finally:
+            self.client = None
             self._initialized = False
 
     async def _ensure_initialized(self):
@@ -425,7 +430,7 @@ async def _execute_agg_search(self) -> Dict[str, Any]:
             if e.response.status_code >= 500:
                 raise  # Let circuit breaker track server errors
             raise ValueError(f"Vulnerability summary query failed: {e.response.status_code}")
-        except httpx.ConnectError:
+        except (httpx.ConnectError, httpx.TimeoutException):
             raise ConnectionError(f"Cannot connect to Wazuh Indexer at {self.host}:{self.port}")
 
     async def health_check(self) -> Dict[str, Any]:

src/wazuh_mcp_server/auth.py

@@ -65,7 +65,11 @@ class AuthManager:
 
     def __init__(self):
         self.secret_key = os.getenv("AUTH_SECRET_KEY", secrets.token_urlsafe(32))
-        self.token_lifetime = int(os.getenv("TOKEN_LIFETIME_HOURS", "24"))
+        try:
+            self.token_lifetime = int(os.getenv("TOKEN_LIFETIME_HOURS", "24"))
+        except (ValueError, TypeError):
+            logger.warning("Invalid TOKEN_LIFETIME_HOURS, defaulting to 24")
+            self.token_lifetime = 24
         self.api_keys: Dict[str, APIKey] = {}
         self.tokens: Dict[str, AuthToken] = {}
         self._default_api_key: Optional[str] = None  # Stores auto-generated key for display
@@ -295,7 +299,7 @@ async def verify_bearer_token(authorization: str) -> AuthToken:
     Raises:
         ValueError: If the token is invalid or expired
     """
-    if not authorization.startswith("Bearer "):
+    if not authorization or not authorization.startswith("Bearer "):
         raise ValueError("Invalid authorization header format")
 
     token = authorization[7:]  # Remove "Bearer " prefix

src/wazuh_mcp_server/monitoring.py

@@ -260,6 +260,9 @@ def __init__(self):
         self.collection_interval = 30  # seconds
         self.last_collection = 0
         self._collection_task = None
+        # Reuse a single Process instance so cpu_percent() measures delta between calls
+        # (psutil.Process.cpu_percent() returns 0.0 on the first call to a new instance)
+        self._process = psutil.Process()
 
     async def start_collection(self):
         """Start metrics collection task."""
@@ -292,12 +295,11 @@ async def _collect_system_metrics(self):
         """Collect system-level metrics."""
         try:
             # Memory usage
-            process = psutil.Process()
-            memory_info = process.memory_info()
+            memory_info = self._process.memory_info()
             SYSTEM_MEMORY_USAGE.set(memory_info.rss)
 
-            # CPU usage
-            cpu_percent = process.cpu_percent()
+            # CPU usage (uses delta since last call on the persistent _process instance)
+            cpu_percent = self._process.cpu_percent()
             SYSTEM_CPU_USAGE.set(cpu_percent)
 
         except Exception as e:

src/wazuh_mcp_server/oauth.py

@@ -430,7 +430,8 @@ async def authorize(
 
         # Validate response_type
         if response_type != "code":
-            return RedirectResponse(f"{redirect_uri}?error=unsupported_response_type&state={state or ''}")
+            params = urlencode({"error": "unsupported_response_type", "state": state or ""})
+            return RedirectResponse(f"{redirect_uri}?{params}")
 
         # For MCP servers, we auto-approve (the user already chose to connect)
         # In production, you might show a consent screen here

src/wazuh_mcp_server/resilience.py

@@ -54,6 +54,7 @@ def __init__(self, config: CircuitBreakerConfig):
         self.last_failure_time: Optional[float] = None
         self.next_retry_time: Optional[float] = None
         self._lock = asyncio.Lock()
+        self._half_open_trial_in_progress = False  # Only allow one trial request in HALF_OPEN
 
     def __call__(self, func: Callable) -> Callable:
         """Decorator to apply circuit breaker to function."""
@@ -71,7 +72,13 @@ async def _call(self, func: Callable, *args, **kwargs) -> Any:
         async with self._lock:
             if self.state == CircuitBreakerState.OPEN:
                 if self._should_attempt_reset():
+                    if self._half_open_trial_in_progress:
+                        # Another coroutine is already running the trial request — reject this one
+                        raise HTTPException(
+                            status_code=503, detail="Service temporarily unavailable - circuit breaker half-open trial in progress"
+                        )
                     self.state = CircuitBreakerState.HALF_OPEN
+                    self._half_open_trial_in_progress = True
                     logger.info(f"Circuit breaker {func.__name__} moved to HALF_OPEN")
                 else:
                     if self.config.fallback_function:
@@ -81,6 +88,11 @@ async def _call(self, func: Callable, *args, **kwargs) -> Any:
                         raise HTTPException(
                             status_code=503, detail="Service temporarily unavailable - circuit breaker open"
                         )
+            elif self.state == CircuitBreakerState.HALF_OPEN and self._half_open_trial_in_progress:
+                # HALF_OPEN with a trial already running — reject concurrent requests
+                raise HTTPException(
+                    status_code=503, detail="Service temporarily unavailable - circuit breaker half-open trial in progress"
+                )
 
         try:
             result = await func(*args, **kwargs)
@@ -104,6 +116,7 @@ def _should_attempt_reset(self) -> bool:
     async def _on_success(self, func_name: str):
         """Handle successful execution."""
         async with self._lock:
+            self._half_open_trial_in_progress = False
             if self.state == CircuitBreakerState.HALF_OPEN:
                 self.state = CircuitBreakerState.CLOSED
                 logger.info(f"Circuit breaker {func_name} reset to CLOSED")
@@ -113,6 +126,7 @@ async def _on_success(self, func_name: str):
     async def _on_failure(self, func_name: str, exception: Exception):
         """Handle failed execution."""
         async with self._lock:
+            self._half_open_trial_in_progress = False
             self.failure_count += 1
             self.last_failure_time = time.time()
             if self.failure_count >= self.config.failure_threshold:

src/wazuh_mcp_server/security.py

@@ -46,7 +46,7 @@ def __init__(self, param_name: str, message: str, suggestion: str = None):
 VALID_COMPLIANCE_FRAMEWORKS = {"PCI-DSS", "HIPAA", "SOX", "GDPR", "NIST"}
 
 # Regex patterns for parameter validation
-AGENT_ID_PATTERN = re.compile(r"^[0-9]{3,5}$")  # Wazuh agent IDs are numeric
+AGENT_ID_PATTERN = re.compile(r"^[0-9]{1,5}$")  # Wazuh agent IDs: 0 (manager) through 99999
 RULE_ID_PATTERN = re.compile(r"^[0-9]{1,6}$")  # Rule IDs are numeric
 ISO_TIMESTAMP_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}(T\d{2}:\d{2}:\d{2}(\.\d+)?(Z|[+-]\d{2}:?\d{2})?)?$")
 # Use ipaddress module for proper validation instead of regex
@@ -685,10 +685,15 @@ class SecurityManager:
 
     def __init__(self):
         self.metrics = SecurityMetrics()
-        self.rate_limiter = RateLimiter(
-            max_requests=int(os.getenv("RATE_LIMIT_REQUESTS", "100")),
-            window_seconds=int(os.getenv("RATE_LIMIT_WINDOW", "60")),
-        )
+        try:
+            max_req = int(os.getenv("RATE_LIMIT_REQUESTS", "100"))
+        except (ValueError, TypeError):
+            max_req = 100
+        try:
+            window = int(os.getenv("RATE_LIMIT_WINDOW", "60"))
+        except (ValueError, TypeError):
+            window = 60
+        self.rate_limiter = RateLimiter(max_requests=max_req, window_seconds=window)
         self.validator = SecurityValidator()
         self.trusted_proxies = {p.strip() for p in os.getenv("TRUSTED_PROXIES", "").split(",") if p.strip()}
 

src/wazuh_mcp_server/server.py

@@ -1861,9 +1861,9 @@ async def handle_tools_list(params: Dict[str, Any], session: MCPSession) -> Dict
         },
     ]
 
-    # Filter tools by session scopes: hide write tools from read-only tokens
+    # Filter tools by session scopes: hide write tools from read-only or unknown tokens
     auth_token = getattr(session, "_auth_token", None)
-    if auth_token and not auth_token.has_scope("wazuh:write"):
+    if not auth_token or not auth_token.has_scope("wazuh:write"):
         tools = [t for t in tools if t["name"] not in WRITE_SCOPE_TOOLS]
 
     # Pagination support per MCP spec
@@ -1881,9 +1881,15 @@ async def handle_tools_call(params: Dict[str, Any], session: MCPSession) -> Dict
     # Validate tool name
     validate_input(tool_name, max_length=100)
 
-    # Scope enforcement: check if the token has the required scope for this tool
+    # Scope enforcement: check if the token has the required scope for this tool.
+    # If auth_token is missing (should not happen in normal flow), deny write tools by default.
     auth_token = getattr(session, "_auth_token", None)
     required_scope = _get_tool_scope(tool_name)
+    if required_scope == "wazuh:write" and not auth_token:
+        raise ValueError(
+            f"Insufficient permissions: tool '{tool_name}' requires '{required_scope}' scope. "
+            f"Authentication token not found on session."
+        )
     if auth_token and not auth_token.has_scope(required_scope):
         raise ValueError(
             f"Insufficient permissions: tool '{tool_name}' requires '{required_scope}' scope. "
@@ -2060,6 +2066,7 @@ def _tool_error(text: str) -> dict:
             result = await wazuh_client.get_vulnerabilities(agent_id=agent_id, severity=severity, limit=limit)
             if compact:
                 result = _compact_vulns_result(result)
+            result = _add_truncation_warning(result, limit)
             _success = True
             return _tool_result(f"Vulnerabilities:\n{json.dumps(result, indent=2 if not compact else None)}")
 
@@ -2070,6 +2077,7 @@ def _tool_error(text: str) -> dict:
             result = await wazuh_client.get_critical_vulnerabilities(limit)
             if compact:
                 result = _compact_vulns_result(result)
+            result = _add_truncation_warning(result, limit)
             _success = True
             return _tool_result(f"Critical Vulnerabilities:\n{json.dumps(result, indent=2 if not compact else None)}")
 
@@ -2564,7 +2572,7 @@ async def mcp_endpoint(
                     status_code=404, detail="Session not found. Please start a new session with InitializeRequest."
                 )
             if existing_session.is_expired():
-                await sessions.delete(mcp_session_id)
+                await sessions.remove(mcp_session_id)
                 _initialized_sessions.pop(mcp_session_id, None)
                 raise HTTPException(
                     status_code=404, detail="Session expired. Please start a new session with InitializeRequest."
@@ -2776,25 +2784,25 @@ async def mcp_sse_endpoint(
         headers = {"Retry-After": str(retry_after)} if retry_after else {}
         raise HTTPException(status_code=429, detail="Rate limit exceeded", headers=headers)
 
-    # Track active connections
+    # Session validation: if client provides session ID but session doesn't exist, return 404
+    # Done BEFORE incrementing ACTIVE_CONNECTIONS to avoid counter leak on early errors.
+    if mcp_session_id:
+        existing_session = await sessions.get(mcp_session_id)
+        if not existing_session:
+            raise HTTPException(status_code=404, detail="Session not found")
+        session = existing_session
+        session.update_activity()
+        await sessions.set(mcp_session_id, session)
+    else:
+        session = await get_or_create_session(None, origin)
+    session.authenticated = True  # Mark as authenticated via bearer token
+    session._auth_token = auth_token  # Store token for scope checks in tool handlers
+
+    # Track active connections — only after validation passes.
+    # The SSE generator will decrement when the stream closes (track_connection=True).
     ACTIVE_CONNECTIONS.inc()
 
     try:
-        # Session validation: if client provides session ID but session doesn't exist, return 404
-        if mcp_session_id:
-            existing_session = await sessions.get(mcp_session_id)
-            if not existing_session:
-                raise HTTPException(status_code=404, detail="Session not found")
-            session = existing_session
-            session.update_activity()
-            await sessions.set(mcp_session_id, session)
-        else:
-            session = await get_or_create_session(None, origin)
-        session.authenticated = True  # Mark as authenticated via bearer token
-        session._auth_token = auth_token  # Store token for scope checks in tool handlers
-
-        # Return SSE stream (track_connection=True so ACTIVE_CONNECTIONS is
-        # decremented when the stream actually closes, not when this function returns)
         response = StreamingResponse(
             generate_sse_events(session, track_connection=True),
             media_type="text/event-stream",
@@ -2868,7 +2876,7 @@ async def mcp_streamable_http_endpoint(
                     status_code=404, detail="Session not found. Please start a new session with InitializeRequest."
                 )
             if existing_session.is_expired():
-                await sessions.delete(mcp_session_id)
+                await sessions.remove(mcp_session_id)
                 _initialized_sessions.pop(mcp_session_id, None)
                 raise HTTPException(
                     status_code=404, detail="Session expired. Please start a new session with InitializeRequest."