NeoAnd/mcp-server-wazuh
1
0
Fork 0
mirror of https://github.com/gbrigandi/mcp-server-wazuh.git synced 2025-12-22 04:52:19 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

* Ported code to RMCP

Browse Source 
* Implemented unit and e2e testing
* Other fixes and enhancements
This commit is contained in:
Gianluca Brigandi
2025-05-22 20:02:41 -07:00
parent 6661523c0f
commit d59d67b8db
28 changed files with 1519 additions and 2778 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
tests/README.md
View File

@@ -1,60 +1,133 @@
# Wazuh MCP Server Tests
This directory contains tests for the Wazuh MCP Server, including end-to-end tests that simulate a client interacting with the server.
This directory contains tests for the Wazuh MCP Server using the rmcp framework, including unit tests, integration tests with mock Wazuh API, and end-to-end MCP protocol tests.
## Test Files
- `e2e_client_test.rs`: End-to-end test for MCP client interacting with Wazuh MCP server
- `integration_test.rs`: Integration test for Wazuh MCP Server with a mock Wazuh API
- `mcp_client.rs`: Reusable MCP client implementation
- `mcp_client_cli.rs`: Command-line tool for interacting with the MCP server
- `rmcp_integration_test.rs`: Integration tests for the rmcp-based MCP server using a mock Wazuh API.
- `mock_wazuh_server.rs`: Mock Wazuh API server implementation, used by the integration tests.
- `mcp_stdio_test.rs`: Tests for MCP protocol communication via stdio, focusing on initialization, compliance, concurrent requests, and error handling for invalid/unsupported messages.
- `run_tests.sh`: A shell script that automates running the various test suites.
## Testing Strategy
### 1. Mock Wazuh Server Tests
Tests the MCP server with a mock Wazuh API to verify:
- Tool registration and schema generation
- Alert retrieval and formatting
- Error handling for API failures
- Parameter validation
### 2. MCP Protocol Tests
Tests the MCP protocol implementation (primarily in `mcp_stdio_test.rs`):
- Initialize handshake.
- Tools listing (basic, without requiring a live Wazuh connection).
- Handling of invalid JSON-RPC requests and unsupported methods.
- Behavior with concurrent requests.
- JSON-RPC 2.0 compliance.
(Note: Full tool execution, like `tools/call`, is primarily tested in `rmcp_integration_test.rs` using the mock Wazuh server.)
### 3. Unit Tests
Tests individual components and modules, typically run via `cargo test --lib`. These may include:
- Wazuh client logic (e.g., authentication, request formation, response parsing).
- Alert data transformation and formatting.
- Internal error handling mechanisms and utility functions.
## Running the Tests
To run all tests:
### Run All Tests
```bash
cargo test
```
To run a specific test:
### Run Specific Test Categories
```bash
cargo test --test e2e_client_test
cargo test --test integration_test
# Integration tests with mock Wazuh
cargo test --test rmcp_integration_test
# MCP protocol tests
cargo test --test mcp_stdio_test
# Unit tests
cargo test --lib
```
## Using the MCP Client CLI
The MCP Client CLI can be used to interact with the MCP server for testing purposes:
### Run Tests with Logging
```bash
# Build the CLI
cargo build --bin mcp_client_cli
# Run the CLI
MCP_SERVER_URL=http://localhost:8000 ./target/debug/mcp_client_cli get-data
MCP_SERVER_URL=http://localhost:8000 ./target/debug/mcp_client_cli health
MCP_SERVER_URL=http://localhost:8000 ./target/debug/mcp_client_cli query '{"severity": "high"}'
RUST_LOG=debug cargo test -- --nocapture
```
## Test Environment Variables
The tests use the following environment variables:
The tests support the following environment variables:
- `MCP_SERVER_URL`: URL of the MCP server (default: http://localhost:8000)
- `WAZUH_HOST`: Hostname of the Wazuh API server
- `WAZUH_PORT`: Port of the Wazuh API server
- `WAZUH_USER`: Username for Wazuh API authentication
- `WAZUH_PASS`: Password for Wazuh API authentication
- `VERIFY_SSL`: Whether to verify SSL certificates (default: false)
- `RUST_LOG`: Log level for the tests (default: info)
- `RUST_LOG`: Log level for tests (default: info)
- `TEST_WAZUH_HOST`: Real Wazuh host for integration tests (optional)
- `TEST_WAZUH_PORT`: Real Wazuh port for integration tests (optional)
- `TEST_WAZUH_USER`: Real Wazuh username for integration tests (optional)
- `TEST_WAZUH_PASS`: Real Wazuh password for integration tests (optional)
## Mock Wazuh API Server
The tests use a mock Wazuh API server to simulate the Wazuh API. The mock server provides:
The mock server simulates a real Wazuh Indexer API with:
- Authentication endpoint: `/security/user/authenticate`
- Alerts endpoint: `/wazuh-alerts-*_search`
### Authentication Endpoint
- `POST /security/user/authenticate`
- Returns mock JWT token
The mock server returns predefined responses for these endpoints, allowing the tests to run without a real Wazuh API server.
### Alerts Endpoint
- `POST /wazuh-alerts-*/_search` (Note: The Wazuh API typically uses POST for search queries with a body)
- Returns configurable mock alert data
- Supports different scenarios (success, empty, error)
### Configurable Responses
The mock server can be configured to return:
- Successful responses with sample alerts
- Empty responses (no alerts)
- Error responses (500, 401, etc.)
- Malformed responses for error testing
## Testing Without Real Wazuh
All tests can run without a real Wazuh instance by using the mock server. This allows for:
- **CI/CD Integration**: Tests run in any environment
- **Deterministic Results**: Predictable test data
- **Error Scenario Testing**: Simulate various failure modes
- **Fast Execution**: No network dependencies
## Testing With a Real Wazuh Instance (Manual End-to-End)
The automated test suites (`cargo test`) use mock servers or no Wazuh connection. To perform end-to-end testing with a real Wazuh instance, you need to run the server application itself and interact with it manually or via a separate client.
1. **Set up your Wazuh environment:** Ensure you have a running Wazuh instance (Indexer/API).
2. **Configure Environment Variables:** Set the standard runtime environment variables for the server to connect to your Wazuh instance:
```bash
export WAZUH_HOST="your-wazuh-indexer-host" # e.g., localhost or an IP address
export WAZUH_PORT="9200" # Or your Wazuh Indexer port
export WAZUH_USER="your-wazuh-api-user"
export WAZUH_PASS="your-wazuh-api-password"
export VERIFY_SSL="false" # Set to "true" if your Wazuh API uses a valid CA-signed SSL certificate
# export RUST_LOG="debug" # For more detailed server logs
```
## Manual Testing
### Using stdio directly
The server communicates over stdin/stdout. You can send commands by piping them to the process:
```bash
# Example: Send an initialize request
echo '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}' | cargo run --bin mcp-server-wazuh
```
### Using the test script
```bash
# Run the provided test script
./tests/run_tests.sh
```
This script will:
1. Start the MCP server with mock Wazuh configuration
2. Send a series of MCP commands
3. Verify responses
4. Clean up processes

194
tests/e2e_client_test.rs
View File

@@ -1,194 +0,0 @@
use anyhow::Result;
use httpmock::prelude::*;
use reqwest::Client;
use serde_json::{json, Value};
use std::process::{Child, Command};
use std::time::Duration;
use tokio::time::sleep;
use uuid::Uuid;
struct MockWazuhServer {
server: MockServer,
}
impl MockWazuhServer {
fn new() -> Self {
let server = MockServer::start();
server.mock(|when, then| {
when.method(POST)
.path("/security/user/authenticate")
.header("Authorization", "Basic YWRtaW46YWRtaW4=");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"jwt": "mock.jwt.token"
}));
});
server.mock(|when, then| {
when.method(GET)
.path("/wazuh-alerts-*_search")
.header("Authorization", "Bearer mock.jwt.token");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"hits": {
"hits": [
{
"_source": {
"id": "12345",
"category": "intrusion_detection",
"severity": "high",
"description": "Possible intrusion attempt detected",
"data": {
"source_ip": "192.168.1.100",
"destination_ip": "10.0.0.1",
"port": 22
},
"notes": "Test alert"
}
},
{
"_source": {
"id": "67890",
"category": "malware",
"severity": "critical",
"description": "Malware detected on system",
"data": {
"file_path": "/tmp/malicious.exe",
"hash": "abcdef123456",
"signature": "EICAR-Test-File"
}
}
}
]
}
}));
});
Self { server }
}
fn url(&self) -> String {
self.server.url("")
}
}
struct McpClient {
client: Client,
base_url: String,
}
impl McpClient {
fn new(base_url: String) -> Self {
let client = Client::builder()
.timeout(Duration::from_secs(10))
.build()
.expect("Failed to create HTTP client");
Self { client, base_url }
}
async fn get_mcp_data(&self) -> Result<Vec<Value>> {
let url = format!("{}/mcp", self.base_url);
let response = self.client.get(&url).send().await?;
if !response.status().is_success() {
anyhow::bail!("MCP request failed with status: {}", response.status());
}
let data = response.json::<Vec<Value>>().await?;
Ok(data)
}
async fn check_health(&self) -> Result<Value> {
let url = format!("{}/health", self.base_url);
let response = self.client.get(&url).send().await?;
if !response.status().is_success() {
anyhow::bail!("Health check failed with status: {}", response.status());
}
let data = response.json::<Value>().await?;
Ok(data)
}
}
fn start_mcp_server(wazuh_url: &str, port: u16) -> Child {
let server_id = Uuid::new_v4().to_string();
let wazuh_host_port: Vec<&str> = wazuh_url.trim_start_matches("http://").split(':').collect();
let wazuh_host = wazuh_host_port[0];
let wazuh_port = wazuh_host_port[1];
Command::new("cargo")
.args(["run", "--"])
.env("WAZUH_HOST", wazuh_host)
.env("WAZUH_PORT", wazuh_port)
.env("WAZUH_USER", "admin")
.env("WAZUH_PASS", "admin")
.env("VERIFY_SSL", "false")
.env("MCP_SERVER_PORT", port.to_string())
.env("RUST_LOG", "info")
.env("SERVER_ID", server_id)
.spawn()
.expect("Failed to start MCP server")
}
#[tokio::test]
async fn test_mcp_client_integration() -> Result<()> {
let mock_wazuh = MockWazuhServer::new();
let wazuh_url = mock_wazuh.url();
let mcp_port = 8765;
let mut mcp_server = start_mcp_server(&wazuh_url, mcp_port);
sleep(Duration::from_secs(2)).await;
let mcp_client = McpClient::new(format!("http://localhost:{}", mcp_port));
let health_data = mcp_client.check_health().await?;
assert_eq!(health_data["status"], "ok");
assert_eq!(health_data["service"], "wazuh-mcp-server");
let mcp_data = mcp_client.get_mcp_data().await?;
assert_eq!(mcp_data.len(), 2);
let first_message = &mcp_data[0];
assert_eq!(first_message["protocol_version"], "1.0");
assert_eq!(first_message["source"], "Wazuh");
assert_eq!(first_message["event_type"], "alert");
let context = &first_message["context"];
assert_eq!(context["id"], "12345");
assert_eq!(context["category"], "intrusion_detection");
assert_eq!(context["severity"], "high");
assert_eq!(
context["description"],
"Possible intrusion attempt detected"
);
let data = &context["data"];
assert_eq!(data["source_ip"], "192.168.1.100");
assert_eq!(data["destination_ip"], "10.0.0.1");
assert_eq!(data["port"], 22);
let second_message = &mcp_data[1];
let context = &second_message["context"];
assert_eq!(context["id"], "67890");
assert_eq!(context["category"], "malware");
assert_eq!(context["severity"], "critical");
assert_eq!(context["description"], "Malware detected on system");
let data = &context["data"];
assert_eq!(data["file_path"], "/tmp/malicious.exe");
assert_eq!(data["hash"], "abcdef123456");
assert_eq!(data["signature"], "EICAR-Test-File");
mcp_server.kill().expect("Failed to kill MCP server");
Ok(())
}

