NeoAnd/mcp-server-wazuh
1
0
Fork 0
mirror of https://github.com/gbrigandi/mcp-server-wazuh.git synced 2025-07-13 15:14:48 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

Wazuh MCP server successfully tested with Claude

Browse Source
This commit is contained in:
Gianluca Brigandi 2025-05-12 19:29:09 -07:00
parent 65e2c55ccb
commit 13494cf101
14 changed files with 1775 additions and 298 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.toml
View File

@ -20,6 +20,7 @@ jsonwebtoken = "9.3"
tower-http = { version = "0.6", features = ["trace"] }
async-trait = "0.1"
anyhow = "1.0"
clap = { version = "4.5", features = ["derive", "env"] }
[dev-dependencies]
mockito = "1.7"

137
src/http_service.rs Normal file
View File

@ -0,0 +1,137 @@
use axum::{
extract::State,
http::StatusCode,
response::{IntoResponse, Response},
routing::{get, post},
Json, Router,
};
use serde_json::{json, Value};
use std::sync::Arc;
use tracing::{debug, error, info};
use crate::logging_utils::{log_mcp_request, log_mcp_response};
use crate::AppState; // Added
pub fn create_http_router(app_state: Arc<AppState>) -> Router {
Router::new()
.route("/health", get(health_check))
.route("/mcp", get(get_mcp_data))
.route("/mcp", post(post_mcp_data))
.with_state(app_state)
}
async fn health_check() -> impl IntoResponse {
Json(json!({
"status": "ok",
"service": "wazuh-mcp-server",
"timestamp": chrono::Utc::now().to_rfc3339()
}))
}
async fn get_mcp_data(
State(app_state): State<Arc<AppState>>,
) -> Result<Json<Vec<Value>>, ApiError> {
info!("Handling GET /mcp request");
let wazuh_client = app_state.wazuh_client.lock().await;
match wazuh_client.get_alerts().await {
Ok(alerts) => {
// Transform Wazuh alerts to MCP messages
let mcp_messages = alerts
.iter()
.map(|alert| {
json!({
"protocol_version": "1.0",
"source": "Wazuh",
"timestamp": chrono::Utc::now().to_rfc3339(),
"event_type": "alert",
"context": alert,
"metadata": {
"integration": "Wazuh-MCP"
}
})
})
.collect::<Vec<_>>();
Ok(Json(mcp_messages))
}
Err(e) => {
error!("Error getting alerts from Wazuh: {}", e);
Err(ApiError::InternalServerError(format!(
"Failed to get alerts from Wazuh: {}",
e
)))
}
}
}
async fn post_mcp_data(
State(app_state): State<Arc<AppState>>,
Json(payload): Json<Value>,
) -> Result<Json<Vec<Value>>, ApiError> {
info!("Handling POST /mcp request with payload");
debug!("Payload: {:?}", payload);
// Log the incoming payload
let request_str = serde_json::to_string(&payload).unwrap_or_else(|e| {
error!(
"Failed to serialize POST request payload for logging: {}",
e
);
format!(
"{{\"error\":\"Failed to serialize request payload: {}\"}}",
e
)
});
log_mcp_request(&request_str);
let result = get_mcp_data(State(app_state)).await;
// Log the response
let response_str = match &result {
Ok(json_response) => serde_json::to_string(&json_response.0).unwrap_or_else(|e| {
error!("Failed to serialize POST response for logging: {}", e);
format!("{{\"error\":\"Failed to serialize response: {}\"}}", e)
}),
Err(api_error) => {
let error_json_surrogate = json!({
"error": format!("{:?}", api_error) // Or a more structured error
});
serde_json::to_string(&error_json_surrogate).unwrap_or_else(|e| {
error!("Failed to serialize POST error response for logging: {}", e);
format!(
"{{\"error\":\"Failed to serialize error response: {}\"}}",
e
)
})
}
};
log_mcp_response(&response_str);
result
}
// API Error handling
#[derive(Debug)]
pub enum ApiError {
BadRequest(String),
NotFound(String),
InternalServerError(String),
}
impl IntoResponse for ApiError {
fn into_response(self) -> Response {
let (status, error_message) = match self {
ApiError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg),
ApiError::NotFound(msg) => (StatusCode::NOT_FOUND, msg),
ApiError::InternalServerError(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg),
};
let body = Json(json!({
"error": error_message
}));
(status, body).into_response()
}
}

13
src/lib.rs Normal file
View File

@ -0,0 +1,13 @@
use crate::wazuh::client::WazuhIndexerClient;
use tokio::sync::Mutex;
pub mod http_service;
pub mod logging_utils;
pub mod mcp;
pub mod stdio_service;
pub mod wazuh;
#[derive(Debug)]
pub struct AppState {
pub wazuh_client: Mutex<WazuhIndexerClient>,
}

27
src/logging_utils.rs Normal file
View File

@ -0,0 +1,27 @@
use std::fs::OpenOptions;
use std::io::Write;
use tracing::error;
const REQUEST_LOG_FILE: &str = "mcp_requests.log";
const RESPONSE_LOG_FILE: &str = "mcp_responses.log";
fn log_to_file(filename: &str, message: &str) {
match OpenOptions::new().create(true).append(true).open(filename) {
Ok(mut file) => {
if let Err(e) = writeln!(file, "{}", message) {
error!("Failed to write to {}: {}", filename, e);
}
}
Err(e) => {
error!("Failed to open {} for appending: {}", filename, e);
}
}
}
pub fn log_mcp_request(request_str: &str) {
log_to_file(REQUEST_LOG_FILE, request_str);
}
pub fn log_mcp_response(response_str: &str) {
log_to_file(RESPONSE_LOG_FILE, response_str);
}

237
src/main.rs
View File

