icinga_mock/
lib.rs

1use std::net::Ipv4Addr;
2
3use hyper::{header, Method};
4use icinga_client::{
5    client::{Client, RequestBuilder},
6    types::{CheckedObject, IcingaObjectResults},
7};
8use serde::de::DeserializeOwned;
9use serde_json::Value;
10use tokio::runtime::Runtime;
11
12#[derive(Debug, Clone)]
13pub struct IcingaOptions {
14    pub address: Ipv4Addr,
15    pub port: u16,
16    pub credentials: Option<(String, String)>,
17}
18
19impl IcingaOptions {
20    pub fn new() -> Self {
21        Self::from_address("127.0.0.1".parse().unwrap())
22    }
23    pub fn from_address(address: Ipv4Addr) -> Self {
24        Self::from_address_and_port(address, 0)
25    }
26
27    pub fn from_address_and_port(address: Ipv4Addr, port: u16) -> Self {
28        IcingaOptions {
29            address,
30            port,
31            credentials: None,
32        }
33    }
34
35    pub fn with_credentials<S: ToString, T: ToString>(
36        &mut self,
37        user: S,
38        password: T,
39    ) -> &mut Self {
40        self.credentials = Some((user.to_string(), password.to_string()));
41        self
42    }
43}
44
45pub struct Session(Runtime, Client);
46
47impl Session {
48    pub fn new() -> Self {
49        let options = IcingaOptions::new();
50        Self::from_options(&options)
51    }
52
53    pub fn from_options(options: &IcingaOptions) -> Self {
54        let rt = Runtime::new().unwrap();
55        let client = start_server(options, &rt);
56        Self(rt, client)
57    }
58
59    pub fn get_objects<T: DeserializeOwned + CheckedObject>(&self) -> Vec<T> {
60        self.get::<IcingaObjectResults<T>>(T::PATH)
61            .results
62            .into_iter()
63            .map(|r| r.attrs)
64            .collect::<Vec<_>>()
65    }
66
67    pub fn post<T: DeserializeOwned>(&self, path: &str, query: &[(&str, &str)], body: Value) -> T {
68        self.request(Method::POST, path, |r| {
69            r.query(query)
70                .header(header::CONTENT_TYPE, "application/json")
71                .body(body.to_string())
72        })
73    }
74
75    pub fn get<T: DeserializeOwned>(&self, path: &str) -> T {
76        self.request(Method::GET, path, |r| r)
77    }
78
79    pub fn request<T: DeserializeOwned>(
80        &self,
81        method: Method,
82        path: &str,
83        make_request: impl FnOnce(RequestBuilder) -> RequestBuilder,
84    ) -> T {
85        self.with_client(|client| {
86            let request = make_request(client.request(method, path).unwrap());
87            client.send_request(request).unwrap()
88        })
89    }
90
91    pub fn with_client<R>(&self, f: impl FnOnce(&Client) -> R) -> R {
92        let Self(_rt, client) = self;
93        f(&client)
94    }
95}
96
97fn start_server(options: &IcingaOptions, rt: &Runtime) -> Client {
98    let options = options.clone();
99    let h = rt.spawn(async move {
100        let server = server::server(&options).await.unwrap();
101        let addr = server.local_addr();
102        tokio::spawn(server);
103        addr
104    });
105    let addr = rt.block_on(h).unwrap();
106    Client::new(format!("http://{}", addr).parse().unwrap(), None, None).unwrap()
107}
108
109pub mod server {
110    use axum::{
111        error_handling::HandleErrorLayer,
112        extract::{Extension, Query},
113        routing::{get, post, IntoMakeService},
114        AddExtensionLayer, BoxError, Json, Router,
115    };
116    use chrono::Utc;
117    use hyper::{server::conn::AddrIncoming, StatusCode};
118    use icinga_client::types::{
119        AckKind, AckState, Acknowledgement, CheckedObject, Host, HostState, IcingaObjectResult,
120        IcingaObjectResults, Service, ServiceState, Timestamp,
121    };
122    use serde::{Deserialize, Serialize};
123    use serde_json::json;
124    use std::{
125        net::SocketAddr,
126        sync::{Arc, Mutex},
127    };
128    use tower::ServiceBuilder;
129    use tower_http::auth::RequireAuthorizationLayer;
130
131    use crate::IcingaOptions;
132
133    pub type Server = axum::Server<AddrIncoming, IntoMakeService<Router>>;
134
135    #[derive(Debug, Clone)]
136    struct State(Arc<Mutex<(Vec<Host>, Vec<Service>)>>);
137
138    impl State {
139        fn with_hosts<R>(&self, f: impl FnOnce(&mut Vec<Host>) -> R) -> R {
140            f(&mut self.0.lock().unwrap().0)
141        }
142
143        fn with_services<R>(&self, f: impl FnOnce(&mut Vec<Service>) -> R) -> R {
144            f(&mut self.0.lock().unwrap().1)
145        }
146    }
147
148    #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
149    #[serde(untagged)]
150    enum Selector {
151        Host(HostSelector),
152        Service(ServiceSelector),
153    }
154
155    #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
156    struct HostSelector {
157        host: String,
158    }
159
160    #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
161    struct ServiceSelector {
162        service: String,
163    }
164
165    #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
166    struct AckBody {
167        expiry: Option<f64>,
168    }
169
170    pub async fn server(opts: &IcingaOptions) -> hyper::Result<Server> {
171        // build our application with a route
172        let addr = SocketAddr::from((opts.address, opts.port));
173        let initial_hosts = hosts();
174        let initial_services = services();
175        let state = State(Arc::new(Mutex::new((initial_hosts, initial_services))));
176        let services = ServiceBuilder::new()
177            .layer(AddExtensionLayer::new(state))
178            .layer(HandleErrorLayer::new(handle_unexpected_error))
179            .option_layer(opts.credentials.as_ref().map(|(username, password)| {
180                RequireAuthorizationLayer::basic(&username, &password)
181            }));
182        let app = Router::new()
183            .route("/", get(|| async { "Welcome to icinga-mock" }))
184            .route(
185                "/v1/objects/hosts",
186                get(|Extension(state): Extension<State>| async move {
187                    Json(state.with_hosts(|hosts| results(hosts.clone())))
188                }),
189            )
190            .route(
191                "/v1/objects/services",
192                get(|Extension(state): Extension<State>| async move {
193                    Json(state.with_services(|services| results(services.clone())))
194                }),
195            )
196            .route(
197                "/v1/actions/acknowledge-problem",
198                post(
199                    |Query(selector): Query<Selector>,
200                     Json(body): Json<AckBody>,
201                     Extension(state): Extension<State>| async move {
202                        fn ack<T: CheckedObject>(
203                            objects: &mut Vec<T>,
204                            selected_name: &str,
205                            ack_state: &AckState,
206                        ) {
207                            for obj in objects {
208                                if obj.name() == selected_name {
209                                    obj.set_acknowledgement(Acknowledgement {
210                                        state: ack_state.clone(),
211                                        last_change: None,
212                                    });
213                                    obj.set_handled(true);
214                                }
215                            }
216                        }
217                        let ack_state = body
218                            .expiry
219                            .map(ack_with_expiry)
220                            .unwrap_or(ack_without_expiry());
221                        match selector {
222                            Selector::Host(HostSelector { host: name }) => {
223                                state.with_hosts(|hosts| ack(hosts, &name, &ack_state))
224                            }
225                            Selector::Service(ServiceSelector { service: name }) => {
226                                state.with_services(|services| ack(services, &name, &ack_state))
227                            }
228                        }
229                        Json(json!({
230                           "results": [
231                               {
232                                   "code": 200,
233                                   "status": "Successfully acknowledged problem ..."
234                               }
235                           ]
236                        }))
237                    },
238                ),
239            )
240            .layer(services);
241        Ok(axum::Server::try_bind(&addr)?.serve(app.into_make_service()))
242    }
243
244    async fn handle_unexpected_error(err: BoxError) -> (StatusCode, String) {
245        (
246            StatusCode::INTERNAL_SERVER_ERROR,
247            format!("Unexpected error in middleware: {}", err),
248        )
249    }
250
251    fn results<T: Serialize>(os: Vec<T>) -> IcingaObjectResults<T> {
252        IcingaObjectResults {
253            results: os
254                .into_iter()
255                .map(|o| IcingaObjectResult { attrs: o })
256                .collect(),
257        }
258    }
259
260    fn hosts() -> Vec<Host> {
261        vec![
262            host(1, HostState::UP, ack_none()),
263            host(2, HostState::DOWN, ack_none()),
264            host(3, HostState::DOWN, ack_with_default_expiry()),
265            host(4, HostState::DOWN, ack_without_expiry()),
266        ]
267    }
268
269    fn services() -> Vec<Service> {
270        vec![
271            service(1, 1, ServiceState::OK, ack_none()),
272            service(2, 1, ServiceState::WARNING, ack_with_default_expiry()),
273            service(3, 1, ServiceState::CRITICAL, ack_without_expiry()),
274            service(4, 1, ServiceState::UNKNOWN, ack_none()),
275        ]
276    }
277
278    fn host(i: usize, state: HostState, ack: AckState) -> Host {
279        let name = format!("host{}", i);
280        let handled = ack != AckState::None;
281        let now = Utc::now().timestamp();
282        let acknowledgement = Acknowledgement {
283            state: ack,
284            last_change: None,
285        };
286        Host {
287            name: name.clone(),
288            display_name: name.clone(),
289            address: format!("127.0.0.{}", i),
290            address6: format!("::{}", i),
291            state,
292            last_state: state,
293            last_hard_state: state,
294            last_state_change: ts(now),
295            last_state_up: ts(now - 5),
296            last_state_down: ts(now - 60),
297            next_check: ts(now + 60),
298            last_check_result: None,
299            acknowledgement,
300            handled,
301        }
302    }
303
304    fn service(i: usize, host_id: usize, state: ServiceState, ack: AckState) -> Service {
305        let name = format!("service{}", i);
306        let host_name = format!("host{}", host_id);
307        let handled = ack != AckState::None;
308        let now = Utc::now().timestamp();
309        let acknowledgement = Acknowledgement {
310            state: ack,
311            last_change: None,
312        };
313        Service {
314            name: name.clone(),
315            display_name: name.clone(),
316            host_name,
317            state,
318            last_state: state,
319            last_hard_state: state,
320            last_state_change: ts(now),
321            last_state_ok: ts(now - 5),
322            last_state_warning: ts(now - 60),
323            last_state_critical: ts(now - 120),
324            last_state_unknown: ts(now - 180),
325            next_check: ts(now + 60),
326            last_check_result: None,
327            last_reachable: true,
328            acknowledgement,
329            handled,
330        }
331    }
332
333    fn ack_none() -> AckState {
334        AckState::None
335    }
336
337    fn ack_with_expiry(t: f64) -> AckState {
338        AckState::Acknowledged {
339            kind: AckKind::Normal,
340            expiry: Some(Timestamp::new(t)),
341        }
342    }
343
344    fn ack_with_default_expiry() -> AckState {
345        ack_with_expiry((Utc::now().timestamp() + 60) as f64)
346    }
347
348    fn ack_without_expiry() -> AckState {
349        AckState::Acknowledged {
350            kind: AckKind::Normal,
351            expiry: None,
352        }
353    }
354
355    fn ts(t: i64) -> Timestamp {
356        Timestamp::new(t as f64)
357    }
358
359    #[cfg(test)]
360    mod test {
361        use super::*;
362
363        #[test]
364        fn test_de_selector_host() {
365            let v = json!({"host": "host2"});
366
367            assert_eq!(
368                serde_json::from_value::<Selector>(v).unwrap(),
369                Selector::Host(HostSelector {
370                    host: "host2".to_string()
371                })
372            );
373        }
374
375        #[test]
376        fn test_de_selector_service() {
377            let v = json!({"service": "service2"});
378
379            assert_eq!(
380                serde_json::from_value::<Selector>(v).unwrap(),
381                Selector::Service(ServiceSelector {
382                    service: "service2".to_string()
383                })
384            );
385        }
386    }
387}