420
tests/integration_test.rs
View File

@@ -1,420 +0,0 @@
use anyhow::Result;
use chrono::{DateTime, Utc};
use httpmock::prelude::*;
use once_cell::sync::Lazy;
use serde_json::json;
use std::net::TcpListener;
use std::process::{Child, Command};
use std::sync::Mutex;
use std::time::Duration;
use tokio::time::sleep;
mod mcp_client;
use mcp_client::{McpClient, McpClientTrait, McpMessage};
static TEST_MUTEX: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
fn find_available_port() -> u16 {
let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind to random port");
let port = listener
.local_addr()
.expect("Failed to get local address")
.port();
drop(listener);
port
}
struct MockWazuhServer {
server: MockServer,
}
impl MockWazuhServer {
fn new() -> Self {
let server = MockServer::start();
server.mock(|when, then| {
when.method(POST).path("/security/user/authenticate");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"jwt": "mock.jwt.token"
}));
});
server.mock(|when, then| {
when.method(GET).path("/wazuh-alerts-*_search");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"hits": {
"hits": [
{
"_source": {
"id": "12345",
"category": "intrusion_detection",
"severity": "high",
"description": "Possible intrusion attempt detected",
"data": {
"source_ip": "192.168.1.100",
"destination_ip": "10.0.0.1",
"port": 22
},
"notes": "Test alert"
}
},
{
"_source": {
"id": "67890",
"category": "malware",
"severity": "critical",
"description": "Malware detected on system",
"data": {
"file_path": "/tmp/malicious.exe",
"hash": "abcdef123456",
"signature": "EICAR-Test-File"
}
}
}
]
}
}));
});
Self { server }
}
fn url(&self) -> String {
self.server.url("")
}
fn host(&self) -> String {
let url = self.url();
let parts: Vec<&str> = url.trim_start_matches("http://").split(':').collect();
parts[0].to_string()
}
fn port(&self) -> u16 {
let url = self.url();
let parts: Vec<&str> = url.trim_start_matches("http://").split(':').collect();
parts[1].parse().unwrap()
}
}
fn setup_mock_wazuh_server() -> MockServer {
let server = MockServer::start();
server.mock(|when, then| {
when.method(POST).path("/security/user/authenticate");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({ "jwt": "mock.jwt.token" }));
});
server.mock(|when, then| {
when.method(GET)
.path_matches(Regex::new(r"/wazuh-alerts-.*_search").unwrap());
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"hits": {
"hits": [
{
"_source": {
"id": "12345",
"timestamp": "2024-01-01T10:00:00.000Z",
"rule": {
"level": 9,
"description": "Possible intrusion attempt detected",
"groups": ["intrusion_detection", "pci_dss"]
},
"agent": { "id": "001", "name": "test-agent" },
"data": {
"source_ip": "192.168.1.100",
"destination_ip": "10.0.0.1",
"port": 22
}
}
},
{
"_source": {
"id": "67890",
"timestamp": "2024-01-01T11:00:00.000Z",
"rule": {
"level": 12,
"description": "Malware detected on system",
"groups": ["malware"]
},
"agent": { "id": "002", "name": "another-agent" },
"data": {
"file_path": "/tmp/malicious.exe",
"hash": "abcdef123456",
"signature": "EICAR-Test-File"
}
}
}
]
}
}));
});
server
}
fn get_host_port(server: &MockServer) -> (String, u16) {
let url = server.url("");
let parts: Vec<&str> = url.trim_start_matches("http://").split(':').collect();
let host = parts[0].to_string();
let port = parts[1].parse().unwrap();
(host, port)
}
fn start_mcp_server(wazuh_host: &str, wazuh_port: u16, mcp_port: u16) -> Child {
Command::new("cargo")
.args(["run", "--"])
.env("WAZUH_HOST", wazuh_host)
.env("WAZUH_PORT", wazuh_port.to_string())
.env("WAZUH_USER", "admin")
.env("WAZUH_PASS", "admin")
.env("VERIFY_SSL", "false")
.env("MCP_SERVER_PORT", mcp_port.to_string())
.env("RUST_LOG", "info")
.spawn()
.expect("Failed to start MCP server")
}
#[tokio::test]
async fn test_mcp_server_with_mock_wazuh() -> Result<()> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_wazuh_server = setup_mock_wazuh_server();
let (wazuh_host, wazuh_port) = get_host_port(&mock_wazuh_server);
let mcp_port = find_available_port();
let mut mcp_server = start_mcp_server(&wazuh_host, wazuh_port, mcp_port);
sleep(Duration::from_secs(2)).await;
let mcp_client = McpClient::new(format!("http://localhost:{}", mcp_port));
let health_data = mcp_client.check_health().await?;
assert_eq!(health_data["status"], "ok");
assert_eq!(health_data["service"], "wazuh-mcp-server");
let mcp_data = mcp_client.get_mcp_data().await?;
assert_eq!(mcp_data.len(), 2);
let first_message: &McpMessage = &mcp_data[0];
assert_eq!(first_message.protocol_version, "1.0");
assert_eq!(first_message.source, "Wazuh");
assert_eq!(first_message.event_type, "alert");
let context = &first_message.context;
assert_eq!(context["id"], "12345");
assert_eq!(context["category"], "intrusion_detection");
assert_eq!(context["severity"], "high");
assert_eq!(
context["description"],
"Possible intrusion attempt detected"
);
assert_eq!(context["agent"]["name"], "test-agent");
let data = &context["data"];
assert_eq!(data["source_ip"], "192.168.1.100");
assert_eq!(data["destination_ip"], "10.0.0.1");
assert_eq!(data["port"], 22);
let second_message = &mcp_data[1];
let context = &second_message.context;
assert_eq!(context["id"], "67890");
assert_eq!(context["category"], "malware");
assert_eq!(context["severity"], "critical");
assert_eq!(context["description"], "Malware detected on system");
assert_eq!(context["agent"]["name"], "another-agent");
let data = &context["data"];
assert_eq!(data["file_path"], "/tmp/malicious.exe");
assert_eq!(data["hash"], "abcdef123456");
assert_eq!(data["signature"], "EICAR-Test-File");
mcp_server.kill().expect("Failed to kill MCP server");
Ok(())
}
#[tokio::test]
async fn test_mcp_server_wazuh_api_error() -> Result<()> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_wazuh_server = setup_mock_wazuh_server();
let (wazuh_host, wazuh_port) = get_host_port(&mock_wazuh_server);
mock_wazuh_server.mock(|when, then| {
when.method(GET)
.path_matches(Regex::new(r"/wazuh-alerts-.*_search").unwrap());
then.status(500)
.header("content-type", "application/json")
.json_body(json!({"error": "Wazuh internal error"}));
});
let mcp_port = find_available_port();
let mut mcp_server = start_mcp_server(&wazuh_host, wazuh_port, mcp_port);
sleep(Duration::from_secs(2)).await;
let mcp_client = McpClient::new(format!("http://localhost:{}", mcp_port));
let result = mcp_client.get_mcp_data().await;
assert!(result.is_err());
let err_string = result.unwrap_err().to_string();
assert!(
err_string.contains("500")
|| err_string.contains("502")
|| err_string.contains("API request failed")
);
let health_result = mcp_client.check_health().await;
assert!(health_result.is_ok());
assert_eq!(health_result.unwrap()["status"], "ok");
mcp_server.kill().expect("Failed to kill MCP server");
Ok(())
}
#[tokio::test]
async fn test_mcp_client_error_handling() -> Result<()> {
let _guard = TEST_MUTEX.lock().unwrap();
let server = MockServer::start();
server.mock(|when, then| {
when.method(GET).path("/mcp");
then.status(500)
.header("content-type", "application/json")
.json_body(json!({
"error": "Internal server error"
}));
});
server.mock(|when, then| {
when.method(GET).path("/health");
then.status(503)
.header("content-type", "application/json")
.json_body(json!({
"error": "Service unavailable"
}));
});
let client = McpClient::new(server.url(""));
let result = client.get_mcp_data().await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("500") || err.to_string().contains("MCP request failed"));
let result = client.check_health().await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("503") || err.to_string().contains("Health check failed"));
Ok(())
}
#[tokio::test]
async fn test_mcp_server_missing_alert_data() -> Result<()> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_wazuh_server = setup_mock_wazuh_server();
let (wazuh_host, wazuh_port) = get_host_port(&mock_wazuh_server);
mock_wazuh_server.mock(|when, then| {
when.method(GET)
.path_matches(Regex::new(r"/wazuh-alerts-.*_search").unwrap());
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"hits": {
"hits": [
{
"_source": {
"id": "missing_all",
"timestamp": "invalid-date-format"
}
},
{
"_source": {
"id": "missing_rule_fields",
"timestamp": "2024-05-05T11:00:00.000Z",
"rule": { },
"agent": { "id": "003", "name": "agent-minimal" },
"data": {}
}
},
{
"id": "no_source_nest",
"timestamp": "2024-05-05T12:00:00.000Z",
"rule": {
"level": 2,
"description": "Low severity event",
"groups": ["low_sev"]
},
"agent": { "id": "004" },
"data": { "info": "some data" }
}
]
}
}));
});
let mcp_port = find_available_port();
let mut mcp_server = start_mcp_server(&wazuh_host, wazuh_port, mcp_port);
sleep(Duration::from_secs(2)).await;
let mcp_client = McpClient::new(format!("http://localhost:{}", mcp_port));
let mcp_data = mcp_client.get_mcp_data().await?;
assert_eq!(mcp_data.len(), 3);
let msg1 = &mcp_data[0];
assert_eq!(msg1.context["id"], "missing_all");
assert_eq!(msg1.context["category"], "unknown_category");
assert_eq!(msg1.context["severity"], "unknown_severity");
assert_eq!(msg1.context["description"], "");
assert!(
msg1.context["agent"].is_object() && msg1.context["agent"].as_object().unwrap().is_empty()
);
assert!(
msg1.context["data"].is_object() && msg1.context["data"].as_object().unwrap().is_empty()
);
let ts1 = DateTime::parse_from_rfc3339(&msg1.timestamp)
.unwrap()
.with_timezone(&Utc);
assert!((Utc::now() - ts1).num_seconds() < 5);
let msg2 = &mcp_data[1];
assert_eq!(msg2.context["id"], "missing_rule_fields");
assert_eq!(msg2.context["category"], "unknown_category");
assert_eq!(msg2.context["severity"], "unknown_severity");
assert_eq!(msg2.context["description"], "");
assert_eq!(msg2.context["agent"]["name"], "agent-minimal");
assert!(
msg2.context["data"].is_object() && msg2.context["data"].as_object().unwrap().is_empty()
);
assert_eq!(msg2.timestamp, "2024-05-05T11:00:00Z");
let msg3 = &mcp_data[2];
assert_eq!(msg3.context["id"], "no_source_nest");
assert_eq!(msg3.context["category"], "low_sev");
assert_eq!(msg3.context["severity"], "low");
assert_eq!(msg3.context["description"], "Low severity event");
assert_eq!(msg3.context["agent"]["id"], "004");
assert!(msg3.context["agent"].get("name").is_none());
assert_eq!(msg3.context["data"]["info"], "some data");
assert_eq!(msg3.timestamp, "2024-05-05T12:00:00Z");
mcp_server.kill().expect("Failed to kill MCP server");
Ok(())
}

