esdiag 0.16.4

Elastic Stack diagnostic collector and processor
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

// 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.

use crate::data::{Auth, KnownHost};
use base64::{Engine, engine::general_purpose::STANDARD};
use elasticsearch::{
    cert::CertificateValidation,
    http::{
        headers,
        transport::{SingleNodeConnectionPool, TransportBuilder},
    },
};
use eyre::Result;
use url::Url;

pub use elasticsearch::Elasticsearch as ElasticsearchClient;

pub struct ElasticsearchBuilder {
    cert_validation: CertificateValidation,
    connection_pool: SingleNodeConnectionPool,
    headers: headers::HeaderMap,
}

/// A builder for the official Elasticsearch client
impl ElasticsearchBuilder {
    pub fn new(url: Url) -> Self {
        let mut headers = headers::HeaderMap::new();
        headers.append(headers::ACCEPT_ENCODING, "gzip".parse().unwrap());

        Self {
            cert_validation: CertificateValidation::Default,
            connection_pool: SingleNodeConnectionPool::new(url),
            headers,
        }
    }

    pub fn apikey(self, apikey: String) -> Self {
        let mut headers = self.headers;
        headers.append(
            headers::AUTHORIZATION,
            format!("ApiKey {}", apikey).parse().expect("Invalid API key"),
        );
        Self { headers, ..self }
    }

    pub fn auth(self, auth: Auth) -> Self {
        tracing::debug!("Setting client auth to {}", auth);
        match auth {
            Auth::Apikey(apikey) => self.apikey(apikey),
            Auth::Basic(username, password) => self.basic_auth(username, password),
            Auth::None => self,
        }
    }

    pub fn basic_auth(self, username: String, password: String) -> Self {
        let mut headers = self.headers;
        headers.append(
            headers::AUTHORIZATION,
            headers::HeaderValue::from_str(&format!(
                "Basic {}",
                STANDARD.encode(format!("{}:{}", username, password))
            ))
            .expect("Invalid basic auth"),
        );
        Self { headers, ..self }
    }

    pub fn build(self) -> Result<ElasticsearchClient> {
        let transport = TransportBuilder::new(self.connection_pool)
            .headers(self.headers)
            .cert_validation(self.cert_validation)
            .request_body_compression(true)
            .build()?;
        Ok(ElasticsearchClient::new(transport))
    }

    pub fn insecure(self, ignore_certs: bool) -> Self {
        let cert_validation = match ignore_certs {
            true => CertificateValidation::None,
            false => CertificateValidation::Default,
        };
        Self {
            cert_validation,
            ..self
        }
    }
}

impl TryFrom<KnownHost> for ElasticsearchClient {
    type Error = eyre::Report;

    fn try_from(host: KnownHost) -> Result<ElasticsearchClient> {
        let url = host.get_url()?;
        let ignore_certs = host.accept_invalid_certs();
        let auth = host.get_auth()?;
        let client = ElasticsearchBuilder::new(url)
            .auth(auth)
            .insecure(ignore_certs)
            .build()?;
        Ok(client)
    }
}