diff --git a/src/collector.rs b/src/collector.rs index 68d9e89..f5f2361 100644 --- a/src/collector.rs +++ b/src/collector.rs @@ -1,5 +1,6 @@ //! Module to define behavior of sys info collection. use serde_json::{Value, json}; +use std::process::Command; use sysinfo::{Disks, Networks, System}; /// Function to get raw uptime. @@ -11,20 +12,26 @@ fn cpu_freq_raw(sys: &System) -> String { let cpu_freq = sys.cpus().first().map(|cpu| cpu.frequency()).unwrap_or(0); cpu_freq.to_string() } -/// Function to collect system metrics as single json object. -pub fn get_sysinfo(sys: &System) -> Value { - let timestamp = std::time::SystemTime::now() + +/// Function to get timestamp +pub fn get_timestamp() -> u64 { + std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() - .as_secs(); + .as_secs() +} - let hostname = System::host_name() +/// Function to get hostname +pub fn get_hostname(sys: &System) -> String { + System::host_name() .unwrap_or_else(|| "unknown".to_string()) - .replace('"', "\\\""); - + .replace('"', "\\\"") +} +/// Function to collect system metrics as single json object. +pub fn get_sysinfo(sys: &System) -> Value { json!({ - "timestamp": timestamp, - "hostname": hostname, + "timestamp": get_timestamp(), + "hostname": get_hostname(sys), "uptime": uptime_raw(sys), "cpu_freq_mhz": cpu_freq_raw(sys), "disk_usage": get_disk_usage(), @@ -79,3 +86,35 @@ pub fn get_disk_usage() -> Vec { }) .collect() } + +/// Function to get status of specific systemd services by name +pub fn get_service_status(services: &[String]) -> Vec { + let mut results = Vec::new(); + + for service in services { + let status = get_service_active_status(&service); + + let service_status = json!({ + "service_name": service, + "status": status + }); + + results.push(service_status); + } + + results +} + +/// Get the active status of a service +fn get_service_active_status(service: &str) -> String { + match Command::new("systemctl") + .args(&["is-active", service]) + .output() + { + Ok(output) => str::from_utf8(&output.stdout) + .unwrap_or("unknown") + .trim() + .to_string(), + Err(_) => "error".to_string(), + } +} diff --git a/src/main.rs b/src/main.rs index c65ce0d..c47b836 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ -//! Main module for tinycollectd. +//! Main module for tinyd. use clap::{Parser, ValueEnum}; +use serde_json::{Value, json}; use std::net::{Ipv4Addr, SocketAddrV4}; use std::time::Duration; use sysinfo::System; @@ -10,14 +11,17 @@ struct Cli { /// destination for metrics (e.g. 127.0.0.1:1555) #[arg(long, default_value_t = SocketAddrV4::new(Ipv4Addr::new(127,0,0,1), 1555))] destination: SocketAddrV4, - /// metrics tinycollectd would collect + /// metrics tinyd would collect #[arg(long, value_enum, value_delimiter = ',', default_value = "All")] metrics: Vec, + /// list of services to pull status + #[arg(long)] + services: Vec, /// interval for data to be collected in seconds. #[arg(long, default_value = "10")] collection_interval: u64, } -#[derive(ValueEnum, Clone)] +#[derive(ValueEnum, Clone, Debug, PartialEq)] enum MetricType { All, DiskUsage, @@ -28,7 +32,7 @@ enum MetricType { } /// Function to add hostname, timestamp, and other metadata to individual metrics -/// Entrypoint for tinycollectd async runtime. +/// Entrypoint for tinyd async runtime. #[tokio::main] async fn main() -> Result<(), Box> { @@ -41,7 +45,62 @@ async fn main() -> Result<(), Box> { loop { sys.refresh_all(); // refresh once on every collection attempt - let bytes = serde_json::to_vec(&collector::get_sysinfo(&sys)).unwrap(); + + let metrics_value = if cli.metrics.contains(&MetricType::All) { + collector::get_sysinfo(&sys) + } else { + let mut combined_object = serde_json::Map::new(); + let mut combined_arrays = Vec::new(); + + if cli.metrics.contains(&MetricType::Service) { + let service_data = collector::get_service_status(&cli.services); + combined_arrays.extend(service_data); + } + + if cli.metrics.contains(&MetricType::DiskUsage) { + let disk_data = collector::get_disk_usage(); + combined_arrays.extend(disk_data); + } + + if cli.metrics.contains(&MetricType::Network) { + let network_data = collector::get_if_data(); + combined_arrays.extend(network_data); + } + + if cli.metrics.contains(&MetricType::Cpufreq) { + let cpu_data = collector::cpu_freq_json(&sys); + if let Value::Object(map) = cpu_data { + combined_object.extend(map); + } + } + + if cli.metrics.contains(&MetricType::Uptime) { + let uptime_data = collector::uptime_json(&sys); + if let Value::Object(map) = uptime_data { + combined_object.extend(map); + } + } + + // Combine single values and arrays + if !combined_object.is_empty() && !combined_arrays.is_empty() { + combined_object.insert("array_data".to_string(), Value::Array(combined_arrays)); + Value::Object(combined_object) + } else if !combined_object.is_empty() { + Value::Object(combined_object) + } else if !combined_arrays.is_empty() { + Value::Array(combined_arrays) + } else { + json!({}) + } + }; + + let combined = json!({ + "timestamp": collector::get_timestamp(), + "hostname": collector::get_hostname(&sys), + "metrics": metrics_value, + }); + + let bytes = serde_json::to_vec(&combined).unwrap(); // Send UDP packet if let Err(e) = socket.send_to(&bytes, cli.destination).await { @@ -50,10 +109,10 @@ async fn main() -> Result<(), Box> { println!( "Sent metrics to {} ({} bytes)", cli.destination, - &bytes.len() + bytes.len() ); } - tokio::time::sleep(Duration::from_secs(10)).await; + tokio::time::sleep(Duration::from_secs(cli.collection_interval)).await; } }