180
tests/mcp_client.rs
View File

@@ -1,180 +0,0 @@
use anyhow::Result;
use async_trait::async_trait;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpMessage {
pub protocol_version: String,
pub source: String,
pub timestamp: String,
pub event_type: String,
pub context: Value,
pub metadata: Value,
}
#[async_trait]
pub trait McpClientTrait {
async fn get_mcp_data(&self) -> Result<Vec<McpMessage>>;
async fn check_health(&self) -> Result<Value>;
async fn query_mcp_data(&self, filters: Value) -> Result<Vec<McpMessage>>;
}
pub struct McpClient {
client: Client,
base_url: String,
}
impl McpClient {
pub fn new(base_url: String) -> Self {
let client = Client::builder()
.timeout(Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client");
Self { client, base_url }
}
}
#[async_trait]
impl McpClientTrait for McpClient {
async fn get_mcp_data(&self) -> Result<Vec<McpMessage>> {
let url = format!("{}/mcp", self.base_url);
let response = self.client.get(&url).send().await?;
match response.status() {
StatusCode::OK => {
let data = response.json::<Vec<McpMessage>>().await?;
Ok(data)
}
status => {
let error_text = response
.text()
.await
.unwrap_or_else(|_| "Unknown error".to_string());
anyhow::bail!("MCP request failed with status {}: {}", status, error_text)
}
}
}
async fn check_health(&self) -> Result<Value> {
let url = format!("{}/health", self.base_url);
let response = self.client.get(&url).send().await?;
match response.status() {
StatusCode::OK => {
let data = response.json::<Value>().await?;
Ok(data)
}
status => {
let error_text = response
.text()
.await
.unwrap_or_else(|_| "Unknown error".to_string());
anyhow::bail!("Health check failed with status {}: {}", status, error_text)
}
}
}
async fn query_mcp_data(&self, filters: Value) -> Result<Vec<McpMessage>> {
let url = format!("{}/mcp", self.base_url);
let response = self.client.post(&url).json(&filters).send().await?;
match response.status() {
StatusCode::OK => {
let data = response.json::<Vec<McpMessage>>().await?;
Ok(data)
}
status => {
let error_text = response
.text()
.await
.unwrap_or_else(|_| "Unknown error".to_string());
anyhow::bail!("MCP query failed with status {}: {}", status, error_text)
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use httpmock::prelude::*;
use serde_json::json;
use tokio;
#[tokio::test]
async fn test_mcp_client_get_data() {
let server = MockServer::start();
server.mock(|when, then| {
when.method(GET).path("/mcp");
then.status(200)
.header("content-type", "application/json")
.json_body(json!([
{
"protocol_version": "1.0",
"source": "Wazuh",
"timestamp": "2023-05-01T12:00:00Z",
"event_type": "alert",
"context": {
"id": "12345",
"category": "intrusion_detection",
"severity": "high",
"description": "Test alert",
"data": {
"source_ip": "192.168.1.100"
}
},
"metadata": {
"integration": "Wazuh-MCP",
"notes": "Test note"
}
}
]));
});
let client = McpClient::new(server.url(""));
let result = client.get_mcp_data().await.unwrap();
assert_eq!(result.len(), 1);
assert_eq!(result[0].protocol_version, "1.0");
assert_eq!(result[0].source, "Wazuh");
assert_eq!(result[0].event_type, "alert");
let context = &result[0].context;
assert_eq!(context["id"], "12345");
assert_eq!(context["category"], "intrusion_detection");
assert_eq!(context["severity"], "high");
}
#[tokio::test]
async fn test_mcp_client_health_check() {
let server = MockServer::start();
server.mock(|when, then| {
when.method(GET).path("/health");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"status": "ok",
"service": "wazuh-mcp-server",
"timestamp": "2023-05-01T12:00:00Z"
}));
});
let client = McpClient::new(server.url(""));
let result = client.check_health().await.unwrap();
assert_eq!(result["status"], "ok");
assert_eq!(result["service"], "wazuh-mcp-server");
}
}

164
tests/mcp_client_cli.rs
View File

@@ -1,164 +0,0 @@
use anyhow::{anyhow, Result};
use clap::Parser;
use serde_json::Value;
use std::io::{self, Write}; // For stdout().flush() and stdin().read_line()
use mcp_server_wazuh::mcp::client::{McpClient, McpClientTrait};
use serde::Deserialize; // For ParsedRequest
#[derive(Parser, Debug)]
#[clap(
name = "mcp-client-cli",
version = "0.1.0",
about = "Interactive CLI for MCP server. Enter JSON-RPC requests, or 'health'/'quit'."
)]
struct CliArgs {
#[clap(long, help = "Path to the MCP server executable for stdio mode.")]
stdio_exe: Option<String>,
#[clap(
long,
env = "MCP_SERVER_URL",
default_value = "http://localhost:8000",
help = "URL of the MCP server for HTTP mode."
)]
http_url: String,
}
// For parsing raw JSON request strings
#[derive(Deserialize, Debug)]
struct ParsedRequest {
// jsonrpc: String, // Not strictly needed for sending
method: String,
params: Option<Value>,
id: Value, // ID can be string or number
}
#[tokio::main]
async fn main() -> Result<()> {
let cli_args = CliArgs::parse();
let mut client: McpClient;
let is_stdio_mode = cli_args.stdio_exe.is_some();
if let Some(ref exe_path) = cli_args.stdio_exe {
println!("Using stdio mode with executable: {}", exe_path);
client = McpClient::new_stdio(exe_path, None).await?;
println!("Sending 'initialize' request to stdio server...");
match client.initialize().await {
Ok(init_result) => {
println!("Initialization successful:");
println!(" Protocol Version: {}", init_result.protocol_version);
println!(" Server Name: {}", init_result.server_info.name);
println!(" Server Version: {}", init_result.server_info.version);
}
Err(e) => {
eprintln!("Stdio Initialization failed: {}. You may need to send a raw 'initialize' JSON-RPC request or check server logs.", e);
// Allow continuing, user might want to send raw init or other commands.
}
}
} else {
println!("Using HTTP mode with URL: {}", cli_args.http_url);
client = McpClient::new_http(cli_args.http_url.clone());
// No automatic initialize for HTTP mode as per McpClientTrait.
// `initialize` is typically a stdio-specific concept in MCP.
}
println!("\nInteractive MCP Client. Enter a JSON-RPC request, 'health' (HTTP only), or 'quit'.");
println!("Press CTRL-D for EOF to exit.");
let mut input_buffer = String::new();
loop {
input_buffer.clear();
print!("mcp> ");
io::stdout().flush().map_err(|e| anyhow!("Failed to flush stdout: {}", e))?;
match io::stdin().read_line(&mut input_buffer) {
Ok(0) => { // EOF (Ctrl-D)
println!("\nEOF detected. Exiting.");
break;
}
Ok(_) => {
let line = input_buffer.trim();
if line.is_empty() {
continue;
}
if line.eq_ignore_ascii_case("quit") {
println!("Exiting.");
break;
}
if line.eq_ignore_ascii_case("health") {
if is_stdio_mode {
println!("'health' command is intended for HTTP mode. For stdio, you would need to send a specific JSON-RPC request if the server supports a health method via stdio.");
} else {
println!("Checking server health (HTTP GET to /health)...");
let health_url = format!("{}/health", cli_args.http_url); // Use the parsed http_url
match reqwest::get(&health_url).await {
Ok(response) => {
let status = response.status();
let response_text = response.text().await.unwrap_or_else(|_| "Failed to read response body".to_string());
if status.is_success() {
match serde_json::from_str::<Value>(&response_text) {
Ok(json_val) => println!("Health response ({}):\n{}", status, serde_json::to_string_pretty(&json_val).unwrap_or_else(|_| response_text.clone())),
Err(_) => println!("Health response ({}):\n{}", status, response_text),
}
} else {
eprintln!("Health check failed with status: {}", status);
eprintln!("Response: {}", response_text);
}
}
Err(e) => eprintln!("Health check request failed: {}", e),
}
}
continue;
}
// Assume it's a JSON-RPC request
println!("Attempting to send as JSON-RPC: {}", line);
match serde_json::from_str::<ParsedRequest>(line) {
Ok(parsed_req) => {
match client
.send_json_rpc_request(
&parsed_req.method,
parsed_req.params.clone(),
parsed_req.id.clone(),
)
.await
{
Ok(response_value) => {
println!(
"Server Response: {}",
serde_json::to_string_pretty(&response_value).unwrap_or_else(
|e_pretty| format!("Failed to pretty-print response ({}): {:?}", e_pretty, response_value)
)
);
}
Err(e) => {
eprintln!("Error processing JSON-RPC request '{}': {}", line, e);
}
}
}
Err(e) => {
eprintln!("Failed to parse input as a JSON-RPC request: {}. Input: '{}'", e, line);
eprintln!("Please enter a valid JSON-RPC request string, 'health', or 'quit'.");
}
}
}
Err(e) => {
eprintln!("Error reading input: {}. Exiting.", e);
break;
}
}
}
if is_stdio_mode {
println!("Sending 'shutdown' request to stdio server...");
match client.shutdown().await {
Ok(_) => println!("Shutdown command acknowledged by server."),
Err(e) => eprintln!("Error during shutdown: {}. Server might have already exited or closed the connection.", e),
}
}
Ok(())
}

