Fix 17 bugs, gaps, and performance issues found in end-to-end audit

Bugs:
- Fix _run_sync catching its own RuntimeError safety guard
- Fix async __contains__ (Python 'in' operator never awaits)
- Fix auth bypass on root / endpoint allowing unauthenticated access
- Remove dead create_auth_endpoints code; fix /auth/token to use auth_manager
- Fix search_security_events silently ignoring query parameter
- Fix cleanup_expired missing timeout_minutes parameter
- Register monitoring middleware for request tracking
- Fix /metrics using default registry instead of custom REGISTRY

Gaps:
- Sync time range values: add 12h/1d across security.py and wazuh_client.py
- Add 'pending' to agent status enum in tool schema

Enhancements:
- Migrate MCPResponse from deprecated Pydantic v1 dict() to model_dump()
- Register security middleware for response headers
- Expand test suite from 10 to 33 tests covering new fixes

Performance:
- Refactor get_alerts to use _search helper with retry logic
- Add max size guard to unbounded _initialized_sessions dict
- Replace O(n*m) json.dumps per-alert with recursive dict search

Author: alokemajumder

src/wazuh_mcp_server/api/wazuh_client.py

@@ -17,7 +17,27 @@
 logger = logging.getLogger(__name__)
 
 # Time range to hours mapping for indexer-based queries
-_TIME_RANGE_HOURS = {"1h": 1, "6h": 6, "12h": 12, "24h": 24, "7d": 168, "30d": 720}
+_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:
@@ -462,9 +482,7 @@ async def analyze_alert_patterns(self, time_range: str, min_frequency: int) -> D
                     "level": rule.get("level", 0),
                 }
             rule_counts[rule_id]["count"] += 1
