From 77914b509746f0ebe31190b1372d2d2bb0b78310 Mon Sep 17 00:00:00 2001
From: Gianluca Brigandi <gbrigand@gmail.com>
Date: Tue, 24 Jun 2025 12:03:58 -0700
Subject: [PATCH] Upgraded to latest wazuh client that fixes fetching critical
 vulnerabilities

---
 Cargo.toml  |  2 +-
 src/main.rs | 17 ++++++++---------
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 9a58f83..40e5aca 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,7 +9,7 @@ repository = "https://github.com/gbrigandi/mcp-server-wazuh"
 readme = "README.md"
 
 [dependencies]
-wazuh-client = "0.1.3" 
+wazuh-client = "0.1.4" 
 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/main.rs b/src/main.rs
index 194b0b1..c2caf2a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -83,7 +83,7 @@ use std::env;
 use clap::Parser;
 use dotenv::dotenv;
 
-use wazuh_client::{WazuhClientFactory, WazuhIndexerClient, RulesClient, VulnerabilityClient, AgentsClient, LogsClient, ClusterClient, Port as WazuhPort};
+use wazuh_client::{WazuhClientFactory, WazuhIndexerClient, RulesClient, VulnerabilityClient, VulnerabilitySeverity, AgentsClient, LogsClient, ClusterClient, Port as WazuhPort};
 
 #[derive(Parser, Debug)]
 #[command(name = "mcp-server-wazuh")]
@@ -532,7 +532,7 @@ impl WazuhToolsServer {
                 &agent_id,
                 Some(1000), // Get more results to filter
                 Some(offset),
-                params.severity.as_deref(),
+                params.severity.as_deref().and_then(VulnerabilitySeverity::from_str)
             ).await {
                 Ok(all_vulns) => {
                     let filtered: Vec<_> = all_vulns
@@ -549,7 +549,7 @@ impl WazuhToolsServer {
                 &agent_id,
                 Some(limit),
                 Some(offset),
-                params.severity.as_deref(),
+                params.severity.as_deref().and_then(VulnerabilitySeverity::from_str)
             ).await
         };
 
@@ -566,12 +566,11 @@ impl WazuhToolsServer {
                 let mcp_content_items: Vec<Content> = vulnerabilities
                     .into_iter()
                     .map(|vuln| {
-                        let severity_indicator = match vuln.severity.to_lowercase().as_str() {
-                            "critical" => "🔴 CRITICAL",
-                            "high" => "🟠 HIGH",
-                            "medium" => "🟡 MEDIUM", 
-                            "low" => "🟢 LOW",
-                            _ => &vuln.severity,
+                        let severity_indicator = match vuln.severity {
+                            VulnerabilitySeverity::Critical => "🔴 CRITICAL",
+                            VulnerabilitySeverity::High => "🟠 HIGH",
+                            VulnerabilitySeverity::Medium => "🟡 MEDIUM", 
+                            VulnerabilitySeverity::Low => "🟢 LOW",
                         };
 
                         let published_info = if let Some(published) = &vuln.published {