361
tests/mcp_stdio_test.rs Normal file
View File

@@ -0,0 +1,361 @@
//! Tests for MCP protocol communication via stdio
//!
//! These tests verify the basic MCP protocol implementation without
//! requiring a Wazuh connection.
use std::process::{Command, Stdio};
use std::io::{BufRead, BufReader, Write};
use std::time::Duration;
use tokio::time::sleep;
use serde_json::{json, Value};
struct McpStdioClient {
child: std::process::Child,
stdin: std::process::ChildStdin,
stdout: BufReader<std::process::ChildStdout>,
}
impl McpStdioClient {
fn start() -> Result<Self, Box<dyn std::error::Error>> {
let mut child = Command::new("cargo")
.args(["run", "--bin", "mcp-server-wazuh"])
.env("WAZUH_HOST", "nonexistent.example.com") // Use non-existent host
.env("WAZUH_PORT", "9999")
.env("WAZUH_USER", "test")
.env("WAZUH_PASS", "test")
.env("VERIFY_SSL", "false")
.env("RUST_LOG", "error") // Minimize logging noise
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit()) // Changed from Stdio::null() to inherit stderr
.spawn()?;
let stdin = child.stdin.take().unwrap();
let stdout = BufReader::new(child.stdout.take().unwrap());
Ok(McpStdioClient {
child,
stdin,
stdout,
})
}
fn send_message(&mut self, message: &Value) -> Result<(), Box<dyn std::error::Error>> {
let message_str = serde_json::to_string(message)?;
writeln!(self.stdin, "{}", message_str)?;
self.stdin.flush()?;
Ok(())
}
fn read_response(&mut self) -> Result<Value, Box<dyn std::error::Error>> {
let mut line = String::new();
self.stdout.read_line(&mut line)?;
let response: Value = serde_json::from_str(&line.trim())?;
Ok(response)
}
fn send_and_receive(&mut self, message: &Value) -> Result<Value, Box<dyn std::error::Error>> {
self.send_message(message)?;
self.read_response()
}
}
impl Drop for McpStdioClient {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
#[tokio::test]
async fn test_mcp_protocol_initialization() -> Result<(), Box<dyn std::error::Error>> {
let mut client = McpStdioClient::start()?;
// Give the server time to start
sleep(Duration::from_millis(500)).await;
// Test initialize request
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {
"name": "test-client",
"version": "1.0.0"
}
}
});
let response = client.send_and_receive(&init_request)?;
// Verify JSON-RPC 2.0 compliance
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 1);
assert!(response["result"].is_object());
assert!(response["error"].is_null());
// Verify MCP initialize response structure
let result = &response["result"];
assert_eq!(result["protocolVersion"], "2024-11-05");
assert!(result["capabilities"].is_object());
assert!(result["serverInfo"].is_object());
// Verify server info
let server_info = &result["serverInfo"];
assert!(server_info["name"].is_string());
assert!(server_info["version"].is_string());
Ok(())
}
#[tokio::test]
async fn test_mcp_tools_list_without_wazuh() -> Result<(), Box<dyn std::error::Error>> {
let mut client = McpStdioClient::start()?;
sleep(Duration::from_millis(500)).await;
// Initialize first
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
client.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
client.send_message(&initialized)?;
// Request tools list
let tools_request = json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/list",
"params": {}
});
let response = client.send_and_receive(&tools_request)?;
// Verify response structure
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 2);
assert!(response["result"].is_object());
let result = &response["result"];
assert!(result["tools"].is_array());
let tools = result["tools"].as_array().unwrap();
assert!(!tools.is_empty());
// Verify tool structure
for tool in tools {
assert!(tool["name"].is_string());
assert!(tool["description"].is_string());
assert!(tool["inputSchema"].is_object());
}
Ok(())
}
#[tokio::test]
async fn test_invalid_json_rpc_request() -> Result<(), Box<dyn std::error::Error>> {
let mut client = McpStdioClient::start()?;
sleep(Duration::from_millis(500)).await;
// 1. Initialize the connection first
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test-client", "version": "1.0.0"}
}
});
let _init_response = client.send_and_receive(&init_request)?; // Read and ignore/assert init response
// assert!(_init_response["result"].is_object()); // Optional: assert successful init
// 2. Send initialized notification
let initialized_notification = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
client.send_message(&initialized_notification)?;
// 3. Send the invalid JSON-RPC request (missing required fields)
let invalid_request = json!({
// "jsonrpc": "2.0", // Missing jsonrpc field to make it invalid
"id": 2, // Use a new ID
"method": "some_method_that_might_not_exist"
});
client.send_message(&invalid_request)?;
// 4. The server currently closes the connection upon such an invalid request (see logs:
// `ERROR rmcp::transport::io ... serde error ...` followed by `input stream terminated`).
// Therefore, subsequent requests should fail. This test verifies this behavior.
// Ideally, the server might send a JSON-RPC error and keep the connection open,
// but that would require changes to the server's error handling logic.
// 5. Attempt to send a subsequent valid request.
let list_tools_request = json!({
"jsonrpc": "2.0",
"id": 3, // New ID
"method": "tools/list",
"params": {}
});
let result = client.send_and_receive(&list_tools_request);
// Assert that the operation failed, indicating the connection was likely closed.
assert!(result.is_err(), "Server should have closed the connection after the invalid request, leading to an error here.");
// Optionally, check the error type more specifically if needed, e.g., for EOF.
if let Err(e) = result {
let error_message = e.to_string().to_lowercase();
assert!(
error_message.contains("eof") || error_message.contains("broken pipe") || error_message.contains("connection reset"),
"Expected EOF, broken pipe, or connection reset error, but got: {}", e
);
}
Ok(())
}
#[tokio::test]
async fn test_unsupported_method() -> Result<(), Box<dyn std::error::Error>> {
let mut client = McpStdioClient::start()?;
sleep(Duration::from_millis(500)).await;
// Initialize first
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
client.send_and_receive(&init_request)?;
// Send initialized notification
let initialized_notification = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
client.send_message(&initialized_notification)?;
let unsupported_request = json!({
"jsonrpc": "2.0",
"id": 2,
"method": "unsupported/method"
// Omitting "params": {} as it might be causing deserialization issues
// in rmcp for unknown methods. The JSON-RPC spec allows params to be omitted.
});
// Send the unsupported request. We don't expect a valid JSON-RPC response.
// Instead, the server is likely to close the connection due to deserialization issues
// in rmcp when encountering an unknown method, as it cannot match it to a known JsonRpcMessage variant.
client.send_message(&unsupported_request)?;
// Attempt to send a subsequent valid request to confirm the connection was dropped.
let list_tools_request = json!({
"jsonrpc": "2.0",
"id": 3, // Use a new ID
"method": "tools/list",
"params": {}
});
let result = client.send_and_receive(&list_tools_request);
// Assert that the operation failed, indicating the connection was likely closed.
assert!(result.is_err(), "Server should have closed the connection after the unsupported method request, leading to an error here.");
// Optionally, check the error type more specifically if needed, e.g., for EOF.
if let Err(e) = result {
let error_message = e.to_string().to_lowercase();
assert!(
error_message.contains("eof") || error_message.contains("broken pipe") || error_message.contains("connection reset"),
"Expected EOF, broken pipe, or connection reset error, but got: {}", e
);
}
Ok(())
}
#[tokio::test]
async fn test_concurrent_requests() -> Result<(), Box<dyn std::error::Error>> {
let mut client = McpStdioClient::start()?;
sleep(Duration::from_millis(500)).await;
// Initialize
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
client.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
client.send_message(&initialized)?;
// Send multiple requests with different IDs
let request1 = json!({
"jsonrpc": "2.0",
"id": 10,
"method": "tools/list",
"params": {}
});
let request2 = json!({
"jsonrpc": "2.0",
"id": 20,
"method": "tools/list",
"params": {}
});
// Send both requests
client.send_message(&request1)?;
client.send_message(&request2)?;
// Read both responses
let response1 = client.read_response()?;
let response2 = client.read_response()?;
// Responses should maintain request IDs (though order might vary)
let ids: Vec<i64> = vec![
response1["id"].as_i64().unwrap(),
response2["id"].as_i64().unwrap(),
];
assert!(ids.contains(&10));
assert!(ids.contains(&20));
Ok(())
}

