diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/error.rs | 16 | ||||
| -rw-r--r-- | src/lib.rs | 61 | ||||
| -rw-r--r-- | src/main.rs | 42 | ||||
| -rw-r--r-- | src/service.rs | 160 |
4 files changed, 279 insertions, 0 deletions
diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..ef30b97 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,16 @@ +pub type Result<T, E = Error> = std::result::Result<T, E>; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("IO error: {0}")] + IO(#[from] std::io::Error), + + #[error("Config error: {0}")] + Toml(#[from] toml::de::Error), + + #[error("Request error: {0}")] + Reqwest(#[from] reqwest::Error), + + #[error("Invalid HTTP method")] + Method, +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..3437dca --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,61 @@ +use std::ops::Range; + +pub use crate::{ + error::{Error, Result}, + service::{Service, Status}, +}; + +pub mod error; +pub mod service; + +pub fn generate(title: String, mut services: Vec<Service>, template: String) -> String { + let client = reqwest::blocking::Client::new(); + + let [up, down, unknown] = std::thread::scope(|s| { + let mut handles = Vec::new(); + for service in services.iter_mut() { + handles.push(s.spawn(|| service.check(client.clone()))); + } + + handles + .into_iter() + .map(|h| h.join().expect("Joining thread")) + .fold([0, 0, 0], |[up, down, unknown], res| match res { + Ok(true) => [up + 1, down, unknown], + Ok(false) => [up, down + 1, unknown], + Err(_) => [up, down, unknown + 1], + }) + }); + + template + .match_indices("{{services}}") + .zip(template.match_indices("{{end}}")) + .map(|(start, stop)| { + ( + start.0 + start.1.len()..stop.0, + start.0..stop.0 + stop.1.len(), + ) + }) + .collect::<Vec<_>>() + .into_iter() + .fold(template, |mut template, (Range { start, end }, outer)| { + let replace_with = services + .iter() + .map(|service| { + template[start..end] + .replace("{name}", &service.name) + .replace("{title}", &service.kind.to_string()) + .replace("{state}", &service.state.to_string()) + .replace("{level}", &service.state.as_level()) + }) + .collect::<String>(); + + template.replace_range(outer, &replace_with); + template + }) + .replace("{title}", &title) + .replace("{status}", if down > 0 { "error" } else { "ok" }) + .replace("{up}", &up.to_string()) + .replace("{down}", &down.to_string()) + .replace("{unknown}", &unknown.to_string()) +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..2ff9fd3 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,42 @@ +use std::{fs::File, io::Write, path::PathBuf}; + +use statsrv::Service; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub title: String, + pub template_path: PathBuf, + pub output_dir: Option<PathBuf>, + pub address: Option<String>, + pub services: Vec<Service>, +} + +fn main() -> Result<(), main_error::MainError> { + let mut args = std::env::args().skip(1); + + let config_path = args + .next() + .unwrap_or_else(|| "/etc/statsrv.toml".to_string()); + let config_file = File::open(config_path)?; + let config_toml = std::io::read_to_string(config_file)?; + let Config { + title, + template_path: template, + output_dir, + address: _, + services, + } = toml::from_str(&config_toml)?; + + let template_file = File::open(template)?; + let template = std::io::read_to_string(template_file)?; + let status_page = statsrv::generate(title, services, template); + + if let Some(output_dir) = output_dir { + std::fs::create_dir_all(&output_dir)?; + let mut html_writer = File::create(output_dir.join("index.html"))?; + + html_writer.write_all(status_page.as_bytes())?; + } + + Ok(()) +} diff --git a/src/service.rs b/src/service.rs new file mode 100644 index 0000000..c5eb0d7 --- /dev/null +++ b/src/service.rs @@ -0,0 +1,160 @@ +use std::{fmt::Display, process::Command}; + +use reqwest::blocking::Client; +use serde::Deserialize; + +use crate::Error; + +#[derive(Debug, Clone, Deserialize)] +pub struct Service { + pub name: String, + #[serde(flatten)] + pub kind: Kind, + #[serde(skip)] + pub state: State, +} + +impl Service { + pub fn check(&mut self, client: Client) -> Result<bool, Error> { + self.state = self.kind.get_state(client)?; + Ok(self.state.is_operational()) + } +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum Kind { + Tcp { + address: String, + }, + Http { + url: String, + #[serde(default = "Kind::default_method")] + method: String, + #[serde(default = "Kind::default_code")] + status_code: u16, + }, + Systemd { + service: String, + }, +} + +impl Kind { + fn default_method() -> String { + "GET".to_string() + } + + fn default_code() -> u16 { + 200 + } + + pub fn get_state(&self, client: Client) -> Result<State, Error> { + let state = match self { + Kind::Tcp { address } => { + if std::net::TcpStream::connect(address).is_ok() { + State::Operational + } else { + State::Down("Unreachable".to_string()) + } + } + Kind::Http { + method, + url, + status_code, + } => { + match client + .request(method.parse().map_err(|_| Error::Method)?, url) + .send()? + .status() + { + s if s.as_u16() == *status_code => State::Operational, + s => State::Down(s.to_string()), + } + } + Kind::Systemd { service } => { + let output = Command::new("systemctl") + .arg("is-active") + .arg(service) + .output()?; + + if output.status.success() { + State::Operational + } else { + State::Down(String::from_utf8_lossy(&output.stdout).to_string()) + } + } + }; + + Ok(state) + } +} + +impl Display for Kind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Kind::Tcp { address } => write!(f, "tcp://{address}"), + Kind::Http { method, url, .. } => write!(f, "{method} {url}"), + Kind::Systemd { service } => write!(f, "{service}"), + } + } +} + +#[derive(Debug, Clone, Default)] +pub struct Status { + pub info: String, + pub state: State, +} + +#[derive(Debug, Clone, Default)] +pub enum State { + #[default] + Unknown, + Operational, + Down(String), +} + +impl State { + /// Returns `true` if this is a `Unknown` variant. + pub fn is_unknown(&self) -> bool { + matches!(self, Self::Unknown) + } + + /// Returns `true` if this is a `Operational` variant. + pub fn is_operational(&self) -> bool { + matches!(self, Self::Operational) + } + + /// Returns `true` if this is a `Down` variant. + pub fn is_down(&self) -> bool { + matches!(self, Self::Down(_)) + } + + /// Converts the `State` into an `Option` containing `String` description if the `State` was + /// `Down` and `None` otherwise. + pub fn down_value(self) -> Option<String> { + match self { + State::Unknown => None, + State::Operational => None, + State::Down(s) => Some(s), + } + } + + pub fn as_level(&self) -> String { + match self { + State::Unknown => "warning", + State::Operational => "ok", + State::Down(_) => "error", + } + .to_string() + } +} + +impl Display for State { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + State::Unknown => write!(f, "Unknown"), + State::Operational => write!(f, "Operational"), + State::Down(s) => write!(f, "{s}"), + } + } +} |