@ -1,136 +1,187 @@
mod wazuh;
mod mcp;
use std::env;
use std::net::SocketAddr;
use std::sync::Arc;
use axum::{
routing::get,
Router,
extract::State,
response::Json,
http::StatusCode,
};
use dotenv::dotenv;
use serde_json::{json, Value};
use tokio::sync::Mutex;
use tracing::{info, error};
use std::backtrace::Backtrace;
use tokio::sync::{oneshot, Mutex};
use tracing::{debug, error, info, Level};
use tracing_subscriber::EnvFilter;
use wazuh::client::WazuhApiClient;
use mcp::transform::transform_to_mcp;
// Application state shared across handlers
struct AppState {
wazuh_client: Mutex<WazuhApiClient>,
}
// Use components from the library crate
use mcp_server_wazuh::http_service::create_http_router;
use mcp_server_wazuh::stdio_service::run_stdio_service;
use mcp_server_wazuh::wazuh::client::WazuhIndexerClient;
use mcp_server_wazuh::AppState;
#[tokio::main]
async fn main() {
dotenv().ok();
tracing_subscriber::fmt::init();
// Set a custom panic hook to ensure panics are logged
std::panic::set_hook(Box::new(|panic_info| {
// Using eprintln directly as tracing might not be available or working during a panic
eprintln!(
"\n================================================================================\n"
);
eprintln!("PANIC OCCURRED IN MCP SERVER");
eprintln!(
"\n--------------------------------------------------------------------------------\n"
);
eprintln!("Panic Info: {:#?}", panic_info);
eprintln!(
"\n--------------------------------------------------------------------------------\n"
);
// Capture and print the backtrace
// Requires RUST_BACKTRACE=1 (or full) to be set in the environment
let backtrace = Backtrace::capture();
eprintln!("Backtrace:\n{:?}", backtrace);
eprintln!(
"\n================================================================================\n"
);
// If tracing is still operational, try to log with it too.
// This might not always work if the panic is deep within tracing or stdio.
error!(panic_info = %panic_info, backtrace = ?backtrace, "Global panic hook caught a panic");
}));
debug!("Custom panic hook set.");
// Configure tracing to output to stderr
tracing_subscriber::fmt()
.with_writer(std::io::stderr)
.with_env_filter(EnvFilter::from_default_env().add_directive(Level::DEBUG.into()))
.init();
info!("Starting Wazuh MCP Server");
debug!("Loading environment variables...");
let wazuh_host = env::var("WAZUH_HOST").unwrap_or_else(|_| "localhost".to_string());
let wazuh_port = env::var("WAZUH_PORT")
.unwrap_or_else(|_| "55000".to_string())
.unwrap_or_else(|_| "9200".to_string())
.parse::<u16>()
.expect("WAZUH_PORT must be a valid port number");
let wazuh_user = env::var("WAZUH_USER").unwrap_or_else(|_| "admin".to_string());
let wazuh_pass = env::var("WAZUH_PASS").unwrap_or_else(|_| "admin".to_string());
let verify_ssl = env::var("VERIFY_SSL")
.unwrap_or_else(|_| "false".to_string())
.to_lowercase() == "true";
.to_lowercase()
== "true";
let mcp_server_port = env::var("MCP_SERVER_PORT")
.unwrap_or_else(|_| "8000".to_string())
.parse::<u16>()
.expect("MCP_SERVER_PORT must be a valid port number");
let wazuh_client = WazuhApiClient::new(
debug!(
wazuh_host,
wazuh_port,
wazuh_user,
// wazuh_pass is sensitive, avoid logging
verify_ssl,
mcp_server_port,
"Environment variables loaded."
);
info!("Initializing Wazuh API client...");
let wazuh_client = WazuhIndexerClient::new(
wazuh_host.clone(),
wazuh_port,
wazuh_user.clone(),
wazuh_pass.clone(),
verify_ssl,
);
let app_state = Arc::new(AppState {
wazuh_client: Mutex::new(wazuh_client),
});
let app = Router::new()
.route("/mcp", get(mcp_endpoint))
.route("/health", get(health_check))
.with_state(app_state);
debug!("AppState created.");
// Set up HTTP routes using the new http_service module
info!("Setting up HTTP routes...");
let app = create_http_router(app_state.clone());
debug!("HTTP routes configured.");
let addr = SocketAddr::from(([0, 0, 0, 0], mcp_server_port));
info!("Attempting to bind server to {}", addr);
let listener = tokio::net::TcpListener::bind(addr).await.unwrap_or_else(|e| {
error!("Failed to bind to address {}: {}", addr, e);
panic!("Failed to bind to address {}: {}", addr, e);
});
info!("Wazuh MCP Server listening on {}", addr);
axum::serve(listener, app.into_make_service())
info!("Attempting to bind HTTP server to {}", addr);
let listener = tokio::net::TcpListener::bind(addr)
.await
.unwrap_or_else(|e| {
error!("Server error: {}", e);
panic!("Server error: {}", e);
error!("Failed to bind to address {}: {}", addr, e);
panic!("Failed to bind to address {}: {}", addr, e);
});
}
info!("Wazuh MCP Server listening on {}", addr);
/// MCP endpoint for Claude Desktop.
/// Retrieves the latest Wazuh alerts, converts them into MCP messages, and returns as JSON.
async fn mcp_endpoint(
State(state): State<Arc<AppState>>,
) -> Result<Json<Vec<Value>>, (StatusCode, Json<Value>)> {
let alert_query = json!({
"query": {
"match_all": {}
}
// Spawn the stdio transport handler using the new stdio_service
info!("Spawning stdio service handler...");
let app_state_for_stdio = app_state.clone();
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let stdio_handle = tokio::spawn(async move {
run_stdio_service(app_state_for_stdio, shutdown_tx).await;
info!("run_stdio_service ASYNC TASK has completed its execution.");
});
let mut wazuh_client = state.wazuh_client.lock().await;
match wazuh_client.get_alerts(alert_query).await {
Ok(alerts_data) => {
let hits_array = alerts_data
.get("hits")
.and_then(|h| h.get("hits"))
.and_then(|h| h.as_array())
.cloned()
.unwrap_or_else(Vec::new);
let mcp_messages = hits_array
.iter()
.filter_map(|hit| {
hit.get("_source").map(|source| {
transform_to_mcp(source.clone(), "alert".to_string())
})
})
.collect::<Vec<_>>();
Ok(Json(mcp_messages))
// Configure Axum with graceful shutdown
let axum_shutdown_signal = async {
shutdown_rx
.await
.map_err(|e| error!("Shutdown signal sender dropped: {}", e))
.ok(); // Wait for the signal, log if sender is dropped
info!("Graceful shutdown signal received for Axum server. Axum will now attempt to shut down.");
};
info!("Starting Axum server with graceful shutdown.");
let axum_task = tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(axum_shutdown_signal)
.await
.unwrap_or_else(|e| {
error!("Axum Server run error: {}", e);
});
info!("Axum server task has completed and shut down.");
});
// Make handles mutable so select can take &mut
let mut stdio_handle = stdio_handle;
let mut axum_task = axum_task;
// Wait for either the stdio service or Axum server to complete.
tokio::select! {
biased; // Prioritize checking stdio_handle first if both are ready
stdio_res = &mut stdio_handle => {
match stdio_res {
Ok(_) => info!("Stdio service task completed. Axum's graceful shutdown should have been triggered if stdio initiated it."),
Err(e) => error!("Stdio service task failed or panicked: {:?}", e),
}
// Stdio has finished. If it didn't send a shutdown signal (e.g. due to panic before sending),
// Axum might still be running. The shutdown_tx being dropped will also trigger axum_shutdown_signal.
info!("Waiting for Axum server to fully shut down after stdio completion...");
match axum_task.await {
Ok(_) => info!("Axum server task completed successfully after stdio completion."),
Err(e) => error!("Axum server task failed or panicked after stdio completion: {:?}", e),
}
}
Err(e) => {
error!("Error in /mcp endpoint: {}", e);
Err((
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": e.to_string() })),
))
axum_res = &mut axum_task => {
match axum_res {
Ok(_) => info!("Axum server task completed (possibly due to graceful shutdown or error)."),
Err(e) => error!("Axum server task failed or panicked: {:?}", e),
}
// Axum has finished. The main function will now exit.
// We should wait for stdio_handle to complete or be cancelled.
info!("Axum finished. Waiting for stdio_handle to complete or be cancelled...");
match stdio_handle.await {
Ok(_) => info!("Stdio service task also completed after Axum finished."),
Err(e) => {
if e.is_cancelled() {
info!("Stdio service task was cancelled after Axum finished (expected if main is exiting).");
} else {
error!("Stdio service task failed or panicked after Axum finished: {:?}", e);
}
}
}
}
}
}
/// Health check endpoint.
/// Returns a simple JSON response to indicate the server is running.
async fn health_check() -> Json<Value> {
Json(json!({
"status": "ok",
"service": "wazuh-mcp-server",
"timestamp": chrono::Utc::now().to_rfc3339()
}))
info!("Main function is exiting.");
}

428
src/mcp/client.rs Normal file
View File

@ -0,0 +1,428 @@
use async_trait::async_trait;
use reqwest::Client;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serde_json::Value;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use thiserror::Error;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
#[derive(Error, Debug)]
pub enum McpClientError {
#[error("HTTP request error: {0}")]
HttpRequestError(#[from] reqwest::Error),
#[error("HTTP API error: status {status}, message: {message}")]
HttpApiError {
status: reqwest::StatusCode,
message: String,
},
#[error("JSON serialization/deserialization error: {0}")]
JsonError(#[from] serde_json::Error),
#[error("IO error: {0}")]
IoError(#[from] std::io::Error),
#[error("Failed to spawn child process: {0}")]
ProcessSpawnError(String),
#[error("Child process stdin/stdout not available")]
ProcessPipeError,
#[error("JSON-RPC error: code {code}, message: {message}, data: {data:?}")]
JsonRpcError {
code: i32,
message: String,
data: Option<Value>,
},
#[error("Received unexpected JSON-RPC response: {0}")]
UnexpectedResponse(String),
#[error("Operation timed out")]
Timeout,
#[error("Operation not supported in current mode: {0}")]
UnsupportedOperation(String),
}
#[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,
}
// --- JSON-RPC Structures (client-side definitions) ---
#[derive(Serialize, Debug)]
struct JsonRpcRequest<T: Serialize> {
jsonrpc: String,
method: String,
params: Option<T>,
id: Value, // Changed from usize to Value
}
#[derive(Deserialize, Debug)]
struct JsonRpcResponse<T> {
jsonrpc: String,
result: Option<T>,
error: Option<JsonRpcErrorData>,
id: Value, // Changed from usize to Value
}
#[derive(Deserialize, Debug)]
struct JsonRpcErrorData {
code: i32,
message: String,
data: Option<Value>,
}
#[derive(Deserialize, Debug, Clone)]
pub struct ServerInfo {
pub name: String,
pub version: String,
}
#[derive(Deserialize, Debug, Clone)]
pub struct InitializeResult {
pub protocol_version: String,
pub server_info: ServerInfo,
}
#[async_trait]
pub trait McpClientTrait {
async fn initialize(&mut self) -> Result<InitializeResult, McpClientError>;
async fn provide_context(
&mut self,
params: Option<Value>,
) -> Result<Vec<McpMessage>, McpClientError>;
async fn shutdown(&mut self) -> Result<(), McpClientError>;
}
enum ClientMode {
Http {
client: Client,
base_url: String,
},
Stdio {
stdin: ChildStdin,
stdout: BufReader<ChildStdout>,
},
}
pub struct McpClient {
mode: ClientMode,
child_process: Option<Child>, // Manages the lifetime of the child process
request_id_counter: AtomicUsize,
}
#[async_trait]
impl McpClientTrait for McpClient {
async fn initialize(&mut self) -> Result<InitializeResult, McpClientError> {
match &mut self.mode {
ClientMode::Http { .. } => Err(McpClientError::UnsupportedOperation(
"initialize is not supported in HTTP mode".to_string(),
)),
ClientMode::Stdio { .. } => {
let request_id = self.next_id();
self.send_stdio_request("initialize", None::<()>, request_id)
.await
}
}
}
async fn provide_context(
&mut self,
params: Option<Value>,
) -> Result<Vec<McpMessage>, McpClientError> {
match &mut self.mode {
ClientMode::Http { client, base_url } => {
let url = format!("{}/mcp", base_url);
let request_builder = if let Some(p) = params {
client.post(&url).json(&p)
} else {
client.get(&url)
};
let response = request_builder
.send()
.await
.map_err(McpClientError::HttpRequestError)?;
if !response.status().is_success() {
let status = response.status();
let message = response.text().await.unwrap_or_else(|_| {
format!("Failed to get error body for status {}", status)
});
return Err(McpClientError::HttpApiError { status, message });
}
response
.json::<Vec<McpMessage>>()
.await
.map_err(McpClientError::HttpRequestError)
}
ClientMode::Stdio { .. } => {
let request_id = self.next_id();
self.send_stdio_request("provideContext", params, request_id)
.await
}
}
}
async fn shutdown(&mut self) -> Result<(), McpClientError> {
match &mut self.mode {
ClientMode::Http { .. } => Err(McpClientError::UnsupportedOperation(
"shutdown is not supported in HTTP mode".to_string(),
)),
ClientMode::Stdio { .. } => {
let request_id = self.next_id();
// Attempt to send shutdown command, ignore error if server already closed pipe
let _result: Result<Option<Value>, McpClientError> = self
.send_stdio_request("shutdown", None::<()>, request_id)
.await;
// Always try to clean up the process
self.close_stdio_process().await
}
}
}
}
impl McpClient {
pub fn new_http(base_url: String) -> Self {
let client = Client::builder()
.timeout(Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client");
Self {
mode: ClientMode::Http { client, base_url },
child_process: None,
request_id_counter: AtomicUsize::new(1),
}
}
pub async fn new_stdio(
executable_path: &str,
envs: Option<Vec<(String, String)>>,
) -> Result<Self, McpClientError> {
let mut command = Command::new(executable_path);
command.stdin(std::process::Stdio::piped());
command.stdout(std::process::Stdio::piped());
command.stderr(std::process::Stdio::inherit()); // Pipe child's stderr to parent's stderr for visibility
if let Some(env_vars) = envs {
for (key, value) in env_vars {
command.env(key, value);
}
}
let mut child = command
.spawn()
.map_err(|e| McpClientError::ProcessSpawnError(e.to_string()))?;
let stdin = child.stdin.take().ok_or(McpClientError::ProcessPipeError)?;
let stdout = child
.stdout
.take()
.ok_or(McpClientError::ProcessPipeError)?;
Ok(Self {
mode: ClientMode::Stdio {
stdin,
stdout: BufReader::new(stdout),
},
child_process: Some(child),
request_id_counter: AtomicUsize::new(1),
})
}
fn next_id(&self) -> Value {
Value::from(self.request_id_counter.fetch_add(1, Ordering::SeqCst))
}
async fn send_stdio_request<P: Serialize, R: DeserializeOwned>(
&mut self,
method: &str,
params: Option<P>,
id: Value, // Added id parameter
) -> Result<R, McpClientError> {
// Removed: let request_id = self.next_id();
let rpc_request = JsonRpcRequest {
jsonrpc: "2.0".to_string(),
method: method.to_string(),
params,
id: id.clone(), // Use the provided id
};
let request_json = serde_json::to_string(&rpc_request)? + "\n";
let (stdin, stdout) = match &mut self.mode {
ClientMode::Stdio { stdin, stdout } => (stdin, stdout),
ClientMode::Http { .. } => {
return Err(McpClientError::UnsupportedOperation(
"send_stdio_request is only for Stdio mode".to_string(),
))
}
};
stdin.write_all(request_json.as_bytes()).await?;
stdin.flush().await?;
let mut response_json = String::new();
match tokio::time::timeout(
Duration::from_secs(10),
stdout.read_line(&mut response_json),
)
.await
{
Ok(Ok(0)) => {
return Err(McpClientError::IoError(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"Server closed stdout",
)))
}
Ok(Ok(_)) => { /* continue */ }
Ok(Err(e)) => return Err(McpClientError::IoError(e)),
Err(_) => return Err(McpClientError::Timeout),
}
let rpc_response: JsonRpcResponse<R> = serde_json::from_str(response_json.trim())?;
// Compare Value IDs. Note: Value implements PartialEq.
if rpc_response.id != id {
return Err(McpClientError::UnexpectedResponse(format!(
"Mismatched request/response IDs. Expected {}, got {}. Response: '{}'",
id, rpc_response.id, response_json
)));
}
if let Some(err_data) = rpc_response.error {
return Err(McpClientError::JsonRpcError {
code: err_data.code,
message: err_data.message,
data: err_data.data,
});
}
rpc_response.result.ok_or_else(|| {
McpClientError::UnexpectedResponse("Missing result in JSON-RPC response".to_string())
})
}
async fn close_stdio_process(&mut self) -> Result<(), McpClientError> {
if let Some(mut child) = self.child_process.take() {
child.kill().await.map_err(McpClientError::IoError)?;
let _ = child.wait().await; // Ensure process is reaped
}
Ok(())
}
// New public method for sending generic JSON-RPC requests
pub async fn send_json_rpc_request(
&mut self,
method: &str,
params: Option<Value>,
id: Value,
) -> Result<Value, McpClientError> {
match &mut self.mode {
ClientMode::Http { .. } => Err(McpClientError::UnsupportedOperation(
"Generic JSON-RPC calls are not supported in HTTP mode by this client.".to_string(),
)),
ClientMode::Stdio { .. } => {
// R (result type) is Value for generic calls
self.send_stdio_request(method, params, id).await
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use httpmock::prelude::*;
use serde_json::json;
use tokio;
#[tokio::test]
async fn test_mcp_client_http_get_data() {
// Renamed to be specific
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 mut client = McpClient::new_http(server.url("")); // Use new_http
// Use provide_context with None params for equivalent of old get_mcp_data
let result = client.provide_context(None).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_http_health_check_equivalent() {
// Renamed
let server = MockServer::start();
server.mock(|when, then| {
when.method(GET).path("/health"); // Assuming /health is still the target for this test
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_http(server.url(""));
// The new trait doesn't have a direct "check_health".
// If `initialize` was to be used for HTTP health, it would be:
// let result = client.initialize().await;
// But initialize is Stdio-only. So this test needs to adapt or be removed
// if there's no direct equivalent in the new trait for HTTP health.
// For now, let's assume we might add a specific http_health method if needed,
// or this test is demonstrating a capability that's no longer directly on the trait.
// To make this test pass with current structure, we'd need a separate HTTP health method.
// Let's simulate calling the /health endpoint directly if that's the intent.
let http_client = reqwest::Client::new();
let response = http_client.get(server.url("/health")).send().await.unwrap();
assert_eq!(response.status(), reqwest::StatusCode::OK);
let health_data: Value = response.json().await.unwrap();
assert_eq!(health_data["status"], "ok");
assert_eq!(health_data["service"], "wazuh-mcp-server");
}
// TODO: Add tests for Stdio mode. This would require a mock executable
// or a more complex test setup. For now, focusing on the client structure.
}

466
src/mcp/mcp_server_core.rs Normal file
View File

@ -0,0 +1,466 @@
use serde_json::{json, Value};
use std::sync::Arc;
use tracing::{debug, error, info};
use crate::mcp::protocol::{error_codes, JsonRpcError, JsonRpcRequest, JsonRpcResponse};
use crate::AppState;
// Structure to parse the parameters for a 'tools/call' request
#[derive(serde::Deserialize, Debug)]
struct ToolCallParams {
#[serde(rename = "name")]
name: String,
#[serde(rename = "arguments")]
arguments: Option<Value>, // Input parameters for the specific tool
#[serde(flatten)]
_extra: std::collections::HashMap<String, Value>,
}
pub struct McpServerCore {
app_state: Arc<AppState>,
}
impl McpServerCore {
pub fn new(app_state: Arc<AppState>) -> Self {
Self { app_state }
}
pub async fn process_request(&self, request: JsonRpcRequest) -> String {
info!("Processing request: method={}", request.method);
let response = match request.method.as_str() {
"initialize" => self.handle_initialize(request).await,
"shutdown" => self.handle_shutdown(request).await,
"provideContext" => self.handle_provide_context(request).await,
// Tool methods (prefix "tools/")
"tools/list" => self.handle_list_tools(request).await,
"tools/call" => self.handle_tool_call(request).await, // Use generic tool call handler
// "tools/wazuhAlerts" => self.handle_wazuh_alerts_tool(request).await,
// Resource methods (prefix "resources/")
"resources/list" => self.handle_get_resources(request).await,
"resources/read" => self.handle_read_resource(request).await,
// Prompt methods (prefix "prompts/")
"prompts/list" => self.handle_list_prompts(request).await,
_ => {
error!("Method not found: {}", request.method);
self.create_error_response(
error_codes::METHOD_NOT_FOUND,
format!("Method '{}' not found", request.method),
None,
request.id.clone(),
)
}
};
response
}
pub fn handle_parse_error(&self, error: serde_json::Error, raw_request: &str) -> String {
error!("Failed to parse JSON-RPC request: {}", error);
// Try to extract the ID from the raw request if possible
let id = serde_json::from_str::<Value>(raw_request)
.and_then(|v| {
if let Some(id) = v.get("id") {
Ok(id.clone())
} else {
// Use a different approach since custom is not available
Err(serde_json::Error::io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"No ID field found",
)))
}
})
.unwrap_or(Value::Null);
self.create_error_response(
error_codes::PARSE_ERROR,
format!("Parse error: {}", error),
None,
id,
)
}
async fn handle_initialize(&self, request: JsonRpcRequest) -> String {
debug!("Handling initialize request");
// Define the wazuhAlertSummary tool - simpler with no output schema
let wazuh_alert_summary_tool = crate::mcp::protocol::ToolDefinition {
name: "wazuhAlertSummary".to_string(),
description: Some("Returns a text summary of all Wazuh alerts.".to_string()),
// Define a minimal valid input schema (empty object)
input_schema: Some(json!({
"type": "object",
"properties": {}
})),
// No output schema needed as per requirements
output_schema: None,
};
// Use the protocol structs for better type safety and structure
let result = crate::mcp::protocol::InitializeResult {
protocol_version: "2024-11-05".to_string(),
capabilities: crate::mcp::protocol::Capabilities {
tools: crate::mcp::protocol::ToolCapability {
supported: true,
definitions: vec![wazuh_alert_summary_tool], // Only include wazuhAlertSummary tool
},
resources: crate::mcp::protocol::SupportedFeature { supported: true },
prompts: crate::mcp::protocol::SupportedFeature { supported: true },
},
server_info: crate::mcp::protocol::ServerInfo {
name: "Wazuh MCP Server".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
},
};
self.create_success_response(result, request.id)
}
async fn handle_shutdown(&self, request: JsonRpcRequest) -> String {
debug!("Handling shutdown request");
self.create_success_response(Value::Null, request.id)
}
async fn handle_provide_context(&self, request: JsonRpcRequest) -> String {
debug!("Handling provideContext request");
// Lock the Wazuh client to make API calls
let mut wazuh_client = self.app_state.wazuh_client.lock().await;
// Get alerts from Wazuh
match wazuh_client.get_alerts().await {
Ok(alerts) => {
let mcp_messages: Vec<Value> = alerts
.into_iter()
.map(|alert| crate::mcp::transform::transform_to_mcp(alert, "alert".to_string()))
.collect();
debug!("Transformed {} alerts into MCP messages for provideContext", mcp_messages.len());
self.create_success_response(json!(mcp_messages), request.id)
}
Err(e) => {
error!("Error getting alerts from Wazuh for provideContext: {}", e);
self.create_error_response(
error_codes::INTERNAL_ERROR,
format!("Failed to get alerts from Wazuh: {}", e),
None,
request.id,
)
}
}
}
async fn handle_get_resources(&self, request: JsonRpcRequest) -> String {
debug!("Handling getResources request");
// Return an empty list for now
let resources_result = crate::mcp::protocol::ResourcesListResult {
resources: vec![],
};
self.create_success_response(resources_result, request.id)
}
async fn handle_read_resource(&self, request: JsonRpcRequest) -> String {
debug!("Handling readResource request: {:?}", request.params);
#[derive(serde::Deserialize, Debug)]
struct ReadResourceParams {
uri: String,
// We can add _meta here if needed later
// _meta: Option<Value>,
}
let params: ReadResourceParams = match request.params {
Some(params_value) => match serde_json::from_value(params_value) {
Ok(p) => p,
Err(e) => {
error!("Failed to parse params for resources/read: {}", e);
return self.create_error_response(
error_codes::INVALID_PARAMS,
format!("Invalid params for resources/read: {}", e),
None,
request.id,
);
}
},
None => {
error!("Missing params for resources/read");
return self.create_error_response(
error_codes::INVALID_PARAMS,
"Missing params for resources/read, 'uri' is required".to_string(),
None,
request.id,
);
}
};
// Currently, no resources are supported for reading
error!("Unsupported URI for resources/read: {}", params.uri);
self.create_error_response(
error_codes::INVALID_PARAMS,
format!("Unsupported or unknown resource URI: {}", params.uri),
None,
request.id,
)
}
// Generic handler for executing tools via 'tools/call'
async fn handle_tool_call(&self, request: JsonRpcRequest) -> String {
debug!("Handling tools/call request: {:?}", request.params);
let params: ToolCallParams = match request.clone().params {
Some(params_value) => match serde_json::from_value(params_value) {
Ok(p) => p,
Err(e) => {
error!("Failed to parse params for tools/call: {}", e);
return self.create_error_response(
error_codes::INVALID_PARAMS,
format!("Invalid params for tools/call: {}", e),
None,
request.id,
);
}
},
None => {
error!("Missing params for tools/call");
return self.create_error_response(
error_codes::INVALID_PARAMS,
"Missing params for tools/call, 'name' and 'arguments' are required".to_string(),
None,
request.id,
);
}
};
// Dispatch based on the tool name
match params.name.as_str() {
"wazuhAlertSummary" => {
info!("Dispatching tools/call to wazuhAlertSummary handler");
self.handle_wazuh_alert_summary_tool(request).await
}
// wazuhAlerts tool is disabled but we keep the handler code
_ => {
error!("Unsupported tool name requested via tools/call: {}", params.name);
self.create_error_response(
error_codes::METHOD_NOT_FOUND, // Or a more specific tool error code if available
format!("Tool '{}' not found", params.name),
None,
request.id,
)
}
}
}
// Handler for listing available tools
async fn handle_list_tools(&self, request: JsonRpcRequest) -> String {
debug!("Handling tools/list request");
// Define the wazuhAlertSummary tool
let wazuh_alert_summary_tool = crate::mcp::protocol::ToolDefinition {
name: "wazuhAlertSummary".to_string(),
description: Some("Returns a text summary of all Wazuh alerts.".to_string()),
// Define a minimal valid input schema (empty object)
input_schema: Some(json!({
"type": "object",
"properties": {}
})),
// No output schema needed as per requirements
output_schema: None,
};
let tools_list = crate::mcp::protocol::ToolsListResult {
tools: vec![wazuh_alert_summary_tool],
};
self.create_success_response(tools_list, request.id)
}
async fn handle_wazuh_alerts_tool(&self, request: JsonRpcRequest) -> String {
debug!("Handling tools/wazuhAlerts request. Params: {:?}", request.params);
let wazuh_client = self.app_state.wazuh_client.lock().await;
match wazuh_client.get_alerts().await {
Ok(raw_alerts) => {
let simplified_alerts: Vec<Value> = raw_alerts
.into_iter()
.map(|alert| {
let source = alert.get("_source").unwrap_or(&alert);
// Extract ID: Try _source.id first, then _id
let id = source.get("id")
.and_then(|v| v.as_str())
.or_else(|| alert.get("_id").and_then(|v| v.as_str()))
.unwrap_or("") // Default to empty string if not found
.to_string();
// Extract Description: Look in _source.rule.description
let description = source.get("rule")
.and_then(|r| r.get("description"))
.and_then(|d| d.as_str())
.unwrap_or("") // Default to empty string if not found
.to_string();
json!({
"id": id,
"description": description,
})
})
.collect();
debug!("Processed {} alerts into simplified format.", simplified_alerts.len());
// Construct the final result with the "alerts" array
let result = json!({
"alerts": simplified_alerts,
"text": "Hello World",
});
self.create_success_response(result, request.id)
}
Err(e) => {
error!("Error getting alerts from Wazuh for tools/wazuhAlerts: {}", e);
self.create_error_response(
error_codes::INTERNAL_ERROR,
format!("Failed to get alerts from Wazuh: {}", e),
None,
request.id,
)
}
}
}
// Handler for the wazuhAlertSummary tool
async fn handle_wazuh_alert_summary_tool(&self, request: JsonRpcRequest) -> String {
debug!("Handling tools/wazuhAlertSummary request. Params: {:?}", request.params);
let mut wazuh_client = self.app_state.wazuh_client.lock().await;
match wazuh_client.get_alerts().await {
Ok(raw_alerts) => {
// Create a content item for each alert
let content_items: Vec<Value> = if raw_alerts.is_empty() {
// If no alerts, return a single "no alerts" message
vec![json!({
"type": "text",
"text": "No Wazuh alerts found."
})]
} else {
// Map each alert to a content item
raw_alerts
.into_iter()
.map(|alert| {
let source = alert.get("_source").unwrap_or(&alert);
// Extract alert ID
let id = source.get("id")
.and_then(|v| v.as_str())
.or_else(|| alert.get("_id").and_then(|v| v.as_str()))
.unwrap_or("Unknown ID");
// Extract rule description
let description = source.get("rule")
.and_then(|r| r.get("description"))
.and_then(|d| d.as_str())
.unwrap_or("No description available");
// Extract timestamp if available
let timestamp = source.get("timestamp")
.and_then(|t| t.as_str())
.unwrap_or("Unknown time");
// Format the alert as a text entry and create a content item
json!({
"type": "text",
"text": format!("Alert ID: {}\nTime: {}\nDescription: {}", id, timestamp, description)
})
})
.collect()
};
debug!("Processed {} alerts into individual content items.", content_items.len());
// Construct the final result with the content array containing multiple text objects
let result = json!({
"content": content_items
});
self.create_success_response(result, request.id)
}
Err(e) => {
error!("Error getting alerts from Wazuh for tools/wazuhAlertSummary: {}", e);
self.create_error_response(
error_codes::INTERNAL_ERROR,
format!("Failed to get alerts from Wazuh: {}", e),
None,
request.id,
)
}
}
}
async fn handle_list_prompts(&self, request: JsonRpcRequest) -> String {
debug!("Handling prompts/list request");
// Define the single prompt according to the new structure
let list_alerts_prompt = crate::mcp::protocol::PromptEntry {
name: "list-wazuh-alerts".to_string(),
description: Some("List the latest security alerts from Wazuh.".to_string()),
arguments: vec![], // This prompt takes no arguments
};
let prompts = vec![list_alerts_prompt];
let result = crate::mcp::protocol::PromptsListResult { prompts };
self.create_success_response(result, request.id)
}
fn create_success_response<T: serde::Serialize>(&self, result: T, id: Value) -> String {
let response = JsonRpcResponse {
jsonrpc: "2.0".to_string(),
result: Some(result),
error: None,
id,
};
serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize JSON-RPC response: {}", e);
format!(
r#"{{"jsonrpc":"2.0","error":{{"code":-32603,"message":"Internal error: Failed to serialize response"}},"id":null}}"#
)
})
}
fn create_error_response(
&self,
code: i32,
message: String,
data: Option<Value>,
id: Value,
) -> String {
let response = JsonRpcResponse::<Value> {
jsonrpc: "2.0".to_string(),
result: None,
error: Some(JsonRpcError {
code,
message,
data,
}),
id,
};
serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize JSON-RPC error response: {}", e);
format!(
r#"{{"jsonrpc":"2.0","error":{{"code":-32603,"message":"Internal error: Failed to serialize error response"}},"id":null}}"#
)
})
}
}

3
src/mcp/mod.rs
View File

@ -1 +1,4 @@
pub mod transform;
pub mod client;
pub mod protocol;
pub mod mcp_server_core;

127
src/mcp/protocol.rs Normal file
View File

@ -0,0 +1,127 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Deserialize, Debug, Clone)]
pub struct JsonRpcRequest {
pub jsonrpc: String,
pub method: String,
pub params: Option<Value>,
pub id: Value,
}
#[derive(Serialize, Debug)]
pub struct JsonRpcResponse<T: Serialize> {
pub jsonrpc: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<T>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<JsonRpcError>,
pub id: Value,
}
#[derive(Serialize, Debug)]
pub struct JsonRpcError {
pub code: i32,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub data: Option<Value>,
}
#[derive(Serialize, Debug, Clone)]
pub struct ToolDefinition {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(rename = "inputSchema", skip_serializing_if = "Option::is_none")]
pub input_schema: Option<Value>, // Added inputSchema
#[serde(rename = "outputSchema", skip_serializing_if = "Option::is_none")]
pub output_schema: Option<Value>, // Added outputSchema
}
#[derive(Serialize, Debug)]
pub struct ToolCapability {
pub supported: bool,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub definitions: Vec<ToolDefinition>, // List available tools
}
#[derive(Serialize, Debug)]
pub struct SupportedFeature {
pub supported: bool,
}
#[derive(Serialize, Debug)]
pub struct Capabilities {
pub tools: ToolCapability, // Use the new structure
pub resources: SupportedFeature,
pub prompts: SupportedFeature,
}
#[derive(Serialize, Debug)]
pub struct ServerInfo {
pub name: String,
pub version: String,
}
#[derive(Serialize, Debug)]
pub struct InitializeResult {
#[serde(rename = "protocolVersion")]
pub protocol_version: String,
pub capabilities: Capabilities,
#[serde(rename = "serverInfo")]
pub server_info: ServerInfo,
}
#[derive(Serialize, Debug)]
pub struct ResourceEntry {
pub uri: String,
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(rename = "mimeType", skip_serializing_if = "Option::is_none")]
pub mime_type: Option<String>,
}
#[derive(Serialize, Debug)]
pub struct ResourcesListResult {
pub resources: Vec<ResourceEntry>,
}
#[derive(Serialize, Debug)]
pub struct ToolsListResult {
pub tools: Vec<ToolDefinition>,
}
#[derive(Serialize, Debug, Clone)]
pub struct PromptArgument {
pub name: String,
pub required: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub default: Option<Value>, // Use Value for flexibility (string, bool, number)
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>, // Optional description for the argument
}
#[derive(Serialize, Debug, Clone)]
pub struct PromptEntry {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub arguments: Vec<PromptArgument>,
}
#[derive(Serialize, Debug)]
pub struct PromptsListResult {
pub prompts: Vec<PromptEntry>,
}
pub mod error_codes {
pub const PARSE_ERROR: i32 = -32700;
pub const INVALID_REQUEST: i32 = -32600;
pub const METHOD_NOT_FOUND: i32 = -32601;
pub const INVALID_PARAMS: i32 = -32602;
pub const INTERNAL_ERROR: i32 = -32603;
pub const SERVER_ERROR_START: i32 = -32000;
pub const SERVER_ERROR_END: i32 = -32099;
}

51
src/mcp/transform.rs
View File

@ -1,27 +1,41 @@
use chrono::{DateTime, Utc, SecondsFormat};
use serde_json::{json, Value};
use tracing::warn;
use tracing::{debug, warn};
pub fn transform_to_mcp(event: Value, event_type: String) -> Value {
debug!(?event, %event_type, "Entering transform_to_mcp");
let source_obj = event.get("_source").unwrap_or(&event);
if event.get("_source").is_some() {
debug!("Event contains '_source' field, using it for transformation.");
} else {
debug!("Event does not contain '_source' field, using the event root for transformation.");
}
let id = source_obj.get("id")
.and_then(|v| v.as_str())
.or_else(|| event.get("_id").and_then(|v| v.as_str()))
.unwrap_or("unknown_id")
.to_string();
debug!(%id, "Transformed: id");
let default_rule = json!({});
let rule = source_obj.get("rule").unwrap_or(&default_rule);
if source_obj.get("rule").is_none() {
debug!("Transformed: rule (defaulted to empty object)");
} else {
debug!(?rule, "Transformed: rule");
}
let category = rule.get("groups")
.and_then(|g| g.as_array())
.and_then(|arr| arr.first())
.and_then(|v| v.as_str())
.unwrap_or("unknown_category")
.to_string();
debug!(%category, "Transformed: category");
let severity = rule.get("level")
.and_then(|v| v.as_u64())
let severity_level = rule.get("level").and_then(|v| v.as_u64());
let severity = severity_level
.map(|level| match level {
0..=3 => "low",
4..=7 => "medium",
@ -30,21 +44,38 @@ pub fn transform_to_mcp(event: Value, event_type: String) -> Value {
})
.unwrap_or("unknown_severity")
.to_string();
debug!(?severity_level, %severity, "Transformed: severity");
let description = rule.get("description")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
debug!(%description, "Transformed: description");
let default_data = json!({});
let data = source_obj.get("data").cloned().unwrap_or(default_data);
let data = source_obj.get("data").cloned().unwrap_or_else(|| {
debug!("Transformed: data (defaulted to empty object)");
default_data.clone()
});
if source_obj.get("data").is_some() {
debug!(?data, "Transformed: data");
}
let default_agent = json!({});
let agent = source_obj.get("agent").cloned().unwrap_or(default_agent);
let agent = source_obj.get("agent").cloned().unwrap_or_else(|| {
debug!("Transformed: agent (defaulted to empty object)");
default_agent.clone()
});
if source_obj.get("agent").is_some() {
debug!(?agent, "Transformed: agent");
}
let timestamp_str = source_obj.get("timestamp")
.and_then(|v| v.as_str())
.unwrap_or("");
debug!(%timestamp_str, "Attempting to parse timestamp");
let timestamp = DateTime::parse_from_rfc3339(timestamp_str)
.map(|dt| dt.with_timezone(&Utc))
@ -53,11 +84,13 @@ pub fn transform_to_mcp(event: Value, event_type: String) -> Value {
warn!("Failed to parse timestamp '{}' for alert ID '{}'. Using current time.", timestamp_str, id);
Utc::now()
});
debug!(%timestamp, "Transformed: timestamp");
let notes = "Data fetched via Wazuh API".to_string();
debug!(%notes, "Transformed: notes");
json!({
"protocol_version": "1.0",
let mcp_message = json!({
"protocolVersion": "1.0", // Match initialize response
"source": "Wazuh",
"timestamp": timestamp.to_rfc3339_opts(SecondsFormat::Secs, true),
"event_type": event_type,
@ -73,7 +106,9 @@ pub fn transform_to_mcp(event: Value, event_type: String) -> Value {
"integration": "Wazuh-MCP",
"notes": notes
}
})
});
debug!(?mcp_message, "Exiting transform_to_mcp with result");
mcp_message
}
#[cfg(test)]

129
src/stdio_service.rs Normal file
View File

@ -0,0 +1,129 @@
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::sync::oneshot::Sender as OneshotSender;
use tracing::{debug, error, info};
use crate::logging_utils::{log_mcp_request, log_mcp_response};
use crate::mcp::mcp_server_core::McpServerCore;
use crate::mcp::protocol::JsonRpcRequest;
use crate::AppState;
pub async fn run_stdio_service(app_state: Arc<AppState>, shutdown_tx: OneshotSender<()>) {
info!("Starting MCP server in stdio mode...");
let mut stdin_reader = BufReader::new(tokio::io::stdin());
let mut stdout_writer = tokio::io::stdout();
let mcp_core = McpServerCore::new(app_state);
let mut line_buffer = String::new();
debug!("run_stdio_service: Initialized readers/writers. Entering main loop.");
loop {
debug!("stdio_service: Top of the loop. Clearing line buffer.");
line_buffer.clear();
debug!("stdio_service: About to read_line from stdin.");
let read_result = stdin_reader.read_line(&mut line_buffer).await;
debug!(?read_result, "stdio_service: read_line completed.");
match read_result {
Ok(0) => {
debug!("stdio_service: read_line returned Ok(0) (EOF).");
info!("Stdin closed (EOF), signaling shutdown and exiting stdio mode.");
let _ = shutdown_tx.send(()); // Signal main to shutdown Axum
debug!("stdio_service read 0 bytes, breaking loop.");
break; // EOF
}
Ok(bytes_read) => {
debug!(%bytes_read, "stdio_service: read_line returned Ok(bytes_read).");
let request_str = line_buffer.trim();
if request_str.is_empty() {
debug!("Received empty line from stdin, continuing.");
continue;
}
info!("Received from stdin (stdio_service): {}", request_str);
// Log the raw request using the utility
log_mcp_request(request_str);
// Process the request using the core module
let response_json = match serde_json::from_str::<JsonRpcRequest>(request_str) {
Ok(rpc_request) => {
// Special handling for shutdown to exit the loop
let is_shutdown = rpc_request.method == "shutdown";
let response = mcp_core.process_request(rpc_request).await;
if is_shutdown {
// Log the response using the utility
log_mcp_response(&response);
// Send the response
if let Err(e) = stdout_writer
.write_all(format!("{}\n", response).as_bytes())
.await
{
error!("Error writing shutdown response to stdout: {}", e);
}
if let Err(e) = stdout_writer.flush().await {
error!("Error flushing stdout for shutdown: {}", e);
}
debug!("Signaling shutdown and exiting stdio_service due to 'shutdown' request.");
let _ = shutdown_tx.send(()); // Signal main to shutdown Axum
return; // Exit the loop and function
}
response
}
Err(e) => mcp_core.handle_parse_error(e, request_str),
};
// Log the raw response using the utility
log_mcp_response(&response_json);
info!("Sending to stdout (stdio_service): {}", response_json);
// Prepare the response string with a newline
let response_to_send = format!("{}\n", response_json);
debug!(
"Attempting to write response to stdout. Length: {} bytes. Preview (up to 200 chars): '{}'",
response_to_send.len(),
response_to_send.chars().take(200).collect::<String>()
);
// Write the response and handle potential errors
match stdout_writer.write_all(response_to_send.as_bytes()).await {
Ok(_) => {
debug!("Successfully wrote response bytes to stdout buffer.");
// Flush immediately after write
if let Err(e) = stdout_writer.flush().await {
error!("Error flushing stdout after successful write: {}", e);
debug!(
"Signaling shutdown and breaking loop due to stdout flush error."
);
let _ = shutdown_tx.send(());
break;
} else {
debug!("Successfully flushed stdout.");
}
}
Err(e) => {
error!("Error writing response to stdout: {}", e);
debug!("Signaling shutdown and breaking loop due to stdout write error.");
let _ = shutdown_tx.send(());
break;
}
}
}
Err(e) => {
debug!(error = %e, "stdio_service: read_line returned Err.");
error!("Error reading from stdin for stdio_service: {}", e);
debug!("Signaling shutdown and breaking loop due to stdin read error.");
let _ = shutdown_tx.send(()); // Signal main to shutdown Axum
break;
}
}
debug!("stdio_service: Bottom of the loop, before next iteration.");
}
info!("run_stdio_service: Exited main loop. stdio_service task is finishing.");
}

195
src/wazuh/client.rs
View File

@ -1,32 +1,35 @@
use reqwest::{header, Client};
use serde_json::Value;
use std::time::{Duration, SystemTime};
use tracing::{info, warn};
use reqwest::{header, Client, Method};
use serde_json::{json, Value};
use std::time::Duration;
use tracing::{debug, error, info};
use super::error::WazuhApiError;
pub struct WazuhApiClient {
#[derive(Debug, Clone)]
pub struct WazuhIndexerClient {
username: String,
password: String,
base_url: String,
jwt_token: Option<String>,
jwt_expiration: Option<SystemTime>,
auth_endpoint: String,
http_client: Client,
}
impl WazuhApiClient {
// Renamed impl block
impl WazuhIndexerClient {
pub fn new(
host: String,
port: u16,
indexer_port: u16,
username: String,
password: String,
verify_ssl: bool,
) -> Self {
let base_url = format!("https://{}:{}", host, port);
debug!(%host, indexer_port, %username, %verify_ssl, "Creating new WazuhIndexerClient");
// Base URL now points to the Indexer
let base_url = format!("https://{}:{}", host, indexer_port);
debug!(%base_url, "Wazuh Indexer base URL set");
let http_client = Client::builder()
.danger_accept_invalid_certs(!verify_ssl)
.timeout(Duration::from_secs(10))
.timeout(Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client");
@ -34,135 +37,93 @@ impl WazuhApiClient {
username,
password,
base_url,
jwt_token: None,
jwt_expiration: None,
auth_endpoint: "/security/user/authenticate".to_string(),
http_client,
}
}
fn is_jwt_valid(&self) -> bool {
match (self.jwt_token.as_ref(), self.jwt_expiration) {
(Some(_), Some(expiration)) => match expiration.duration_since(SystemTime::now()) {
Ok(remaining) => remaining.as_secs() > 60,
Err(_) => false,
},
_ => false,
}
}
pub async fn get_jwt(&mut self) -> Result<String, WazuhApiError> {
if self.is_jwt_valid() {
return Ok(self.jwt_token.clone().unwrap());
}
let auth_url = format!("{}{}", self.base_url, self.auth_endpoint);
info!("Requesting new JWT token from {}", auth_url);
let response = self
.http_client
.post(&auth_url)
.basic_auth(&self.username, Some(&self.password))
.send()
.await?;
if !response.status().is_success() {
let status = response.status();
let error_text = response
.text()
.await
.unwrap_or_else(|_| "Unknown error".to_string());
return Err(WazuhApiError::AuthenticationError(format!(
"Authentication failed with status {}: {}",
status, error_text
)));
}
let data: Value = response.json().await?;
let token = data
.get("jwt")
.and_then(|t| t.as_str())
.ok_or(WazuhApiError::JwtNotFound)?
.to_string();
self.jwt_token = Some(token.clone());
self.jwt_expiration = Some(SystemTime::now() + Duration::from_secs(5 * 60));
info!("Obtained new JWT token valid for 5 minutes");
Ok(token)
}
async fn make_request(
&mut self,
method: reqwest::Method,
async fn make_indexer_request(
&self,
method: Method,
endpoint: &str,
body: Option<Value>,
) -> Result<Value, WazuhApiError> {
let jwt_token = self.get_jwt().await?;
debug!(?method, %endpoint, ?body, "Making request to Wazuh Indexer");
let url = format!("{}{}", self.base_url, endpoint);
debug!(%url, "Constructed Indexer request URL");
let mut request_builder = self
.http_client
.request(method.clone(), &url)
.header(header::AUTHORIZATION, format!("Bearer {}", jwt_token));
.basic_auth(&self.username, Some(&self.password)); // Use Basic Auth
if let Some(json_body) = &body {
request_builder = request_builder.json(json_body);
request_builder = request_builder
.header(header::CONTENT_TYPE, "application/json")
.json(json_body);
}
debug!("Request builder configured with Basic Auth");
let response = request_builder.send().await?;
let status = response.status();
debug!(%status, "Received response from Indexer endpoint");
if response.status() == reqwest::StatusCode::UNAUTHORIZED {
warn!("JWT expired. Re-authenticating and retrying request.");
self.jwt_token = None;
let new_jwt_token = self.get_jwt().await?;
let mut retry_builder = self
.http_client
.request(method, &url)
.header(header::AUTHORIZATION, format!("Bearer {}", new_jwt_token));
if let Some(json_body) = &body {
retry_builder = retry_builder.json(json_body);
}
let retry_response = retry_builder.send().await?;
if !retry_response.status().is_success() {
let status = retry_response.status();
let error_text = retry_response
.text()
.await
.unwrap_or_else(|_| "Unknown error".to_string());
return Err(WazuhApiError::ApiError(format!(
"API request failed with status {}: {}",
status, error_text
)));
}
Ok(retry_response.json().await?)
} else if !response.status().is_success() {
let status = response.status();
if !status.is_success() {
let error_text = response
.text()
.await
.unwrap_or_else(|_| "Unknown error".to_string());
Err(WazuhApiError::ApiError(format!(
"API request failed with status {}: {}",
status, error_text
)))
} else {
Ok(response.json().await?)
.unwrap_or_else(|_| "Unknown error reading response body".to_string());
error!(%url, %status, %error_text, "Indexer API request failed");
// Provide more context in the error
return Err(WazuhApiError::ApiError(format!(
"Indexer request to {} failed with status {}: {}",
url, status, error_text
)));
}
debug!("Indexer API request successful");
response.json().await.map_err(|e| {
error!("Failed to parse JSON response from Indexer: {}", e);
WazuhApiError::RequestError(e) // Use appropriate error variant
})
}
pub async fn get_alerts(&mut self, query: Value) -> Result<Value, WazuhApiError> {
let index_pattern = "wazuh-alerts-*";
let endpoint = format!("/{}_search", index_pattern);
pub async fn get_alerts(&self) -> Result<Vec<Value>, WazuhApiError> {
let endpoint = "/wazuh-alerts*/_search";
let query_body = json!({
"size": 100,
"query": {
"match_all": {}
},
});
info!("Retrieving alerts with index pattern '{}'", index_pattern);
self.make_request(reqwest::Method::GET, &endpoint, Some(query))
.await
debug!(%endpoint, ?query_body, "Preparing to get alerts from Wazuh Indexer");
info!("Retrieving up to 100 alerts from Wazuh Indexer");
let response = self
.make_indexer_request(Method::POST, endpoint, Some(query_body))
.await?;
let hits = response
.get("hits")
.and_then(|h| h.get("hits"))
.and_then(|h_array| h_array.as_array())
.ok_or_else(|| {
error!(
?response,
"Failed to find 'hits.hits' array in Indexer response"
);
WazuhApiError::ApiError("Indexer response missing 'hits.hits' array".to_string())
})?;
let alerts: Vec<Value> = hits
.iter()
.filter_map(|hit| hit.get("_source").cloned())
.collect();
debug!(
"Successfully retrieved {} alerts from Indexer",
alerts.len()
);
Ok(alerts)
}
}

215
tests/mcp_client_cli.rs
View File

@ -1,95 +1,162 @@
use anyhow::Result;
use anyhow::{anyhow, Result};
use clap::Parser;
use serde_json::Value;
use std::env;
use std::process;
use std::io::{self, Write}; // For stdout().flush() and stdin().read_line()
mod mcp_client;
use mcp_client::{McpClient, McpClientTrait};
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 args: Vec<String> = env::args().collect();
let cli_args = CliArgs::parse();
let mut client: McpClient;
let is_stdio_mode = cli_args.stdio_exe.is_some();
if args.len() < 2 {
eprintln!("Usage: {} <command> [options]", args[0]);
eprintln!("Commands:");
eprintln!(" get-data - Get MCP data from the server");
eprintln!(" health - Check server health");
eprintln!(" query - Query MCP data with filters");
process::exit(1);
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.
}
let command = &args[1];
let mcp_url =
env::var("MCP_SERVER_URL").unwrap_or_else(|_| "http://localhost:8000".to_string());
println!("\nInteractive MCP Client. Enter a JSON-RPC request, 'health' (HTTP only), or 'quit'.");
println!("Press CTRL-D for EOF to exit.");
println!("Connecting to MCP server at: {}", mcp_url);
let client = McpClient::new(mcp_url);
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 command.as_str() {
"get-data" => {
println!("Fetching MCP data...");
let data = client.get_mcp_data().await?;
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;
}
println!("Received {} MCP messages:", data.len());
for (i, message) in data.iter().enumerate() {
println!("\nMessage {}:", i + 1);
println!(" Source: {}", message.source);
println!(" Event Type: {}", message.event_type);
println!(" Timestamp: {}", message.timestamp);
if line.eq_ignore_ascii_case("quit") {
println!("Exiting.");
break;
}
let context = &message.context;
println!(" Context:");
println!(" ID: {}", context["id"]);
println!(" Category: {}", context["category"]);
println!(" Severity: {}", context["severity"]);
println!(" Description: {}", context["description"]);
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;
}
if let Some(data) = context.get("data").and_then(|d| d.as_object()) {
println!(" Data:");
for (key, value) in data {
println!(" {}: {}", key, value);
// 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'.");
}
}
}
}
"health" => {
println!("Checking server health...");
let health = client.check_health().await?;
println!("Health status: {}", health["status"]);
println!("Service: {}", health["service"]);
println!("Timestamp: {}", health["timestamp"]);
}
"query" => {
if args.len() < 3 {
eprintln!("Error: Missing query parameters");
eprintln!("Usage: {} query <json_filter>", args[0]);
process::exit(1);
}
let filter_str = &args[2];
let filters: Value = serde_json::from_str(filter_str)?;
println!("Querying MCP data with filters: {}", filters);
let data = client.query_mcp_data(filters).await?;
println!("Received {} MCP messages:", data.len());
for (i, message) in data.iter().enumerate() {
println!("\nMessage {}:", i + 1);
println!(" Source: {}", message.source);
println!(" Event Type: {}", message.event_type);
let context = &message.context;
println!(" Context:");
println!(" ID: {}", context["id"]);
println!(" Category: {}", context["category"]);
println!(" Severity: {}", context["severity"]);
Err(e) => {
eprintln!("Error reading input: {}. Exiting.", e);
break;
}
}
_ => {
eprintln!("Error: Unknown command '{}'", command);
process::exit(1);
}
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),
}
}

44
tests/run_tests.sh
View File

@ -6,10 +6,42 @@ cargo test
echo "Building MCP client CLI..."
cargo build --bin mcp_client_cli
if nc -z localhost 8000 2>/dev/null; then
echo "Testing MCP client CLI against running server..."
./target/debug/mcp_client_cli health
./target/debug/mcp_client_cli get-data
else
echo "MCP server is not running. Start it with 'cargo run' to test the CLI."
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 "Testing MCP client CLI in stdio mode..."
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 "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
# 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
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 "MCP client CLI stdio tests completed."