340
tests/mock_wazuh_server.rs Normal file
View File

@@ -0,0 +1,340 @@
//! Mock Wazuh API server for testing
//!
//! This module provides a configurable mock server that simulates the Wazuh Indexer API
//! for testing purposes. It supports various response scenarios including success,
//! empty results, and error conditions.
use httpmock::prelude::*;
use serde_json::json;
pub struct MockWazuhServer {
server: MockServer,
}
impl MockWazuhServer {
pub fn new() -> Self {
let server = MockServer::start();
// Setup default authentication endpoint
server.mock(|when, then| {
when.method(POST).path("/security/user/authenticate");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"jwt": "mock.jwt.token.eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9"
}));
});
server.mock(|when, then| {
when.method(POST)
.path_matches(Regex::new(r"/wazuh-alerts.*/_search").unwrap());
then.status(200)
.header("content-type", "application/json")
.json_body(Self::sample_alerts_response());
});
Self { server }
}
pub fn with_empty_alerts() -> Self {
let server = MockServer::start();
server.mock(|when, then| {
when.method(POST).path("/security/user/authenticate");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"jwt": "mock.jwt.token"
}));
});
server.mock(|when, then| {
when.method(POST)
.path_matches(Regex::new(r"/wazuh-alerts.*/_search").unwrap());
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"hits": {
"hits": []
}
}));
});
Self { server }
}
pub fn with_auth_error() -> Self {
let server = MockServer::start();
server.mock(|when, then| {
when.method(POST).path("/security/user/authenticate");
then.status(401)
.header("content-type", "application/json")
.json_body(json!({
"error": "Invalid credentials"
}));
});
Self { server }
}
pub fn with_alerts_error() -> Self {
let server = MockServer::start();
server.mock(|when, then| {
when.method(POST).path("/security/user/authenticate");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"jwt": "mock.jwt.token"
}));
});
server.mock(|when, then| {
when.method(POST)
.path_matches(Regex::new(r"/wazuh-alerts.*/_search").unwrap());
then.status(500)
.header("content-type", "application/json")
.json_body(json!({
"error": "Internal server error"
}));
});
Self { server }
}
pub fn with_malformed_alerts() -> Self {
let server = MockServer::start();
server.mock(|when, then| {
when.method(POST).path("/security/user/authenticate");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"jwt": "mock.jwt.token"
}));
});
server.mock(|when, then| {
when.method(POST)
.path_matches(Regex::new(r"/wazuh-alerts.*/_search").unwrap());
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"hits": {
"hits": [
{
"_source": {
"id": "missing_fields",
"timestamp": "invalid-date-format"
// Missing rule, agent, etc.
}
},
{
"_source": {
"id": "partial_data",
"timestamp": "2024-01-15T10:30:45.123Z",
"rule": {
"level": 5
// Missing description
},
"agent": {
"id": "001"
// Missing name
}
}
}
]
}
}));
});
Self { server }
}
pub fn url(&self) -> String {
self.server.url("")
}
pub fn host(&self) -> String {
let url = self.url();
let parts: Vec<&str> = url.trim_start_matches("http://").split(':').collect();
parts[0].to_string()
}
pub fn port(&self) -> u16 {
let url = self.url();
let parts: Vec<&str> = url.trim_start_matches("http://").split(':').collect();
parts[1].parse().unwrap()
}
fn sample_alerts_response() -> serde_json::Value {
json!({
"hits": {
"hits": [
{
"_source": {
"id": "1747091815.1212763",
"timestamp": "2024-01-15T10:30:45.123Z",
"rule": {
"level": 7,
"description": "Attached USB Storage",
"groups": ["usb", "pci_dss"]
},
"agent": {
"id": "001",
"name": "web-server-01"
},
"data": {
"device": "/dev/sdb1",
"mount_point": "/media/usb"
}
}
},
{
"_source": {
"id": "1747066333.1207112",
"timestamp": "2024-01-15T10:25:12.456Z",
"rule": {
"level": 5,
"description": "New dpkg (Debian Package) installed.",
"groups": ["package_management", "debian"]
},
"agent": {
"id": "002",
"name": "database-server"
},
"data": {
"package": "nginx",
"version": "1.18.0-6ubuntu14.4"
}
}
},
{
"_source": {
"id": "1747055444.1205998",
"timestamp": "2024-01-15T10:20:33.789Z",
"rule": {
"level": 12,
"description": "Multiple authentication failures",
"groups": ["authentication_failed", "pci_dss"]
},
"agent": {
"id": "003",
"name": "ssh-gateway"
},
"data": {
"source_ip": "192.168.1.100",
"user": "admin",
"attempts": 5
}
}
}
]
}
})
}
}
impl Default for MockWazuhServer {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_mock_server_creation() {
let mock_server = MockWazuhServer::new();
assert!(!mock_server.url().is_empty());
assert!(!mock_server.host().is_empty());
assert!(mock_server.port() > 0);
}
#[tokio::test]
async fn test_mock_server_auth_endpoint() {
let mock_server = MockWazuhServer::new();
let client = reqwest::Client::new();
let response = client
.post(&format!("{}/security/user/authenticate", mock_server.url()))
.json(&json!({"username": "admin", "password": "admin"}))
.send()
.await
.unwrap();
assert_eq!(response.status(), 200);
let body: serde_json::Value = response.json().await.unwrap();
assert!(body.get("jwt").is_some());
}
#[tokio::test]
async fn test_mock_server_alerts_endpoint() {
let mock_server = MockWazuhServer::new();
let client = reqwest::Client::new();
let response = client
.post(&format!("{}/wazuh-alerts*/_search", mock_server.url()))
.json(&json!({"query": {"match_all": {}}}))
.send()
.await
.unwrap();
assert_eq!(response.status(), 200);
let body: serde_json::Value = response.json().await.unwrap();
assert!(body.get("hits").is_some());
let hits = body["hits"]["hits"].as_array().unwrap();
assert!(!hits.is_empty());
}
#[tokio::test]
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()))
.json(&json!({"query": {"match_all": {}}}))
.send()
.await
.unwrap();
assert_eq!(response.status(), 200);
let body: serde_json::Value = response.json().await.unwrap();
let hits = body["hits"]["hits"].as_array().unwrap();
assert!(hits.is_empty());
}
#[tokio::test]
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()))
.json(&json!({"username": "admin", "password": "wrong"}))
.send()
.await
.unwrap();
assert_eq!(response.status(), 401);
}
#[tokio::test]
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()))
.json(&json!({"query": {"match_all": {}}}))
.send()
.await
.unwrap();
assert_eq!(response.status(), 500);
}
}

