From bfffdbfb52537f9734aea016e975af812d94ffd4 Mon Sep 17 00:00:00 2001 From: Gianluca Brigandi Date: Tue, 17 Jun 2025 22:46:55 -0700 Subject: [PATCH] Replaced bespoke indexer client with full fledges indexer and manager API crate. --- Cargo.toml | 1 + src/lib.rs | 10 ++-- src/main.rs | 51 ++++++++++------ src/wazuh/client.rs | 140 -------------------------------------------- src/wazuh/error.rs | 26 -------- src/wazuh/mod.rs | 2 - 6 files changed, 41 insertions(+), 189 deletions(-) delete mode 100644 src/wazuh/client.rs delete mode 100644 src/wazuh/error.rs delete mode 100644 src/wazuh/mod.rs diff --git a/Cargo.toml b/Cargo.toml index 0305a6d..17d71e3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ repository = "https://github.com/gbrigandi/mcp-server-wazuh" readme = "README.md" [dependencies] +wazuh-client = "0.1.0" rmcp = { version = "0.1.5", features = ["server", "transport-io"] } tokio = { version = "1", features = ["full"] } reqwest = { version = "0.12", features = ["json", "rustls-tls"], default-features = false } diff --git a/src/lib.rs b/src/lib.rs index 507b8d5..bb8e9b4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,6 @@ -pub mod wazuh; - -pub use wazuh::client::WazuhIndexerClient; -pub use wazuh::error::WazuhApiError; +// Re-export the wazuh-client crate types for convenience +pub use wazuh_client::{ + WazuhClientFactory, WazuhClients, WazuhIndexerClient, WazuhApiError, + AgentsClient, RulesClient, ConfigurationClient, VulnerabilityClient, + ActiveResponseClient, ClusterClient, LogsClient, ConnectivityStatus +}; diff --git a/src/main.rs b/src/main.rs index 3b507ff..4883392 100644 --- a/src/main.rs +++ b/src/main.rs @@ -49,12 +49,7 @@ use std::env; use clap::Parser; use dotenv::dotenv; -mod wazuh { - pub mod client; - pub mod error; -} - -use wazuh::client::WazuhIndexerClient; +use wazuh_client::{WazuhClientFactory, WazuhIndexerClient}; #[derive(Parser, Debug)] #[command(name = "mcp-server-wazuh")] @@ -72,7 +67,9 @@ struct GetAlertSummaryParams { #[derive(Clone)] struct WazuhToolsServer { - wazuh_client: Arc, + #[allow(dead_code)] // Kept for future expansion to other Wazuh clients + wazuh_factory: Arc, + wazuh_indexer_client: Arc, } #[tool(tool_box)] @@ -81,10 +78,21 @@ impl WazuhToolsServer { dotenv().ok(); let wazuh_host = env::var("WAZUH_HOST").unwrap_or_else(|_| "localhost".to_string()); - let wazuh_port = env::var("WAZUH_PORT") + let wazuh_api_port = env::var("WAZUH_API_PORT") + .unwrap_or_else(|_| "55000".to_string()) + .parse::() + .map_err(|e| anyhow::anyhow!("Invalid WAZUH_API_PORT: {}", e))?; + let wazuh_indexer_port = env::var("WAZUH_INDEXER_PORT") .unwrap_or_else(|_| "9200".to_string()) .parse::() - .map_err(|e| anyhow::anyhow!("Invalid WAZUH_PORT: {}", e))?; + .map_err(|e| anyhow::anyhow!("Invalid WAZUH_INDEXER_PORT: {}", e))?; + + // For backward compatibility, also check WAZUH_PORT + let wazuh_port = env::var("WAZUH_PORT") + .ok() + .and_then(|p| p.parse::().ok()) + .unwrap_or(wazuh_indexer_port); + 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") @@ -95,17 +103,26 @@ impl WazuhToolsServer { let protocol = env::var("WAZUH_TEST_PROTOCOL").unwrap_or_else(|_| "https".to_string()); tracing::debug!(?protocol, "Using Wazuh protocol for client from WAZUH_TEST_PROTOCOL or default"); - let wazuh_client = WazuhIndexerClient::new_with_protocol( - wazuh_host, - wazuh_port, - wazuh_user, - wazuh_pass, + // Create the factory with both API and Indexer configurations + let wazuh_factory = WazuhClientFactory::new( + wazuh_host.clone(), // API host + wazuh_api_port, // API port + wazuh_user.clone(), // API username + wazuh_pass.clone(), // API password + wazuh_host, // Indexer host (same as API host) + wazuh_port, // Indexer port (use WAZUH_PORT for backward compatibility) + wazuh_user, // Indexer username + wazuh_pass, // Indexer password verify_ssl, - &protocol, + Some(protocol), ); + // Create the indexer client using the factory + let wazuh_indexer_client = wazuh_factory.create_indexer_client(); + Ok(Self { - wazuh_client: Arc::new(wazuh_client), + wazuh_factory: Arc::new(wazuh_factory), + wazuh_indexer_client: Arc::new(wazuh_indexer_client), }) } @@ -121,7 +138,7 @@ impl WazuhToolsServer { tracing::info!(limit = %limit, "Retrieving Wazuh alert summary"); - match self.wazuh_client.get_alerts().await { + match self.wazuh_indexer_client.get_alerts().await { Ok(raw_alerts) => { let alerts_to_process: Vec<_> = raw_alerts.into_iter().take(limit as usize).collect(); diff --git a/src/wazuh/client.rs b/src/wazuh/client.rs deleted file mode 100644 index 264e5ee..0000000 --- a/src/wazuh/client.rs +++ /dev/null @@ -1,140 +0,0 @@ -use reqwest::{header, Client, Method}; -use serde_json::{json, Value}; -use std::time::Duration; -use tracing::{debug, error, info}; - -use super::error::WazuhApiError; - -#[derive(Debug, Clone)] -pub struct WazuhIndexerClient { - username: String, - password: String, - base_url: String, - http_client: Client, -} - -impl WazuhIndexerClient { - #[allow(dead_code)] - pub fn new( - host: String, - indexer_port: u16, - username: String, - password: String, - verify_ssl: bool, - ) -> Self { - Self::new_with_protocol(host, indexer_port, username, password, verify_ssl, "https") - } - - pub fn new_with_protocol( - host: String, - indexer_port: u16, - username: String, - password: String, - verify_ssl: bool, - protocol: &str, - ) -> Self { - debug!(%host, indexer_port, %username, %verify_ssl, %protocol, "Creating new WazuhIndexerClient"); - // Base URL now points to the Indexer - let base_url = format!("{}://{}:{}", protocol, 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(30)) - .build() - .expect("Failed to create HTTP client"); - - Self { - username, - password, - base_url, - http_client, - } - } - - async fn make_indexer_request( - &self, - method: Method, - endpoint: &str, - body: Option, - ) -> Result { - 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) - .basic_auth(&self.username, Some(&self.password)); // Use Basic Auth - - if let Some(json_body) = &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 !status.is_success() { - let error_text = response - .text() - .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(&self) -> Result, WazuhApiError> { - let endpoint = "/wazuh-alerts*/_search"; - let query_body = json!({ - "size": 100, - "query": { - "match_all": {} - }, - }); - - 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 = hits - .iter() - .filter_map(|hit| hit.get("_source").cloned()) - .collect(); - - debug!( - "Successfully retrieved {} alerts from Indexer", - alerts.len() - ); - Ok(alerts) - } -} diff --git a/src/wazuh/error.rs b/src/wazuh/error.rs deleted file mode 100644 index 4563600..0000000 --- a/src/wazuh/error.rs +++ /dev/null @@ -1,26 +0,0 @@ -use thiserror::Error; - -#[derive(Error, Debug)] -#[allow(dead_code)] -pub enum WazuhApiError { - #[error("Failed to create HTTP client: {0}")] - HttpClientCreationError(reqwest::Error), - - #[error("HTTP request error: {0}")] - RequestError(#[from] reqwest::Error), - - #[error("JWT token not found in Wazuh API response")] - JwtNotFound, - - #[error("Wazuh API Authentication failed: {0}")] - AuthenticationError(String), - - #[error("JSON parsing error: {0}")] - JsonError(#[from] serde_json::Error), - - #[error("Wazuh API error: {0}")] - ApiError(String), - - #[error("Alert with ID '{0}' not found")] - AlertNotFound(String), -} diff --git a/src/wazuh/mod.rs b/src/wazuh/mod.rs deleted file mode 100644 index f27e761..0000000 --- a/src/wazuh/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod client; -pub mod error;