wazuh-client 0.1.8

use reqwest::Method;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::fmt;
use tracing::{debug, info, warn};

use super::indexer_client::WazuhIndexerClient;
use super::wazuh_client::WazuhApiClient;

mod string_or_number_as_string {
    use serde::{Deserialize, Deserializer};

    pub fn deserialize<'de, D>(deserializer: D) -> Result<String, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(untagged)]
        enum StringOrNumber {
            Str(String),
            Num(serde_json::Number),
        }

        match StringOrNumber::deserialize(deserializer)? {
            StringOrNumber::Str(s) => Ok(s),
            StringOrNumber::Num(n) => Ok(n.to_string()),
        }
    }

    pub fn deserialize_optional<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(untagged)]
        enum StringOrNumberOrNull {
            Str(String),
            Num(serde_json::Number),
        }

        let intermediate: Option<StringOrNumberOrNull> = Option::deserialize(deserializer)?;

        match intermediate {
            Some(StringOrNumberOrNull::Str(s)) => Ok(Some(s)),
            Some(StringOrNumberOrNull::Num(n)) => Ok(Some(n.to_string())),
            None => Ok(None),
        }
    }
}
use super::error::WazuhApiError;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub enum VulnerabilitySeverity {
    Critical,
    High,
    Medium,
    Low,
}

impl VulnerabilitySeverity {
    /// Convert to the format expected by the Wazuh indexer (first letter uppercase, rest lowercase)
    pub fn to_indexer_format(&self) -> &'static str {
        match self {
            VulnerabilitySeverity::Critical => "Critical",
            VulnerabilitySeverity::High => "High",
            VulnerabilitySeverity::Medium => "Medium",
            VulnerabilitySeverity::Low => "Low",
        }
    }

    /// Parse from string (case-insensitive)
    pub fn parse(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "critical" => Some(VulnerabilitySeverity::Critical),
            "high" => Some(VulnerabilitySeverity::High),
            "medium" => Some(VulnerabilitySeverity::Medium),
            "low" => Some(VulnerabilitySeverity::Low),
            _ => None,
        }
    }

    /// Get all severity levels in order from highest to lowest
    pub fn all() -> [VulnerabilitySeverity; 4] {
        [
            VulnerabilitySeverity::Critical,
            VulnerabilitySeverity::High,
            VulnerabilitySeverity::Medium,
            VulnerabilitySeverity::Low,
        ]
    }
}

impl fmt::Display for VulnerabilitySeverity {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.to_indexer_format())
    }
}