546
tests/rmcp_integration_test.rs Normal file
View File

@@ -0,0 +1,546 @@
//! 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.
use std::process::{Child, Command, Stdio};
use std::io::{BufRead, BufReader, Write};
use std::time::Duration;
use tokio::time::sleep;
use serde_json::{json, Value};
use once_cell::sync::Lazy;
use std::sync::Mutex;
mod mock_wazuh_server;
use mock_wazuh_server::MockWazuhServer;
static TEST_MUTEX: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
struct McpServerProcess {
child: Child,
stdin: std::process::ChildStdin,
stdout: BufReader<std::process::ChildStdout>,
}
impl McpServerProcess {
fn start_with_mock_wazuh(mock_server: &MockWazuhServer) -> Result<Self, Box<dyn std::error::Error>> {
let mut child = Command::new("cargo")
.args(["run", "--bin", "mcp-server-wazuh"])
.env("WAZUH_HOST", mock_server.host())
.env("WAZUH_PORT", mock_server.port().to_string())
.env("WAZUH_USER", "admin")
.env("WAZUH_PASS", "admin")
.env("VERIFY_SSL", "false")
.env("WAZUH_TEST_PROTOCOL", "http")
.env("RUST_LOG", "warn") // Reduce noise in tests
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit()) // Inherit stderr to see server logs
.spawn()?;
let stdin = child.stdin.take().unwrap();
let stdout = BufReader::new(child.stdout.take().unwrap());
Ok(McpServerProcess {
child,
stdin,
stdout,
})
}
fn send_message(&mut self, message: &Value) -> Result<(), Box<dyn std::error::Error>> {
let message_str = serde_json::to_string(message)?;
writeln!(self.stdin, "{}", message_str)?;
self.stdin.flush()?;
Ok(())
}
fn read_response(&mut self) -> Result<Value, Box<dyn std::error::Error>> {
let mut line = String::new();
self.stdout.read_line(&mut line)?;
let response: Value = serde_json::from_str(line.trim())?;
Ok(response)
}
fn send_and_receive(&mut self, message: &Value) -> Result<Value, Box<dyn std::error::Error>> {
self.send_message(message)?;
self.read_response()
}
}
impl Drop for McpServerProcess {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
#[tokio::test]
async fn test_mcp_server_initialization() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::new();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
// Give the server time to start
sleep(Duration::from_millis(500)).await;
// Send initialize request
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {
"name": "test-client",
"version": "1.0.0"
}
}
});
let response = mcp_server.send_and_receive(&init_request)?;
// Verify response structure
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 1);
assert!(response["result"].is_object());
let result = &response["result"];
assert_eq!(result["protocolVersion"], "2024-11-05");
assert!(result["capabilities"].is_object());
assert!(result["serverInfo"].is_object());
assert!(result["instructions"].is_string());
Ok(())
}
#[tokio::test]
async fn test_tools_list() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::new();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
sleep(Duration::from_millis(500)).await;
// Initialize first
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
mcp_server.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
mcp_server.send_message(&initialized)?;
// Request tools list
let tools_request = json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/list",
"params": {}
});
let response = mcp_server.send_and_receive(&tools_request)?;
// Verify response
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 2);
assert!(response["result"]["tools"].is_array());
let tools = response["result"]["tools"].as_array().unwrap();
assert!(!tools.is_empty());
// Check for our Wazuh alert summary tool
let alert_tool = tools.iter()
.find(|tool| tool["name"] == "get_wazuh_alert_summary")
.expect("get_wazuh_alert_summary tool should be present");
assert!(alert_tool["description"].is_string());
assert!(alert_tool["inputSchema"].is_object());
// Verify input schema structure
let input_schema = &alert_tool["inputSchema"];
assert_eq!(input_schema["type"], "object");
assert!(input_schema["properties"].is_object());
assert!(input_schema["properties"]["limit"].is_object());
Ok(())
}
#[tokio::test]
async fn test_get_wazuh_alert_summary_success() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::new();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
sleep(Duration::from_millis(500)).await;
// Initialize
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
mcp_server.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
mcp_server.send_message(&initialized)?;
// Call the tool
let tool_call = json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "get_wazuh_alert_summary",
"arguments": {
"limit": 2
}
}
});
let response = mcp_server.send_and_receive(&tool_call)?;
// Verify response structure
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 3);
assert!(response["result"].is_object());
let result = &response["result"];
assert!(result["content"].is_array());
assert_eq!(result["isError"], false);
let content = result["content"].as_array().unwrap();
assert!(!content.is_empty());
// Verify content format
for item in content {
assert_eq!(item["type"], "text");
assert!(item["text"].is_string());
let text = item["text"].as_str().unwrap();
assert!(text.contains("Alert ID:"));
assert!(text.contains("Time:"));
assert!(text.contains("Agent:"));
assert!(text.contains("Level:"));
assert!(text.contains("Description:"));
}
Ok(())
}
#[tokio::test]
async fn test_get_wazuh_alert_summary_empty_results() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::with_empty_alerts();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
sleep(Duration::from_millis(500)).await;
// Initialize
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
mcp_server.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
mcp_server.send_message(&initialized)?;
// Call the tool
let tool_call = json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "get_wazuh_alert_summary",
"arguments": {}
}
});
let response = mcp_server.send_and_receive(&tool_call)?;
// Verify response
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 3);
let result = &response["result"];
assert!(result["content"].is_array());
assert_eq!(result["isError"], false);
let content = result["content"].as_array().unwrap();
assert_eq!(content.len(), 1);
assert_eq!(content[0]["type"], "text");
assert_eq!(content[0]["text"], "No Wazuh alerts found.");
Ok(())
}
#[tokio::test]
async fn test_get_wazuh_alert_summary_api_error() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::with_alerts_error();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
sleep(Duration::from_millis(500)).await;
// Initialize
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
mcp_server.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
mcp_server.send_message(&initialized)?;
// Call the tool
let tool_call = json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "get_wazuh_alert_summary",
"arguments": {
"limit": 5
}
}
});
let response = mcp_server.send_and_receive(&tool_call)?;
// Verify error response
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 3);
let result = &response["result"];
assert!(result["content"].is_array());
assert_eq!(result["isError"], true);
let content = result["content"].as_array().unwrap();
assert_eq!(content.len(), 1);
assert_eq!(content[0]["type"], "text");
let error_text = content[0]["text"].as_str().unwrap();
assert!(error_text.contains("Error retrieving alerts from Wazuh"));
Ok(())
}
#[tokio::test]
async fn test_invalid_tool_call() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::new();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
sleep(Duration::from_millis(500)).await;
// Initialize
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
mcp_server.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
mcp_server.send_message(&initialized)?;
// Call non-existent tool
let tool_call = json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "non_existent_tool",
"arguments": {}
}
});
let response = mcp_server.send_and_receive(&tool_call)?;
// Should get an error response
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 3);
assert!(response["error"].is_object());
Ok(())
}
#[tokio::test]
async fn test_parameter_validation() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::new();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
sleep(Duration::from_millis(500)).await;
// Initialize
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
mcp_server.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
mcp_server.send_message(&initialized)?;
// Test with invalid parameter type (string instead of number)
let tool_call = json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "get_wazuh_alert_summary",
"arguments": {
"limit": "invalid"
}
}
});
let response = mcp_server.send_and_receive(&tool_call)?;
// Should get an error response for invalid parameters
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 3);
// The response might be an error or a successful response with error content
// depending on how rmcp handles parameter validation
assert!(response["error"].is_object() ||
(response["result"]["isError"] == true));
Ok(())
}
#[tokio::test]
async fn test_malformed_alert_data_handling() -> Result<(), Box<dyn std::error::Error>> {
let _guard = TEST_MUTEX.lock().unwrap();
let mock_server = MockWazuhServer::with_malformed_alerts();
let mut mcp_server = McpServerProcess::start_with_mock_wazuh(&mock_server)?;
sleep(Duration::from_millis(500)).await;
// Initialize
let init_request = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
}
});
mcp_server.send_and_receive(&init_request)?;
// Send initialized notification
let initialized = json!({
"jsonrpc": "2.0",
"method": "notifications/initialized"
});
mcp_server.send_message(&initialized)?;
// Call the tool
let tool_call = json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "get_wazuh_alert_summary",
"arguments": {
"limit": 5
}
}
});
let response = mcp_server.send_and_receive(&tool_call)?;
// Should handle malformed data gracefully
assert_eq!(response["jsonrpc"], "2.0");
assert_eq!(response["id"], 3);
let result = &response["result"];
assert!(result["content"].is_array());
// Should not error out, but handle missing fields gracefully
assert_eq!(result["isError"], false);
let content = result["content"].as_array().unwrap();
assert!(!content.is_empty());
// Verify that missing fields are handled with defaults
for item in content {
assert_eq!(item["type"], "text");
let text = item["text"].as_str().unwrap();
// Should contain default values for missing fields
assert!(text.contains("Alert ID:"));
assert!(text.contains("Unknown") || text.contains("missing_fields") || text.contains("partial_data"));
}
Ok(())
}

