From 020d19600db8fcea120778bd585402ab6afcae73 Mon Sep 17 00:00:00 2001 From: Gianluca Brigandi Date: Fri, 5 Dec 2025 15:52:19 -0800 Subject: [PATCH] * Using latest 0.1.8 wazuh-client-rs craate which fixes issue with ordering: (#17) * Improved unmarshaling for indexer responses * Other minor changes. --- Cargo.toml | 2 +- src/tools/alerts.rs | 40 +++++++++++++++++++++++++++++++++- src/tools/vulnerabilities.rs | 2 +- tests/mcp_stdio_test.rs | 2 +- tests/mock_wazuh_server.rs | 16 +++++++------- tests/rmcp_integration_test.rs | 5 ++++- 6 files changed, 54 insertions(+), 13 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 0484507..30872f2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ repository = "https://github.com/gbrigandi/mcp-server-wazuh" readme = "README.md" [dependencies] -wazuh-client = "0.1.7" +wazuh-client = "0.1.8" rmcp = { version = "0.1.5", features = ["server", "transport-io"] } tokio = { version = "1", features = ["full"] } reqwest = { version = "0.12", features = ["json", "rustls-tls"], default-features = false } diff --git a/src/tools/alerts.rs b/src/tools/alerts.rs index 2212f5a..b6ac524 100644 --- a/src/tools/alerts.rs +++ b/src/tools/alerts.rs @@ -79,10 +79,48 @@ impl AlertTools { .and_then(|l| l.as_u64()) .unwrap_or(0); - let formatted_text = format!( + // Extract source IP from data.srcip (common for SSH, network alerts) + let src_ip = source.get("data") + .and_then(|d| d.get("srcip")) + .and_then(|ip| ip.as_str()) + .or_else(|| source.get("data") + .and_then(|d| d.get("src_ip")) + .and_then(|ip| ip.as_str())) + .unwrap_or(""); + + // Extract destination IP if available + let dst_ip = source.get("data") + .and_then(|d| d.get("dstip")) + .and_then(|ip| ip.as_str()) + .or_else(|| source.get("data") + .and_then(|d| d.get("dst_ip")) + .and_then(|ip| ip.as_str())) + .unwrap_or(""); + + // Extract source user if available + let src_user = source.get("data") + .and_then(|d| d.get("srcuser")) + .and_then(|u| u.as_str()) + .or_else(|| source.get("data") + .and_then(|d| d.get("dstuser")) + .and_then(|u| u.as_str())) + .unwrap_or(""); + + // Build formatted text with optional fields + let mut formatted_text = format!( "Alert ID: {}\nTime: {}\nAgent: {}\nLevel: {}\nDescription: {}", id, timestamp, agent_name, rule_level, description ); + + if !src_ip.is_empty() { + formatted_text.push_str(&format!("\nSource IP: {}", src_ip)); + } + if !dst_ip.is_empty() { + formatted_text.push_str(&format!("\nDestination IP: {}", dst_ip)); + } + if !src_user.is_empty() { + formatted_text.push_str(&format!("\nUser: {}", src_user)); + } Content::text(formatted_text) }) .collect(); diff --git a/src/tools/vulnerabilities.rs b/src/tools/vulnerabilities.rs index cb78e90..dc520bb 100644 --- a/src/tools/vulnerabilities.rs +++ b/src/tools/vulnerabilities.rs @@ -101,7 +101,7 @@ impl VulnerabilityTools { params .severity .as_deref() - .and_then(VulnerabilitySeverity::from_str), + .and_then(VulnerabilitySeverity::parse), ) .await; diff --git a/tests/mcp_stdio_test.rs b/tests/mcp_stdio_test.rs index b4e4b5c..cf1438b 100644 --- a/tests/mcp_stdio_test.rs +++ b/tests/mcp_stdio_test.rs @@ -54,7 +54,7 @@ impl McpStdioClient { fn read_response(&mut self) -> Result> { let mut line = String::new(); self.stdout.read_line(&mut line)?; - let response: Value = serde_json::from_str(&line.trim())?; + let response: Value = serde_json::from_str(line.trim())?; Ok(response) } diff --git a/tests/mock_wazuh_server.rs b/tests/mock_wazuh_server.rs index 0aa3da9..a907d3d 100644 --- a/tests/mock_wazuh_server.rs +++ b/tests/mock_wazuh_server.rs @@ -260,7 +260,7 @@ mod tests { let client = reqwest::Client::new(); let response = client - .post(&format!("{}/security/user/authenticate", mock_server.url())) + .post(format!("{}/security/user/authenticate", mock_server.url())) .json(&json!({"username": "admin", "password": "admin"})) .send() .await @@ -277,7 +277,7 @@ mod tests { let client = reqwest::Client::new(); let response = client - .post(&format!("{}/wazuh-alerts*/_search", mock_server.url())) + .post(format!("{}/wazuh-alerts*/_search", mock_server.url())) .json(&json!({"query": {"match_all": {}}})) .send() .await @@ -294,9 +294,9 @@ mod tests { async fn test_empty_alerts_server() { let mock_server = MockWazuhServer::with_empty_alerts(); let client = reqwest::Client::new(); - + let response = client - .post(&format!("{}/wazuh-alerts*/_search", mock_server.url())) + .post(format!("{}/wazuh-alerts*/_search", mock_server.url())) .json(&json!({"query": {"match_all": {}}})) .send() .await @@ -312,9 +312,9 @@ mod tests { async fn test_auth_error_server() { let mock_server = MockWazuhServer::with_auth_error(); let client = reqwest::Client::new(); - + let response = client - .post(&format!("{}/security/user/authenticate", mock_server.url())) + .post(format!("{}/security/user/authenticate", mock_server.url())) .json(&json!({"username": "admin", "password": "wrong"})) .send() .await @@ -327,9 +327,9 @@ mod tests { async fn test_alerts_error_server() { let mock_server = MockWazuhServer::with_alerts_error(); let client = reqwest::Client::new(); - + let response = client - .post(&format!("{}/wazuh-alerts*/_search", mock_server.url())) + .post(format!("{}/wazuh-alerts*/_search", mock_server.url())) .json(&json!({"query": {"match_all": {}}})) .send() .await diff --git a/tests/rmcp_integration_test.rs b/tests/rmcp_integration_test.rs index c8de7ce..686039e 100644 --- a/tests/rmcp_integration_test.rs +++ b/tests/rmcp_integration_test.rs @@ -1,8 +1,11 @@ //! Integration tests for the rmcp-based Wazuh MCP Server -//! +//! //! These tests verify the MCP server functionality using a mock Wazuh API server. //! Tests cover tool registration, parameter validation, alert retrieval, and error handling. +// Allow holding mutex guard across await in tests - this is intentional for test serialization +#![allow(clippy::await_holding_lock)] + use std::process::{Child, Command, Stdio}; use std::io::{BufRead, BufReader, Write}; use std::time::Duration;