-        patterns = [
-            {"rule_id": k, **v} for k, v in rule_counts.items() if v["count"] >= min_frequency
-        ]
+        patterns = [{"rule_id": k, **v} for k, v in rule_counts.items() if v["count"] >= min_frequency]
         patterns.sort(key=lambda x: x["count"], reverse=True)
         return {
             "data": {
@@ -476,11 +494,20 @@ async def analyze_alert_patterns(self, time_range: str, min_frequency: int) -> D
         }
 
     async def search_security_events(self, query: str, time_range: str, limit: int) -> Dict[str, Any]:
-        """Search security events via the Wazuh Indexer."""
+        """Search security events via the Wazuh Indexer with query filtering."""
         if not self._indexer_client:
             raise IndexerNotConfiguredError()
         start = self._time_range_to_start(time_range)
-        return await self._indexer_client.get_alerts(limit=limit, timestamp_start=start)
+        # Fetch a larger batch from the indexer, then filter by query
+        fetch_limit = min(limit * 5, 2000)
+        result = await self._indexer_client.get_alerts(limit=fetch_limit, timestamp_start=start)
+        if query:
+            alerts = result.get("data", {}).get("affected_items", [])
+            query_lower = query.lower()
+            filtered = [a for a in alerts if _dict_contains_text(a, query_lower)]
+            result["data"]["affected_items"] = filtered[:limit]
+            result["data"]["total_affected_items"] = len(filtered)
+        return result
 
     async def get_running_agents(self) -> Dict[str, Any]:
         """Get running agents."""
@@ -589,11 +616,8 @@ async def analyze_security_threat(self, indicator: str, indicator_type: str) ->
             raise IndexerNotConfiguredError()
         result = await self._indexer_client.get_alerts(limit=100)
         alerts = result.get("data", {}).get("affected_items", [])
-        matches = []
-        for alert in alerts:
-            alert_str = json.dumps(alert)
-            if indicator.lower() in alert_str.lower():
-                matches.append(alert)
+        indicator_lower = indicator.lower()
+        matches = [a for a in alerts if _dict_contains_text(a, indicator_lower)]
         return {
             "data": {
                 "indicator": indicator,
@@ -611,9 +635,9 @@ async def check_ioc_reputation(self, indicator: str, indicator_type: str) -> Dic
         alerts = result.get("data", {}).get("affected_items", [])
         occurrences = 0
         max_level = 0
+        indicator_lower = indicator.lower()
         for alert in alerts:
-            alert_str = json.dumps(alert)
-            if indicator.lower() in alert_str.lower():
+            if _dict_contains_text(alert, indicator_lower):
                 occurrences += 1
                 level = alert.get("rule", {}).get("level", 0)
                 if isinstance(level, int) and level > max_level:
@@ -887,14 +911,12 @@ async def check_blocked_ip(self, ip_address: str, agent_id: str = None) -> Dict[
             raise IndexerNotConfiguredError()
         result = await self._indexer_client.get_alerts(limit=50)
         alerts = result.get("data", {}).get("affected_items", [])
-        matches = [a for a in alerts if ip_address in json.dumps(a) and "firewall-drop" in json.dumps(a)]
+        matches = [a for a in alerts if _dict_contains_text(a, ip_address) and _dict_contains_text(a, "firewall-drop")]
         return {"data": {"ip_address": ip_address, "blocked": len(matches) > 0, "matching_alerts": len(matches)}}
 
     async def check_agent_isolation(self, agent_id: str) -> Dict[str, Any]:
         """Check agent isolation status."""
-        result = await self._request(
-            "GET", "/agents", params={"agents_list": agent_id, "select": "id,name,status"}
-        )
+        result = await self._request("GET", "/agents", params={"agents_list": agent_id, "select": "id,name,status"})
         agents = result.get("data", {}).get("affected_items", [])
         if not agents:
             raise ValueError(f"Agent {agent_id} not found")
@@ -910,9 +932,7 @@ async def check_agent_isolation(self, agent_id: str) -> Dict[str, Any]:
 
     async def check_process(self, agent_id: str, process_id: int) -> Dict[str, Any]:
         """Check if a process is still running on an agent."""
-        result = await self._request(
-            "GET", f"/syscollector/{agent_id}/processes", params={"limit": 500}
-        )
+        result = await self._request("GET", f"/syscollector/{agent_id}/processes", params={"limit": 500})
         processes = result.get("data", {}).get("affected_items", [])
         running = any(str(p.get("pid")) == str(process_id) for p in processes)
         return {"data": {"agent_id": agent_id, "process_id": process_id, "running": running}}
@@ -930,9 +950,7 @@ async def check_user_status(self, agent_id: str, username: str) -> Dict[str, Any
 
     async def check_file_quarantine(self, agent_id: str, file_path: str) -> Dict[str, Any]:
         """Check if a file has been quarantined via FIM events."""
-        result = await self._request(
-            "GET", "/syscheck", params={"agents_list": agent_id, "q": f"file={file_path}"}
-        )
+        result = await self._request("GET", "/syscheck", params={"agents_list": agent_id, "q": f"file={file_path}"})
         events = result.get("data", {}).get("affected_items", [])
         quarantined = any(e.get("type") == "deleted" or "quarantine" in str(e) for e in events)
         return {"data": {"agent_id": agent_id, "file_path": file_path, "quarantined": quarantined}}

src/wazuh_mcp_server/api/wazuh_indexer.py

@@ -95,22 +95,27 @@ async def _ensure_initialized(self):
         retry=retry_if_exception_type((httpx.RequestError, httpx.HTTPStatusError)),
         reraise=True,
     )
-    async def _search(self, index: str, query: Dict[str, Any], size: int = 100) -> Dict[str, Any]:
+    async def _search(
+        self, index: str, query: Dict[str, Any], size: int = 100, sort: Optional[list] = None
+    ) -> Dict[str, Any]:
         """
         Execute a search query against the Wazuh Indexer.
 
         Args:
             index: Index pattern to search
             query: Elasticsearch query DSL
             size: Maximum number of results
+            sort: Optional sort specification
 
         Returns:
             Search results from the indexer
         """
         await self._ensure_initialized()
 
         url = f"{self.base_url}/{index}/_search"
-        body = {"query": query, "size": size}
+        body: Dict[str, Any] = {"query": query, "size": size}
+        if sort:
+            body["sort"] = sort
 
         try:
             response = await self.client.post(url, json=body, headers={"Content-Type": "application/json"})
@@ -151,8 +156,6 @@ async def get_alerts(
         Returns:
             Alert data in standard Wazuh format
         """
-        await self._ensure_initialized()
-
         # Build bool query with must clauses for each non-empty filter
         must_clauses: list = []
 
@@ -183,28 +186,8 @@ async def get_alerts(
         else:
             query = {"match_all": {}}
 
-        # Build the full search body with sort by timestamp desc
-        url = f"{self.base_url}/{ALERTS_INDEX}/_search"
-        body = {
-            "query": query,
-            "size": limit,
-            "sort": [{"timestamp": {"order": "desc"}}],
-        }
-
-        try:
-            response = await self.client.post(url, json=body, headers={"Content-Type": "application/json"})
-            response.raise_for_status()
-            try:
-                result = response.json()
-            except (json.JSONDecodeError, ValueError):
-                raise ValueError("Invalid JSON response from Wazuh Indexer alerts query")
-        except httpx.HTTPStatusError as e:
-            logger.error(f"Alerts query failed: {e.response.status_code} - {e.response.text}")
-            raise ValueError(f"Alerts query failed: {e.response.status_code}")
-        except httpx.ConnectError:
-            raise ConnectionError(f"Cannot connect to Wazuh Indexer at {self.host}:{self.port}")
-        except httpx.TimeoutException:
-            raise ConnectionError(f"Timeout connecting to Wazuh Indexer at {self.host}:{self.port}")
+        # Use _search helper for consistent retry logic (sorted by timestamp desc)
+        result = await self._search(ALERTS_INDEX, query, size=limit, sort=[{"timestamp": {"order": "desc"}}])
 
         # Transform to standard Wazuh format
         hits = result.get("hits", {})

src/wazuh_mcp_server/auth.py

@@ -14,7 +14,6 @@
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, List, Optional
 
-from fastapi import Header, HTTPException
 from jose import jwt
 from jose.exceptions import ExpiredSignatureError, JWTError
 from pydantic import BaseModel, Field
@@ -268,20 +267,6 @@ def get_stats(self) -> Dict[str, Any]:
 auth_manager = AuthManager()
 
 
-class TokenRequest(BaseModel):
-    """Token request model."""
-
-    api_key: str = Field(description="API key to exchange for token")
-
-
-class TokenResponse(BaseModel):
-    """Token response model."""
-
-    token: str = Field(description="Authentication token")
-    expires_in: int = Field(description="Token lifetime in seconds")
-    token_type: str = Field(default="Bearer")
-
-
 async def verify_bearer_token(authorization: str) -> AuthToken:
     """
     Verify bearer token from Authorization header.
@@ -379,57 +364,3 @@ def verify_token(token: str, secret_key: str) -> Dict[str, Any]:
         raise ValueError("Token has expired")
     except JWTError:
         raise ValueError("Invalid token")
-
-
-async def create_auth_endpoints(app):
-    """Add authentication endpoints to FastAPI app."""
-
-    @app.post("/auth/token", response_model=TokenResponse)
-    async def create_token(request: TokenRequest):
-        """Exchange API key for authentication token."""
-        token = auth_manager.create_token(request.api_key)
-        if not token:
-            raise HTTPException(status_code=401, detail="Invalid API key")
-
-        token_obj = auth_manager.tokens[token]
-        expires_in = int((token_obj.expires_at - datetime.now(timezone.utc)).total_seconds())
-
-        return TokenResponse(token=token, expires_in=expires_in)
-
-    @app.get("/auth/validate")
-    async def validate_token(authorization: str = Header(description="Bearer token")):
-        """Validate authentication token."""
-        try:
-            token_obj = await verify_bearer_token(authorization)
-            return {
-                "valid": True,
-                "api_key_id": token_obj.api_key_id,
-                "scopes": token_obj.scopes,
-                "expires_at": token_obj.expires_at.isoformat() if token_obj.expires_at else None,
-            }
-        except ValueError as e:
-            raise HTTPException(status_code=401, detail=str(e))
-
-    @app.post("/auth/revoke")
-    async def revoke_token(authorization: str = Header(description="Bearer token")):
-        """Revoke authentication token."""
-        if not authorization.startswith("Bearer "):
-            raise HTTPException(status_code=400, detail="Invalid authorization header")
-
-        token = authorization[7:]
-        if auth_manager.revoke_token(token):
-            return {"revoked": True}
-        else:
-            raise HTTPException(status_code=404, detail="Token not found")
-
-    @app.get("/auth/stats")
-    async def auth_stats(authorization: str = Header(description="Bearer token")):
-        """Get authentication statistics (requires admin scope)."""
-        try:
-            token_obj = await verify_bearer_token(authorization)
-            if not token_obj.has_scope("admin"):
-                raise HTTPException(status_code=403, detail="Admin scope required")
-
-            return auth_manager.get_stats()
-        except ValueError as e:
-            raise HTTPException(status_code=401, detail=str(e))

src/wazuh_mcp_server/security.py

@@ -39,7 +39,7 @@ def __init__(self, param_name: str, message: str, suggestion: str = None):
 
 
 # Valid enum values for tool parameters
-VALID_TIME_RANGES = {"1h", "6h", "24h", "7d", "1d", "30d"}
+VALID_TIME_RANGES = {"1h", "6h", "12h", "24h", "7d", "1d", "30d"}
 VALID_SEVERITIES = {"low", "medium", "high", "critical"}
 VALID_AGENT_STATUSES = {"active", "disconnected", "never_connected", "pending"}
 VALID_INDICATOR_TYPES = {"ip", "hash", "domain", "url"}
@@ -354,7 +354,9 @@ def validate_file_path(value: Any, required: bool = False, param_name: str = "fi
         raise ToolValidationError(param_name, "contains path traversal", "Path must not contain '..'")
 
     if len(file_path) > 500:
-        raise ToolValidationError(param_name, f"too long ({len(file_path)} chars)", "Path must be 500 characters or less")
+        raise ToolValidationError(
+            param_name, f"too long ({len(file_path)} chars)", "Path must be 500 characters or less"
+        )
 
     return file_path