127
tests/run_tests.sh
View File

@@ -1,47 +1,100 @@
#!/bin/bash
echo "Running all tests..."
cargo test
# Test script for Wazuh MCP Server (rmcp-based)
# This script runs various tests to ensure the server is working correctly
echo "Building MCP client CLI..."
cargo build --bin mcp_client_cli
set -e
echo "Building main server binary for stdio CLI tests..."
# Build the server executable that mcp_client_cli will run
cargo build --bin mcp-server-wazuh # Output: target/debug/mcp-server-wazuh
echo "Starting Wazuh MCP Server tests (rmcp-based)..."
echo "Testing MCP client CLI in stdio mode..."
# Set test environment variables
export RUST_LOG=info
echo "Executing: ./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh initialize"
./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh initialize
if [ $? -ne 0 ]; then
echo "CLI 'initialize' command failed!"
exit 1
fi
echo "Environment variables set:"
echo " RUST_LOG: $RUST_LOG"
echo "Executing: ./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh provideContext"
./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh provideContext
if [ $? -ne 0 ]; then
echo "CLI 'provideContext' command failed!"
exit 1
fi
# Function to cleanup background processes
cleanup() {
echo "Cleaning up..."
if [ ! -z "$SERVER_PID" ]; then
kill $SERVER_PID 2>/dev/null || true
wait $SERVER_PID 2>/dev/null || true
fi
}
# Example of provideContext with empty JSON params (optional to uncomment and test)
# echo "Executing: ./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh provideContext '{}'"
# ./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh provideContext '{}'
# if [ $? -ne 0 ]; then
# echo "CLI 'provideContext {}' command failed!"
# exit 1
# fi
# Set trap to cleanup on exit
trap cleanup EXIT
echo "Executing: ./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh shutdown"
./target/debug/mcp_client_cli --stdio-exe ./target/debug/mcp-server-wazuh shutdown
if [ $? -ne 0 ]; then
# Shutdown might return an error if the server closes the pipe before the client fully processes the response,
# but the primary goal is that the server process is terminated.
# For this script, we'll be lenient on shutdown's exit code for now,
# as long as initialize and provideContext worked.
echo "CLI 'shutdown' command executed (non-zero exit code is sometimes expected if server closes pipe quickly)."
fi
echo ""
echo "=== Running Unit Tests ==="
cargo test --lib
echo "MCP client CLI stdio tests completed."
echo ""
echo "=== Running MCP Protocol Tests ==="
cargo test --test mcp_stdio_test
echo ""
echo "=== Running Integration Tests with Mock Wazuh ==="
cargo test --test rmcp_integration_test
echo ""
echo "=== Manual MCP Server Testing ==="
# Test the server with a simple MCP interaction
echo "Testing MCP server initialization..."
# Create a temporary test script
cat > /tmp/test_mcp_server.sh << 'INNER_EOF'
#!/bin/bash
# Start the server in background
WAZUH_HOST=mock.example.com RUST_LOG=error cargo run --bin mcp-server-wazuh &
SERVER_PID=$!
# Give server time to start
sleep 1
# Test MCP initialization
echo '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}' | timeout 5s nc -l 0 2>/dev/null || {
# If nc doesn't work, try a different approach
echo "Testing server startup..."
sleep 2
}
# Clean up
kill $SERVER_PID 2>/dev/null || true
wait $SERVER_PID 2>/dev/null || true
echo "Manual test completed"
INNER_EOF
chmod +x /tmp/test_mcp_server.sh
/tmp/test_mcp_server.sh
rm /tmp/test_mcp_server.sh
echo ""
echo "=== Testing Server Binary ==="
echo "Verifying server binary can start and show help..."
# Test that the binary can start and show help
timeout 5s cargo run --bin mcp-server-wazuh -- --help || echo "Help command test completed"
echo ""
echo "=== All Tests Complete ==="
echo ""
echo "Test Summary:"
echo "✓ Unit tests for library components"
echo "✓ Wazuh client tests with mock HTTP server"
echo "✓ MCP protocol tests via stdio"
echo "✓ Integration tests with mock Wazuh API"
echo "✓ Server binary startup test"
echo ""
echo "To test manually with a real Wazuh instance:"
echo " export WAZUH_HOST=your-wazuh-host"
echo " export WAZUH_PORT=9200"
echo " export WAZUH_USER=admin"
echo " export WAZUH_PASS=your-password"
echo " cargo run --bin mcp-server-wazuh"
echo ""
echo "Then send MCP commands via stdin, for example:"
echo ' echo '"'"'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}'"'"' | cargo run --bin mcp-server-wazuh'