diff --git a/CHANGELOG.md b/CHANGELOG.md index ddbea5f..a52d43d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added + +- `RequestCode::ProcessList` to request agent to collect host's process list. + ### Changed - Fix clippy warnings. diff --git a/src/lib.rs b/src/lib.rs index 181782d..8374b83 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,7 +5,7 @@ pub mod request; mod test; use num_enum::{FromPrimitive, IntoPrimitive}; -pub use request::{Configuration, ResourceUsage}; +pub use request::{Configuration, Process, ResourceUsage}; /// Numeric representation of the message types. #[derive(Clone, Copy, Debug, Eq, FromPrimitive, IntoPrimitive, PartialEq)] @@ -68,6 +68,9 @@ pub enum RequestCode { /// Update the list of trusted domains TrustedDomainList = 0, + /// Collect process list + ProcessList = 19, + /// Unknown request #[num_enum(default)] Unknown = u32::MAX, diff --git a/src/request.rs b/src/request.rs index 4d522ba..efddb35 100644 --- a/src/request.rs +++ b/src/request.rs @@ -32,6 +32,15 @@ pub struct ResourceUsage { pub used_disk_space: u64, } +#[derive(Debug, Deserialize, Serialize)] +pub struct Process { + pub user: String, + pub cpu_usage: f32, + pub mem_usage: f64, + pub start_time: i64, + pub command: String, +} + #[derive(Debug, Deserialize, Serialize)] pub struct Configuration { pub hog_sources: Option>, @@ -61,6 +70,7 @@ pub enum HandlerError { } /// A request handler that can handle a request to an agent. +#[allow(clippy::diverging_sub_expression)] #[async_trait] pub trait Handler: Send { async fn dns_start(&mut self) -> Result<(), String> { @@ -137,6 +147,9 @@ pub trait Handler: Send { async fn trusted_user_agent_list(&mut self, _list: &[&str]) -> Result<(), String> { return Err("not supported".to_string()); } + async fn process_list(&mut self) -> Result, String> { + return Err("not supported".to_string()); + } } /// Handles requests to an agent. @@ -278,6 +291,9 @@ pub async fn handle( let result = handler.set_config(conf).await; send_response(send, &mut buf, result).await?; } + RequestCode::ProcessList => { + send_response(send, &mut buf, handler.process_list().await).await?; + } RequestCode::Unknown => { let err_msg = format!("unknown request code: {code}"); message::send_err(send, &mut buf, err_msg).await?; @@ -304,7 +320,7 @@ mod tests { frame, message, request::HostNetworkGroup, test::{channel, TOKEN}, - RequestCode, + Process, RequestCode, }; use async_trait::async_trait; use ipnet::IpNet; @@ -326,6 +342,7 @@ mod tests { allow_list: usize, block_list: usize, trusted_user_agents: usize, + process_list_cnt: usize, } #[async_trait] @@ -374,6 +391,11 @@ mod tests { self.trusted_user_agents = list.len(); Ok(()) } + + async fn process_list(&mut self) -> Result, String> { + self.process_list_cnt += 1; + Ok(Vec::new()) + } } let mut handler = TestHandler::default(); @@ -506,6 +528,16 @@ mod tests { .await; assert!(res.is_ok()); + let mut buf = Vec::new(); + let res = message::send_request( + &mut channel.client.send, + &mut buf, + RequestCode::ProcessList, + (), + ) + .await; + assert!(res.is_ok()); + channel.client.send.finish().await.unwrap(); assert_eq!(handler.reboot_count, 0); @@ -515,6 +547,7 @@ mod tests { assert_eq!(handler.allow_list, 0); assert_eq!(handler.block_list, 0); assert_eq!(handler.trusted_user_agents, 0); + assert_eq!(handler.process_list_cnt, 0); let res = super::handle( &mut handler, &mut channel.server.send, @@ -529,6 +562,7 @@ mod tests { assert_eq!(handler.allow_list, 2); assert_eq!(handler.block_list, 1); assert_eq!(handler.trusted_user_agents, 4); + assert_eq!(handler.process_list_cnt, 1); frame::recv_raw(&mut channel.client.recv, &mut buf) .await