esdiag 0.16.4

Elastic Stack diagnostic collector and processor
// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
// or more contributor license agreements. Licensed under the Elastic License 2.0;
// you may not use this file except in compliance with the Elastic License 2.0.

/// Builder for the Elasticsearch client
mod elasticsearch;
/// Client for Kibana APIs
mod kibana;
/// Client for Logstash APIs
mod logstash;

pub use elasticsearch::{ElasticsearchBuilder, ElasticsearchClient};
pub(crate) use kibana::KIBANA_REQUEST_CONCURRENCY;
pub use kibana::KibanaClient;
pub use logstash::LogstashClient;

extern crate elasticsearch as es;
use crate::data::{Product, Uri};
use eyre::{Result, eyre};
use reqwest::Method;
use serde::de::DeserializeOwned;
use std::collections::HashMap;

/// A response from an Elastic Stack client.
pub enum ClientResponse {
    Elasticsearch(es::http::response::Response),
    Http(reqwest::Response),
}

impl ClientResponse {
    pub fn status(&self) -> reqwest::StatusCode {
        match self {
            Self::Elasticsearch(response) => {
                reqwest::StatusCode::from_u16(response.status_code().as_u16()).expect("valid HTTP status code")
            }
            Self::Http(response) => response.status(),
        }
    }

    pub async fn text(self) -> Result<String> {
        match self {
            Self::Elasticsearch(response) => Ok(response.text().await?),
            Self::Http(response) => Ok(response.text().await?),
        }
    }

    pub async fn bytes(self) -> Result<bytes::Bytes> {
        match self {
            Self::Elasticsearch(response) => Ok(response.bytes().await?),
            Self::Http(response) => Ok(response.bytes().await?),
        }
    }

    pub async fn json<T: DeserializeOwned>(self) -> Result<T> {
        match self {
            Self::Elasticsearch(response) => Ok(response.json().await?),
            Self::Http(response) => Ok(response.json().await?),
        }
    }
}

/// A standardized client for interacting with Elastic Stack APIs
pub enum Client {
    Elasticsearch(ElasticsearchClient),
    Kibana(KibanaClient),
    Logstash(LogstashClient),
}

impl Client {
    /// Send an HTTP request to a path on the client's base URL
    pub async fn request(
        &self,
        method: Method,
        headers: &HashMap<String, String>,
        path: &str,
        body: Option<&[u8]>,
    ) -> Result<ClientResponse> {
        tracing::debug!("Request: {method} {path}");
        match self {
            Client::Elasticsearch(client) => {
                let method = match method {
                    Method::GET => es::http::Method::Get,
                    Method::POST => es::http::Method::Post,
                    Method::PUT => es::http::Method::Put,
                    Method::DELETE => es::http::Method::Delete,
                    Method::HEAD => es::http::Method::Head,
                    _ => return Err(eyre!("Unsupported http method for Elasticsearch client")),
                };
                let header_map: es::http::headers::HeaderMap = headers
                    .iter()
                    .filter_map(|(k, v)| match (k.parse(), v.parse()) {
                        (Ok(k), Ok(v)) => Some((k, v)),
                        x => {
                            tracing::warn!("Failed to parse header: {:?}", x);
                            None
                        }
                    })
                    .collect();
                use es::http::request::JsonBody;
                let body: Option<JsonBody<serde_json::Value>> =
                    body.map(serde_json::from_slice).transpose()?.map(JsonBody::new);
                let response = client
                    .send(method, path, header_map, Option::<&serde_json::Value>::None, body, None)
                    .await?;
                Ok(ClientResponse::Elasticsearch(response))
            }
            Client::Kibana(client) => client
                .request(method, headers, path, body)
                .await
                .map(ClientResponse::Http),
            Client::Logstash(client) => client
                .request(method, headers, path, body)
                .await
                .map(ClientResponse::Http),
        }
    }

    /// Verify the connection and authentication to the stack component
    pub async fn test_connection(&self) -> std::result::Result<String, String> {
        match self {
            Client::Elasticsearch(client) => {
                let response = client
                    .send(
                        es::http::Method::Get,
                        "/",
                        es::http::headers::HeaderMap::new(),
                        Option::<&serde_json::Value>::None,
                        Option::<es::http::request::JsonBody<serde_json::Value>>::None,
                        None,
                    )
                    .await
                    .map_err(|e| format!("{e}"))?;

                let status = response.status_code();
                let json: serde_json::Value = response
                    .json::<serde_json::Value>()
                    .await
                    .map_err(|e| format!("Failed to read test body: {e}"))?;
                tracing::debug!("Test response {} ", json);

                if json.get("tagline").is_some() {
                    Ok(format!("{} ✅ Elasticsearch", status))
                } else {
                    Err(format!("{} ❌ Root response did not match Elasticsearch", status))
                }
            }
            Client::Kibana(client) => {
                let response = client.test_connection().await.map_err(|e| format!("{e}"))?;
                let status = response.status();
                let json: serde_json::Value = response
                    .json::<serde_json::Value>()
                    .await
                    .map_err(|e| format!("Failed to read test body: {e}"))?;
                tracing::debug!("Test response {} ", json);
                match json.get("name") {
                    Some(name) => Ok(format!("{status} ✅ Kibana: {name}")),
                    None => Err(format!("{status} ❌ Host is not an Kibana node!")),
                }
            }
            Client::Logstash(client) => {
                let response = client.test_connection().await.map_err(|e| format!("{e}"))?;
                let status = response.status();
                let json: serde_json::Value = response
                    .json::<serde_json::Value>()
                    .await
                    .map_err(|e| format!("Failed to read test body: {e}"))?;
                tracing::debug!("Test response {} ", json);

                if let Some(version) = json.get("version").and_then(|v| v.as_str()) {
                    let name = json.get("name").and_then(|v| v.as_str()).unwrap_or("unknown");
                    Ok(format!("{status} ✅ Logstash: {name} ({version})"))
                } else {
                    Err(format!("{} ❌ Root response did not match Logstash", status))
                }
            }
        }
    }