impl From<VulnerabilitySeverity> for String {
    fn from(severity: VulnerabilitySeverity) -> Self {
        severity.to_indexer_format().to_string()
    }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Vulnerability {
    pub cve: String,
    pub title: String,
    #[serde(deserialize_with = "deserialize_severity")]
    pub severity: VulnerabilitySeverity,
    pub published: Option<String>,
    pub updated: Option<String>,
    pub reference: Option<String>,
    pub description: Option<String>,
    pub cvss: Option<CvssScore>,
    pub detection_time: Option<String>,
    pub agent_id: Option<String>,
    pub agent_name: Option<String>,
}

fn deserialize_severity<'de, D>(deserializer: D) -> Result<VulnerabilitySeverity, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let s = String::deserialize(deserializer)?;
    VulnerabilitySeverity::parse(&s)
        .ok_or_else(|| serde::de::Error::custom(format!("Invalid severity: {}", s)))
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CvssScore {
    pub cvss2: Option<CvssDetails>,
    pub cvss3: Option<CvssDetails>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CvssDetails {
    pub vector: Option<String>,
    pub base_score: Option<f64>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Package {
    pub name: String,
    pub version: String,
    pub architecture: Option<String>,
    pub format: Option<String>,
    pub description: Option<String>,
    pub size: Option<u64>,
    pub vendor: Option<String>,
    pub multiarch: Option<String>,
    pub source: Option<String>,
    pub priority: Option<String>,
    pub scan_id: Option<u64>,
    pub section: Option<String>,
    pub agent_id: Option<String>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Process {
    #[serde(deserialize_with = "string_or_number_as_string::deserialize")]
    pub pid: String,
    pub name: String,
    pub state: Option<String>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub ppid: Option<String>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub utime: Option<String>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub stime: Option<String>,
    pub cmd: Option<String>,
    pub argvs: Option<String>,
    pub euser: Option<String>,
    pub ruser: Option<String>,
    pub suser: Option<String>,
    pub egroup: Option<String>,
    pub rgroup: Option<String>,
    pub sgroup: Option<String>,
    pub fgroup: Option<String>,
    pub priority: Option<i32>,
    pub nice: Option<i32>,
    pub size: Option<u64>,
    pub vm_size: Option<u64>,
    pub resident: Option<u64>,
    pub share: Option<u64>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub start_time: Option<String>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub pgrp: Option<String>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub session: Option<String>,
    pub nlwp: Option<u32>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub tgid: Option<String>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub tty: Option<String>,
    #[serde(
        default,
        deserialize_with = "string_or_number_as_string::deserialize_optional"
    )]
    pub processor: Option<String>,
    pub scan_id: Option<u64>,
    pub agent_id: Option<String>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Port {
    pub local: PortInfo,
    pub remote: Option<PortInfo>,
    pub state: Option<String>,
    pub protocol: String,
    pub tx_queue: Option<u32>,
    pub rx_queue: Option<u32>,
    pub inode: Option<u64>,
    pub process: Option<String>,
    pub scan_id: Option<u64>,
    pub agent_id: Option<String>,
    pub pid: Option<u32>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct VulnerabilitySummaryCounts {
    pub critical: i32,
    pub high: i32,
    pub medium: i32,
    pub low: i32,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct VulnerabilitySummaryResponseData {
    pub agent_id: String,
    pub summary: VulnerabilitySummaryCounts,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PortInfo {
    pub ip: Option<String>,
    pub port: u16,
}

#[derive(Debug, Clone)]
pub struct VulnerabilityClient {
    api_client: WazuhApiClient,
    indexer_client: WazuhIndexerClient,
}

impl VulnerabilityClient {
    pub fn new(api_client: WazuhApiClient, indexer_client: WazuhIndexerClient) -> Self {
        Self {
            api_client,
            indexer_client,
        }
    }

    pub async fn get_agent_vulnerabilities(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
        offset: Option<u32>,
        severity: Option<VulnerabilitySeverity>,
    ) -> Result<Vec<Vulnerability>, WazuhApiError> {
        debug!(%agent_id, ?severity, "Getting vulnerabilities for agent");

        let size = limit.unwrap_or(100);
        let from = offset.unwrap_or(0);

        let mut must_clauses = vec![json!({
            "term": {
                "agent.id": agent_id
            }
        })];

        if let Some(severity) = severity {
            must_clauses.push(json!({
                "term": {
                    "vulnerability.severity": severity.to_indexer_format()
                }
            }));
        }

        let query_body = json!({
            "size": size,
            "from": from,
            "query": {
                "bool": {
                    "must": must_clauses
                }
            }
        });

        let endpoint = "/wazuh-states-vulnerabilities*/_search";
        let response = self
            .indexer_client
            .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(|| {
                WazuhApiError::ApiError("Missing 'hits.hits' in vulnerability response".to_string())
            })?;

        let mut vulnerabilities = Vec::new();
        for hit in hits {
            if let Some(source) = hit.get("_source") {
                if let Ok(vuln) = self.parse_vulnerability_from_source(source, agent_id) {
                    vulnerabilities.push(vuln);
                }
            }
        }

        info!(%agent_id, "Retrieved {} vulnerabilities", vulnerabilities.len());
        Ok(vulnerabilities)
    }

    fn parse_vulnerability_from_source(
        &self,
        source: &Value,
        agent_id: &str,
    ) -> Result<Vulnerability, WazuhApiError> {
        let vulnerability_data = source.get("vulnerability").ok_or_else(|| {
            WazuhApiError::ApiError("Missing vulnerability data in source".to_string())
        })?;

        let cve = vulnerability_data
            .get("cve")
            .and_then(|v| v.as_str())
            .unwrap_or("Unknown")
            .to_string();

        let title = vulnerability_data
            .get("title")
            .and_then(|v| v.as_str())
            .unwrap_or("No title available")
            .to_string();

        let severity = vulnerability_data
            .get("severity")
            .and_then(|v| v.as_str())
            .and_then(VulnerabilitySeverity::parse)
            .unwrap_or(VulnerabilitySeverity::Low); // Default to Low if parsing fails

        let published = vulnerability_data
            .get("published")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let updated = vulnerability_data
            .get("updated")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let reference = vulnerability_data
            .get("reference")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let description = vulnerability_data
            .get("description")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let detection_time = source
            .get("timestamp")
            .or_else(|| source.get("@timestamp"))
            .or_else(|| source.get("vulnerability.detection_time"))
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let agent_name = source
            .get("agent")
            .and_then(|a| a.get("name"))
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let cvss = self.parse_cvss_scores(vulnerability_data);

        Ok(Vulnerability {
            cve,
            title,
            severity,
            published,
            updated,
            reference,
            description,
            cvss,
            detection_time,
            agent_id: Some(agent_id.to_string()),
            agent_name,
        })
    }

    fn parse_cvss_scores(&self, vulnerability_data: &Value) -> Option<CvssScore> {
        let cvss2 = vulnerability_data
            .get("cvss2")
            .map(|cvss2_data| CvssDetails {
                vector: cvss2_data
                    .get("vector")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string()),
                base_score: cvss2_data.get("base_score").and_then(|v| v.as_f64()),
            });

        let cvss3 = vulnerability_data
            .get("cvss3")
            .map(|cvss3_data| CvssDetails {
                vector: cvss3_data
                    .get("vector")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string()),
                base_score: cvss3_data.get("base_score").and_then(|v| v.as_f64()),
            });

        if cvss2.is_some() || cvss3.is_some() {
            Some(CvssScore { cvss2, cvss3 })
        } else {
            None
        }
    }

    pub async fn get_critical_vulnerabilities(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
    ) -> Result<Vec<Vulnerability>, WazuhApiError> {
        debug!(%agent_id, "Getting critical vulnerabilities");
        self.get_agent_vulnerabilities(agent_id, limit, None, Some(VulnerabilitySeverity::Critical))
            .await
    }

    pub async fn get_high_vulnerabilities(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
    ) -> Result<Vec<Vulnerability>, WazuhApiError> {
        debug!(%agent_id, "Getting high severity vulnerabilities");
        self.get_agent_vulnerabilities(agent_id, limit, None, Some(VulnerabilitySeverity::High))
            .await
    }

    pub async fn get_medium_vulnerabilities(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
    ) -> Result<Vec<Vulnerability>, WazuhApiError> {
        debug!(%agent_id, "Getting medium severity vulnerabilities");
        self.get_agent_vulnerabilities(agent_id, limit, None, Some(VulnerabilitySeverity::Medium))
            .await
    }

    pub async fn get_low_vulnerabilities(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
    ) -> Result<Vec<Vulnerability>, WazuhApiError> {
        debug!(%agent_id, "Getting low severity vulnerabilities");
        self.get_agent_vulnerabilities(agent_id, limit, None, Some(VulnerabilitySeverity::Low))
            .await
    }

    pub async fn get_agent_packages(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
        offset: Option<u32>,
        search: Option<&str>,
    ) -> Result<Vec<Package>, WazuhApiError> {
        debug!(%agent_id, ?search, "Getting packages for agent");

        let mut query_params = Vec::new();

        if let Some(limit) = limit {
            query_params.push(("limit", limit.to_string()));
        }
        if let Some(offset) = offset {
            query_params.push(("offset", offset.to_string()));
        }
        if let Some(search) = search {
            query_params.push(("search", search.to_string()));
        }

        let query_params_ref: Vec<(&str, &str)> =
            query_params.iter().map(|(k, v)| (*k, v.as_str())).collect();

        let endpoint = format!("/syscollector/{}/packages", agent_id);
        let response = self
            .api_client
            .make_request(
                Method::GET,
                &endpoint,
                None,
                if query_params_ref.is_empty() {
                    None
                } else {
                    Some(&query_params_ref)
                },
            )
            .await?;

        let packages_data = response
            .get("data")
            .and_then(|d| d.get("affected_items"))
            .ok_or_else(|| {
                WazuhApiError::ApiError(
                    "Missing 'data.affected_items' in packages response".to_string(),
                )
            })?;

        let packages: Vec<Package> = serde_json::from_value(packages_data.clone())?;
        info!(%agent_id, "Retrieved {} packages", packages.len());
        Ok(packages)
    }

    pub async fn get_agent_processes(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
        offset: Option<u32>,
        search: Option<&str>,
    ) -> Result<Vec<Process>, WazuhApiError> {
        debug!(%agent_id, ?search, "Getting processes for agent");

        let mut query_params = Vec::new();

        if let Some(limit) = limit {
            query_params.push(("limit", limit.to_string()));
        }
        if let Some(offset) = offset {
            query_params.push(("offset", offset.to_string()));
        }
        if let Some(search) = search {
            query_params.push(("search", search.to_string()));
        }

        let query_params_ref: Vec<(&str, &str)> =
            query_params.iter().map(|(k, v)| (*k, v.as_str())).collect();

        let endpoint = format!("/syscollector/{}/processes", agent_id);
        let response = self
            .api_client
            .make_request(
                Method::GET,
                &endpoint,
                None,
                if query_params_ref.is_empty() {
                    None
                } else {
                    Some(&query_params_ref)
                },
            )
            .await?;

        let processes_data = response
            .get("data")
            .and_then(|d| d.get("affected_items"))
            .ok_or_else(|| {
                WazuhApiError::ApiError(
                    "Missing 'data.affected_items' in processes response".to_string(),
                )
            })?;

        let processes: Vec<Process> = serde_json::from_value(processes_data.clone())?;
        info!(%agent_id, "Retrieved {} processes", processes.len());
        Ok(processes)
    }

    pub async fn get_agent_ports(
        &mut self,
        agent_id: &str,
        limit: Option<u32>,
        offset: Option<u32>,
        protocol: Option<&str>,
    ) -> Result<Vec<Port>, WazuhApiError> {
        debug!(%agent_id, ?protocol, "Getting ports for agent");

        let mut query_params = Vec::new();

        if let Some(limit) = limit {
            query_params.push(("limit", limit.to_string()));
        }
        if let Some(offset) = offset {
            query_params.push(("offset", offset.to_string()));
        }
        if let Some(protocol) = protocol {
            query_params.push(("protocol", protocol.to_string()));
        }

        let query_params_ref: Vec<(&str, &str)> =
            query_params.iter().map(|(k, v)| (*k, v.as_str())).collect();

        let endpoint = format!("/syscollector/{}/ports", agent_id);
        let response = self
            .api_client
            .make_request(
                Method::GET,
                &endpoint,
                None,
                if query_params_ref.is_empty() {
                    None
                } else {
                    Some(&query_params_ref)
                },
            )
            .await?;

        let ports_data = response
            .get("data")
            .and_then(|d| d.get("affected_items"))
            .ok_or_else(|| {
                WazuhApiError::ApiError(
                    "Missing 'data.affected_items' in ports response".to_string(),
                )
            })?;

        let ports: Vec<Port> = serde_json::from_value(ports_data.clone())?;
        info!(%agent_id, "Retrieved {} ports", ports.len());
        Ok(ports)
    }

    pub async fn get_vulnerability_summary(
        &mut self,
        agent_id: &str,
    ) -> Result<Option<VulnerabilitySummaryResponseData>, WazuhApiError> {
        debug!(%agent_id, "Getting vulnerability summary");

        let query_body = json!({
            "size": 0,
            "query": {
                "bool": {
                    "must": [
                        {
                            "term": {
                                "agent.id": agent_id
                            }
                        }
                    ]
                }
            },
            "aggs": {
                "severity_counts": {
                    "terms": {
                        "field": "vulnerability.severity",
                        "size": 10
                    }
                }
            }
        });

        let endpoint = "/wazuh-states-vulnerabilities*/_search";
        match self
            .indexer_client
            .make_indexer_request(Method::POST, endpoint, Some(query_body))
            .await
        {
            Ok(response_value) => {
                let aggregations = response_value
                    .get("aggregations")
                    .and_then(|aggs| aggs.get("severity_counts"))
                    .and_then(|severity| severity.get("buckets"))
                    .and_then(|buckets| buckets.as_array());

                if let Some(buckets) = aggregations {
                    let mut critical = 0;
                    let mut high = 0;
                    let mut medium = 0;
                    let mut low = 0;

                    for bucket in buckets {
                        if let (Some(key), Some(count)) = (
                            bucket.get("key").and_then(|k| k.as_str()),
                            bucket.get("doc_count").and_then(|c| c.as_i64()),
                        ) {
                            if let Some(severity) = VulnerabilitySeverity::parse(key) {
                                match severity {
                                    VulnerabilitySeverity::Critical => critical = count as i32,
                                    VulnerabilitySeverity::High => high = count as i32,
                                    VulnerabilitySeverity::Medium => medium = count as i32,
                                    VulnerabilitySeverity::Low => low = count as i32,
                                }
                            }
                        }
                    }

                    let summary = VulnerabilitySummaryResponseData {
                        agent_id: agent_id.to_string(),
                        summary: VulnerabilitySummaryCounts {
                            critical,
                            high,
                            medium,
                            low,
                        },
                    };

                    info!(%agent_id, "Retrieved vulnerability summary");
                    Ok(Some(summary))
                } else {
                    info!(%agent_id, "No vulnerability data found for summary");
                    Ok(None)
                }
            }
            Err(e) => {
                warn!(%agent_id, "Failed to get vulnerability summary: {}", e);
                Ok(None)
            }
        }
    }

    pub async fn search_package(
        &mut self,
        package_name: &str,
        agent_ids: Option<&[String]>,
    ) -> Result<Vec<(String, Vec<Package>)>, WazuhApiError> {
        debug!(%package_name, "Searching for package across agents");

        let mut results = Vec::new();

        if let Some(agent_ids) = agent_ids {
            for agent_id in agent_ids {
                match self
                    .get_agent_packages(agent_id, None, None, Some(package_name))
                    .await
                {
                    Ok(packages) => {
                        if !packages.is_empty() {
                            results.push((agent_id.clone(), packages));
                        }
                    }
                    Err(e) => {
                        debug!(%agent_id, %package_name, "Failed to get packages: {}", e);
                    }
                }
            }
        }

        info!(%package_name, "Found package in {} agents", results.len());
        Ok(results)
    }

    pub async fn get_agents_with_vulnerability(
        &mut self,
        cve: &str,
        agent_ids: &[String],
        limit: Option<u32>,
    ) -> Result<Vec<String>, WazuhApiError> {
        debug!(%cve, "Finding agents with specific vulnerability");

        let mut affected_agents = Vec::new();

        for agent_id in agent_ids {
            match self
                .get_agent_vulnerabilities(agent_id, limit, None, None)
                .await
            {
                Ok(vulnerabilities) => {
                    if vulnerabilities.iter().any(|v| v.cve == cve) {
                        affected_agents.push(agent_id.clone());
                    }
                }
                Err(e) => {
                    debug!(%agent_id, %cve, "Failed to get vulnerabilities: {}", e);
                }
            }
        }

        info!(%cve, "Found {} agents with vulnerability", affected_agents.len());
        Ok(affected_agents)
    }

    /// Get vulnerabilities with severity at or above the specified level
    pub async fn get_vulnerabilities_by_min_severity(
        &mut self,
        agent_id: &str,
        min_severity: VulnerabilitySeverity,
        limit: Option<u32>,
        offset: Option<u32>,
    ) -> Result<Vec<Vulnerability>, WazuhApiError> {
        debug!(%agent_id, ?min_severity, "Getting vulnerabilities with minimum severity");

        // Get all vulnerabilities and filter by severity level
        let all_vulnerabilities = self
            .get_agent_vulnerabilities(agent_id, limit, offset, None)
            .await?;

        let filtered: Vec<Vulnerability> = all_vulnerabilities
            .into_iter()
            .filter(|vuln| match (vuln.severity, min_severity) {
                (VulnerabilitySeverity::Critical, _) => true,
                (VulnerabilitySeverity::High, VulnerabilitySeverity::Critical) => false,
                (VulnerabilitySeverity::High, _) => true,
                (
                    VulnerabilitySeverity::Medium,
                    VulnerabilitySeverity::Critical | VulnerabilitySeverity::High,
                ) => false,
                (VulnerabilitySeverity::Medium, _) => true,
                (VulnerabilitySeverity::Low, VulnerabilitySeverity::Low) => true,
                (VulnerabilitySeverity::Low, _) => false,
            })
            .collect();

        info!(%agent_id, ?min_severity, "Found {} vulnerabilities at or above severity level", filtered.len());
        Ok(filtered)
    }

    /// Get vulnerabilities for multiple severity levels
    pub async fn get_vulnerabilities_by_severities(
        &mut self,
        agent_id: &str,
        severities: &[VulnerabilitySeverity],
        limit: Option<u32>,
        offset: Option<u32>,
    ) -> Result<Vec<Vulnerability>, WazuhApiError> {
        debug!(%agent_id, ?severities, "Getting vulnerabilities for multiple severity levels");

        let mut all_vulnerabilities = Vec::new();

        for &severity in severities {
            let vulns = self
                .get_agent_vulnerabilities(agent_id, limit, offset, Some(severity))
                .await?;
            all_vulnerabilities.extend(vulns);
        }

        // Remove duplicates based on CVE
        all_vulnerabilities.sort_by(|a, b| a.cve.cmp(&b.cve));
        all_vulnerabilities.dedup_by(|a, b| a.cve == b.cve);

        info!(%agent_id, ?severities, "Found {} unique vulnerabilities", all_vulnerabilities.len());
        Ok(all_vulnerabilities)
    }
}