1use std::cmp::Ordering;
5use std::fmt;
6use std::fmt::Formatter;
7
8use crate::ServiceState::{Critical, Warning};
9use std::str::FromStr;
10
11#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
12pub enum ServiceState {
14 Ok,
15 Warning,
16 Critical,
17 #[default]
18 Unknown,
19}
20
21impl ServiceState {
22 pub fn exit_code(&self) -> i32 {
24 match self {
25 ServiceState::Ok => 0,
26 ServiceState::Warning => 1,
27 ServiceState::Critical => 2,
28 ServiceState::Unknown => 3,
29 }
30 }
31
32 fn order_number(&self) -> u8 {
35 match self {
36 ServiceState::Ok => 0,
37 ServiceState::Unknown => 1,
38 ServiceState::Warning => 2,
39 ServiceState::Critical => 3,
40 }
41 }
42}
43
44impl PartialOrd for ServiceState {
45 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
46 self.order_number().partial_cmp(&other.order_number())
47 }
48}
49
50impl Ord for ServiceState {
51 fn cmp(&self, other: &Self) -> Ordering {
52 self.order_number().cmp(&other.order_number())
53 }
54}
55
56impl fmt::Display for ServiceState {
57 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
58 let s = match self {
59 ServiceState::Ok => "OK",
60 ServiceState::Warning => "WARNING",
61 ServiceState::Critical => "CRITICAL",
62 ServiceState::Unknown => "UNKNOWN",
63 };
64
65 f.write_str(s)
66 }
67}
68
69#[derive(Debug, thiserror::Error)]
70#[error("expected one of: ok, warning, critical, unknown")]
71pub struct ServiceStateFromStrError;
73
74impl FromStr for ServiceState {
75 type Err = ServiceStateFromStrError;
76
77 fn from_str(s: &str) -> Result<Self, Self::Err> {
78 match s.to_lowercase().as_str() {
79 "ok" => Ok(ServiceState::Ok),
80 "warning" => Ok(ServiceState::Warning),
81 "critical" => Ok(ServiceState::Critical),
82 "unknown" => Ok(ServiceState::Unknown),
83 _ => Err(ServiceStateFromStrError),
84 }
85 }
86}
87
88#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
89pub enum Unit {
92 #[default]
93 None,
94 Seconds,
95 Milliseconds,
96 Microseconds,
97 Percentage,
98 Bytes,
99 Kilobytes,
100 Megabytes,
101 Gigabytes,
102 Terabytes,
103 Counter,
104 Other(UnitString),
105}
106
107impl Unit {
108 fn as_str(&self) -> &str {
109 match self {
110 Unit::None => "",
111 Unit::Seconds => "s",
112 Unit::Milliseconds => "ms",
113 Unit::Microseconds => "us",
114 Unit::Percentage => "%",
115 Unit::Bytes => "B",
116 Unit::Kilobytes => "KB",
117 Unit::Megabytes => "MB",
118 Unit::Gigabytes => "GB",
119 Unit::Terabytes => "TB",
120 Unit::Counter => "c",
121 Unit::Other(s) => &s.0,
122 }
123 }
124}
125
126#[derive(Debug, thiserror::Error)]
127#[non_exhaustive]
128pub enum UnitStringCreateError {
130 #[error("expected string to not include numbers, semicolons or quotes")]
132 InvalidCharacters,
133}
134
135#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)]
136pub struct UnitString(String);
138
139impl UnitString {
140 pub fn new(s: impl Into<String>) -> Result<Self, UnitStringCreateError> {
141 let s = s.into();
142 if ('0'..='9').chain(['"', ';']).any(|c| s.contains(c)) {
143 Err(UnitStringCreateError::InvalidCharacters)
144 } else {
145 Ok(UnitString::new_unchecked(s))
146 }
147 }
148
149 pub fn new_unchecked(s: impl Into<String>) -> Self {
150 UnitString(s.into())
151 }
152}
153
154impl FromStr for UnitString {
155 type Err = UnitStringCreateError;
156
157 fn from_str(s: &str) -> Result<Self, Self::Err> {
158 UnitString::new(s)
159 }
160}
161
162#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
164pub enum TriggerIfValue {
165 Greater,
166 Less,
167}
168
169impl From<&TriggerIfValue> for Ordering {
170 fn from(v: &TriggerIfValue) -> Self {
171 match v {
172 TriggerIfValue::Greater => Ordering::Greater,
173 TriggerIfValue::Less => Ordering::Less,
174 }
175 }
176}
177
178#[derive(Debug, Clone)]
181pub struct Metric<T> {
182 name: String,
183 value: T,
184 unit: Unit,
185 thresholds: Option<(Option<T>, Option<T>, TriggerIfValue)>,
186 min: Option<T>,
187 max: Option<T>,
188 fixed_state: Option<ServiceState>,
189}
190
191impl<T> Metric<T> {
192 pub fn new(name: impl Into<String>, value: T) -> Self {
193 Self {
194 name: name.into(),
195 value,
196 unit: Default::default(),
197 thresholds: Default::default(),
198 min: Default::default(),
199 max: Default::default(),
200 fixed_state: Default::default(),
201 }
202 }
203
204 pub fn with_thresholds(
205 mut self,
206 warning: impl Into<Option<T>>,
207 critical: impl Into<Option<T>>,
208 trigger_if_value: TriggerIfValue,
209 ) -> Self {
210 self.thresholds = Some((warning.into(), critical.into(), trigger_if_value));
211 self
212 }
213
214 pub fn with_minimum(mut self, minimum: T) -> Self {
215 self.min = Some(minimum);
216 self
217 }
218
219 pub fn with_maximum(mut self, maximum: T) -> Self {
220 self.max = Some(maximum);
221 self
222 }
223
224 pub fn with_fixed_state(mut self, state: ServiceState) -> Self {
227 self.fixed_state = Some(state);
228 self
229 }
230
231 pub fn with_unit(mut self, unit: Unit) -> Self {
232 self.unit = unit;
233 self
234 }
235}
236
237#[derive(Debug, Clone)]
239pub struct PerfData<T> {
240 name: String,
241 value: T,
242 unit: Unit,
243 warning: Option<T>,
244 critical: Option<T>,
245 minimum: Option<T>,
246 maximum: Option<T>,
247}
248
249impl<T: ToPerfString> PerfData<T> {
250 pub fn new(name: impl Into<String>, value: T) -> Self {
251 Self {
252 name: name.into(),
253 value,
254 unit: Default::default(),
255 warning: Default::default(),
256 critical: Default::default(),
257 minimum: Default::default(),
258 maximum: Default::default(),
259 }
260 }
261
262 pub fn with_thresholds(mut self, warning: Option<T>, critical: Option<T>) -> Self {
263 self.warning = warning;
264 self.critical = critical;
265 self
266 }
267
268 pub fn with_minimum(mut self, minimum: T) -> Self {
269 self.minimum = Some(minimum);
270 self
271 }
272
273 pub fn with_maximum(mut self, maximum: T) -> Self {
274 self.maximum = Some(maximum);
275 self
276 }
277
278 pub fn with_unit(mut self, unit: Unit) -> Self {
279 self.unit = unit;
280 self
281 }
282}
283
284impl<T: ToPerfString> From<PerfData<T>> for PerfString {
285 fn from(perf_data: PerfData<T>) -> Self {
286 let s = PerfString::new(
287 &perf_data.name,
288 &perf_data.value,
289 perf_data.unit,
290 perf_data.warning.as_ref(),
291 perf_data.critical.as_ref(),
292 perf_data.minimum.as_ref(),
293 perf_data.maximum.as_ref(),
294 );
295 s
296 }
297}
298
299#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
302pub struct PerfString(String);
303
304impl PerfString {
305 pub fn new<T>(
306 name: &str,
307 value: &T,
308 unit: Unit,
309 warning: Option<&T>,
310 critical: Option<&T>,
311 minimum: Option<&T>,
312 maximum: Option<&T>,
313 ) -> Self
314 where
315 T: ToPerfString,
316 {
317 let value = value.to_perf_string();
319 let warning = warning.map_or_else(|| "".to_owned(), |v| v.to_perf_string());
320 let critical = critical.map_or_else(|| "".to_owned(), |v| v.to_perf_string());
321 let minimum = minimum.map_or_else(|| "".to_owned(), |v| v.to_perf_string());
322 let maximum = maximum.map_or_else(|| "".to_owned(), |v| v.to_perf_string());
323 PerfString(format!(
324 "'{}'={}{};{};{};{};{}",
325 name,
326 value,
327 unit.as_str(),
328 warning,
329 critical,
330 minimum,
331 maximum
332 ))
333 }
334}
335
336#[derive(Debug, Clone, PartialEq, Eq)]
338pub struct CheckResult {
339 state: Option<ServiceState>,
340 message: Option<String>,
341 perf_string: Option<PerfString>,
342}
343
344impl CheckResult {
345 pub fn new() -> Self {
347 Self {
348 state: Default::default(),
349 message: Default::default(),
350 perf_string: Default::default(),
351 }
352 }
353
354 pub fn with_state(mut self, state: ServiceState) -> Self {
355 self.state = Some(state);
356 self
357 }
358
359 pub fn with_message(mut self, message: impl Into<String>) -> Self {
360 self.message = Some(message.into());
361 self
362 }
363
364 pub fn with_perf_data(mut self, perf_data: impl Into<PerfString>) -> Self {
367 self.perf_string = Some(perf_data.into());
368 self
369 }
370}
371
372impl Default for CheckResult {
373 fn default() -> Self {
374 Self::new()
375 }
376}
377
378impl<T: PartialOrd + ToPerfString> From<Metric<T>> for CheckResult {
379 fn from(metric: Metric<T>) -> Self {
380 let state = if let Some(state) = metric.fixed_state {
381 Some(state)
382 } else if let Some((warning, critical, trigger)) = &metric.thresholds {
383 let ord: Ordering = trigger.into();
384 let warning_cmp = warning.as_ref().and_then(|w| metric.value.partial_cmp(w));
385 let critical_cmp = critical.as_ref().and_then(|w| metric.value.partial_cmp(w));
386
387 [(critical_cmp, Critical), (warning_cmp, Warning)]
388 .iter()
389 .filter_map(|(cmp, state)| cmp.as_ref().map(|cmp| (cmp, state)))
390 .filter_map(|(&cmp, &state)| {
391 if cmp == ord || cmp == Ordering::Equal {
392 Some(state)
393 } else {
394 None
395 }
396 })
397 .next()
398 } else {
399 None
400 };
401
402 let message = match state {
403 Some(state) if state != ServiceState::Ok => {
404 let (warning, critical, _) = metric.thresholds.as_ref().unwrap();
405 let threshold = match state {
406 ServiceState::Warning => warning.as_ref().unwrap(),
407 ServiceState::Critical => critical.as_ref().unwrap(),
408 _ => unreachable!(),
409 };
410 Some(format!(
411 "metric '{}' is {}: value '{}' has exceeded threshold of '{}'",
412 &metric.name,
413 state,
414 metric.value.to_perf_string(),
415 threshold.to_perf_string(),
416 ))
417 }
418 _ => None,
419 };
420
421 let perf_string = {
422 let (warning, critical) = if let Some((warning, critical, _)) = &metric.thresholds {
423 (warning.as_ref(), critical.as_ref())
424 } else {
425 (None, None)
426 };
427
428 PerfString::new(
429 &metric.name,
430 &metric.value,
431 metric.unit,
432 warning,
433 critical,
434 metric.min.as_ref(),
435 metric.max.as_ref(),
436 )
437 };
438
439 CheckResult {
440 state,
441 message,
442 perf_string: Some(perf_string),
443 }
444 }
445}
446
447pub trait ToPerfString {
449 fn to_perf_string(&self) -> String;
450}
451
452macro_rules! impl_to_perf_string {
453 ($t:ty) => {
454 impl ToPerfString for $t {
455 fn to_perf_string(&self) -> String {
456 self.to_string()
457 }
458 }
459 };
460}
461
462impl_to_perf_string!(usize);
463impl_to_perf_string!(isize);
464impl_to_perf_string!(u8);
465impl_to_perf_string!(u16);
466impl_to_perf_string!(u32);
467impl_to_perf_string!(u64);
468impl_to_perf_string!(u128);
469impl_to_perf_string!(i8);
470impl_to_perf_string!(i16);
471impl_to_perf_string!(i32);
472impl_to_perf_string!(i64);
473impl_to_perf_string!(i128);
474impl_to_perf_string!(f32);
475impl_to_perf_string!(f64);
476
477#[derive(Debug, PartialEq, Eq)]
479pub struct Resource {
480 name: String,
481 results: Vec<CheckResult>,
482 fixed_state: Option<ServiceState>,
483 description: Option<String>,
484}
485
486impl Resource {
487 pub fn new(name: impl Into<String>) -> Self {
489 Self {
490 name: name.into(),
491 results: Default::default(),
492 fixed_state: Default::default(),
493 description: Default::default(),
494 }
495 }
496
497 pub fn with_fixed_state(mut self, state: ServiceState) -> Self {
500 self.fixed_state = Some(state);
501 self
502 }
503
504 pub fn with_result(mut self, result: impl Into<CheckResult>) -> Self {
505 self.push_result(result);
506 self
507 }
508
509 pub fn with_description(mut self, description: impl Into<String>) -> Self {
510 self.set_description(description);
511 self
512 }
513
514 pub fn set_description(&mut self, description: impl Into<String>) {
515 self.description = Some(description.into());
516 }
517
518 pub fn push_result(&mut self, result: impl Into<CheckResult>) {
519 self.results.push(result.into());
520 }
521
522 pub fn nagios_result(self) -> (ServiceState, String) {
524 let (state, messages, perf_string) = {
525 let mut final_state = ServiceState::Ok;
526
527 let mut messages = String::new();
528 let mut perf_string = String::new();
529
530 for result in self.results {
531 if let Some(state) = result.state {
532 if final_state < state {
533 final_state = state;
534 }
535 }
536
537 if let Some(message) = result.message {
538 messages.push_str(message.trim());
539 messages.push('\n');
540 }
541
542 if let Some(s) = result.perf_string {
543 perf_string.push(' ');
544 perf_string.push_str(s.0.trim());
545 }
546 }
547
548 if let Some(state) = self.fixed_state {
549 final_state = state;
550 }
551
552 (final_state, messages, perf_string)
553 };
554
555 let description = {
556 let mut s = String::new();
557 s.push_str(&self.name);
558 s.push_str(" is ");
559 s.push_str(&state.to_string());
560
561 if let Some(description) = self.description {
562 s.push_str(": ");
563 s.push_str(description.trim());
564 }
565 s
566 };
567
568 let mut result = String::new();
569 result.push_str(&description);
570
571 if !messages.is_empty() {
572 result.push_str("\n\n");
573 result.push_str(&messages);
574 }
575
576 if !perf_string.is_empty() {
577 result.push_str("|");
578 result.push_str(perf_string.trim());
579 }
580
581 (state, result)
582 }
583
584 fn print_and_exit(self) -> ! {
587 let (state, s) = self.nagios_result();
588 println!("{}", &s);
589 std::process::exit(state.exit_code());
590 }
591}
592
593pub fn safe_run<E>(
617 f: impl FnOnce() -> Result<Resource, E>,
618 error_state: ServiceState,
619) -> RunResult<E> {
620 match f() {
621 Ok(resource) => RunResult::Ok(resource),
622 Err(err) => RunResult::Err(error_state, err),
623 }
624}
625
626#[derive(Debug)]
628pub enum RunResult<E> {
629 Ok(Resource),
631 Err(ServiceState, E),
633}
634
635impl<E: std::fmt::Display> RunResult<E> {
636 pub fn print_and_exit(self) -> ! {
637 match self {
638 RunResult::Ok(resource) => resource.print_and_exit(),
639 RunResult::Err(state, msg) => {
640 println!("{}: {}", state, msg);
641 std::process::exit(state.exit_code());
642 }
643 }
644 }
645}
646
647#[cfg(test)]
648mod tests {
649 use super::*;
650
651 #[test]
652 fn test_resource_nagios_result() {
653 let (state, s) = Resource::new("foo")
654 .with_description("i am bar")
655 .with_result(
656 CheckResult::new()
657 .with_state(ServiceState::Warning)
658 .with_message("flubblebar"),
659 )
660 .with_result(CheckResult::new().with_state(ServiceState::Critical))
661 .nagios_result();
662
663 assert_eq!(state, ServiceState::Critical);
664 assert!(s.contains("i am bar"));
665 assert!(s.contains("flubblebar"));
666 assert!(s.contains(&ServiceState::Critical.to_string()));
667 }
668
669 #[test]
670 fn test_resource_with_fixed_state() {
671 let (state, _) = Resource::new("foo")
672 .with_fixed_state(ServiceState::Critical)
673 .nagios_result();
674 assert_eq!(state, ServiceState::Critical);
675 }
676
677 #[test]
678 fn test_resource_with_ok_result() {
679 let (state, msg) = Resource::new("foo")
680 .with_result(
681 CheckResult::new()
682 .with_message("test")
683 .with_state(ServiceState::Ok),
684 )
685 .nagios_result();
686
687 assert_eq!(ServiceState::Ok, state);
688 assert!(msg.contains("test"));
689 }
690
691 #[test]
692 fn test_perf_string_new() {
693 let s = PerfString::new("foo", &12, Unit::None, Some(&42), None, None, Some(&60));
694 assert_eq!(&s.0, "'foo'=12;42;;;60")
695 }
696
697 #[test]
698 fn test_metric_into_check_result_complete() {
699 let metric = Metric::new("test", 42)
700 .with_minimum(0)
701 .with_maximum(100)
702 .with_thresholds(40, 50, TriggerIfValue::Greater);
703
704 let result: CheckResult = metric.into();
705 assert_eq!(result.state, Some(ServiceState::Warning));
706
707 let message = result.message.expect("no message set");
708 assert!(message.contains(&ServiceState::Warning.to_string()));
709 assert!(message.contains("test"));
710 assert!(message.contains("threshold"));
711 }
712
713 #[test]
714 fn test_metric_into_check_result_threshold_less() {
715 let result: CheckResult = Metric::new("test", 40)
716 .with_thresholds(50, 30, TriggerIfValue::Less)
717 .into();
718
719 assert_eq!(result.state, Some(ServiceState::Warning));
720 }
721
722 #[test]
723 fn test_metric_into_check_result_threshold_greater() {
724 let result: CheckResult = Metric::new("test", 40)
725 .with_thresholds(30, 50, TriggerIfValue::Greater)
726 .into();
727
728 assert_eq!(result.state, Some(ServiceState::Warning));
729 }
730
731 #[test]
732 fn test_metric_into_check_result_threshold_equal_to_val() {
733 let result: CheckResult = Metric::new("foo", 30)
734 .with_thresholds(30, 40, TriggerIfValue::Greater)
735 .into();
736
737 assert_eq!(result.state, Some(ServiceState::Warning));
738 }
739
740 #[test]
741 fn test_metric_into_check_result_threshold_only_warning() {
742 let result: CheckResult = Metric::new("foo", 30)
743 .with_thresholds(25, None, TriggerIfValue::Greater)
744 .into();
745
746 assert_eq!(result.state, Some(ServiceState::Warning));
747
748 let result: CheckResult = Metric::new("foo", 30)
749 .with_thresholds(35, None, TriggerIfValue::Greater)
750 .into();
751
752 assert_eq!(result.state, None);
753 }
754
755 #[test]
756 fn test_metric_into_check_result_with_unit() {
757 let result: CheckResult = Metric::new("foo", 20)
758 .with_thresholds(25, None, TriggerIfValue::Greater)
759 .with_unit(Unit::Megabytes)
760 .into();
761
762 result.perf_string.unwrap().0.contains("MB");
763
764 assert_eq!(result.state, None);
765 }
766
767 #[derive(Debug, thiserror::Error)]
768 #[error("woops")]
769 struct EmptyError;
770
771 fn do_check(success: bool) -> Result<Resource, EmptyError> {
772 if success {
773 Ok(Resource::new("test"))
774 } else {
775 Err(EmptyError {})
776 }
777 }
778
779 #[test]
780 fn test_safe_run_ok() {
781 let result = safe_run(|| do_check(true), ServiceState::Critical);
782
783 matches!(result, RunResult::Ok(_));
784 }
785
786 #[test]
787 fn test_safe_run_error() {
788 let result = safe_run(|| do_check(false), ServiceState::Critical);
789
790 matches!(result, RunResult::Err(_, _));
791 }
792}