    /// Check if security is enabled on the cluster.
    ///
    /// For Elasticsearch, this checks the `security.enabled` flag in `/_xpack/usage`.
    /// For Kibana and Logstash, this currently always returns `true`.
    pub async fn has_security_enabled(&self) -> Result<bool> {
        match self {
            Client::Elasticsearch(client) => {
                let response = client
                    .send(
                        es::http::Method::Get,
                        "/_xpack/usage",
                        es::http::headers::HeaderMap::new(),
                        Option::<&serde_json::Value>::None,
                        Option::<es::http::request::JsonBody<serde_json::Value>>::None,
                        None,
                    )
                    .await?;

                let status = response.status_code();
                if status.is_success() {
                    let json: serde_json::Value = response.json().await?;
                    let enabled = json
                        .get("security")
                        .and_then(|s| s.get("enabled"))
                        .and_then(|e| e.as_bool())
                        .unwrap_or(true);
                    Ok(enabled)
                } else {
                    match status.as_u16() {
                        401 | 403 => {
                            tracing::debug!(
                                "Security detection returned {status}. Security is enabled but access to /_xpack/usage is restricted."
                            );
                            Ok(true)
                        }
                        404 => {
                            tracing::debug!(
                                "Security detection returned 404. Assuming security is disabled or not supported."
                            );
                            Ok(false)
                        }
                        _ => {
                            tracing::warn!("Failed to check security status (HTTP {status}).");
                            Err(eyre!("Failed to check security status: HTTP {status}"))
                        }
                    }
                }
            }
            Client::Kibana(_) => {
                // For Kibana we assume true for now as requested
                Ok(true)
            }
            Client::Logstash(_) => {
                // Logstash does not expose the Elasticsearch security usage endpoint
                Ok(true)
            }
        }
    }
}

impl From<Client> for Product {
    fn from(client: Client) -> Self {
        match client {
            Client::Elasticsearch(_) => Product::Elasticsearch,
            Client::Kibana(_) => Product::Kibana,
            Client::Logstash(_) => Product::Logstash,
        }
    }
}

impl From<&Client> for Product {
    fn from(client: &Client) -> Self {
        match client {
            Client::Elasticsearch(_) => Product::Elasticsearch,
            Client::Kibana(_) => Product::Kibana,
            Client::Logstash(_) => Product::Logstash,
        }
    }
}

impl std::fmt::Display for Client {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Client::Elasticsearch(_) => write!(f, "elasticsearch"),
            Client::Kibana(_) => write!(f, "kibana"),
            Client::Logstash(_) => write!(f, "logstash"),
        }
    }
}

impl TryFrom<Uri> for Client {
    type Error = eyre::Report;

    fn try_from(uri: Uri) -> Result<Self, Self::Error> {
        match uri {
            Uri::KnownHost(host) => match host.app() {
                Product::Kibana => Ok(Client::Kibana(KibanaClient::try_from(host)?)),
                Product::Elasticsearch => Ok(Client::Elasticsearch(ElasticsearchClient::try_from(host)?)),
                Product::Logstash => Ok(Client::Logstash(LogstashClient::try_from(host)?)),
                _ => Err(eyre!("Unsupported product: {}", host.app())),
            },
            _ => Err(eyre!("Unsupported URI")),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{Json, Router, http::StatusCode, routing::get};
    use serde_json::json;
    use tokio::net::TcpListener;

    #[tokio::test]
    async fn elasticsearch_request_uses_compatibility_response() {
        let app = Router::new().route(
            "/_compatibility",
            get(|| async { (StatusCode::CREATED, Json(json!({ "created": true }))) }),
        );
        let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener");
        let address = listener.local_addr().expect("local address");
        tokio::spawn(async move {
            axum::serve(listener, app).await.expect("serve test application");
        });

        let client = Client::Elasticsearch(
            ElasticsearchBuilder::new(format!("http://{address}").parse().expect("server URL"))
                .build()
                .expect("Elasticsearch client"),
        );
        let response = client
            .request(Method::GET, &HashMap::new(), "_compatibility", None)
            .await
            .expect("Elasticsearch request");

        assert_eq!(response.status(), StatusCode::CREATED);
        assert_eq!(
            response.json::<serde_json::Value>().await.expect("JSON response"),
            json!({ "created": true })
        );
    }
}