📝 Add docstrings to agent/bug-fixes-01

agent/crates/agent-bin/src/main.rs

@@ -39,6 +39,16 @@ struct Args {
     enroll: bool,
 }
 
+/// Updates or inserts the `node_id` field in the `[agent]` section of the TOML configuration file.
+///
+/// # Examples
+///
+/// ```
+/// use std::path::Path;
+/// use uuid::Uuid;
+/// let node_id = Uuid::new_v4();
+/// // save_node_id_to_config(Path::new("config.toml"), node_id).unwrap();
+/// ```
 fn save_node_id_to_config(path: &Path, node_id: Uuid) -> anyhow::Result<()> {
     let content = std::fs::read_to_string(path)?;
     let mut lines: Vec<String> = content.lines().map(String::from).collect();
@@ -68,6 +78,16 @@ fn save_node_id_to_config(path: &Path, node_id: Uuid) -> anyhow::Result<()> {
     Ok(())
 }
 
+/// Extracts IP address and port from an endpoint string.
+///
+/// Invalid IP addresses default to 127.0.0.1, and missing or invalid ports default to 50051.
+///
+/// # Examples
+///
+/// ```
+/// let (ip, port) = parse_endpoint("http://192.168.1.1:8080");
+/// assert_eq!(port, 8080);
+/// ```
 fn parse_endpoint(endpoint: &str) -> (std::net::IpAddr, u16) {
     let clean = endpoint
         .trim_start_matches("http://")
@@ -83,6 +103,18 @@ fn parse_endpoint(endpoint: &str) -> (std::net::IpAddr, u16) {
     (ip, port)
 }
 
+/// Initializes and runs the Aigis-Zero agent until interrupted.
+///
+/// Loads configuration, validates the environment, enrolls with the fleet server,
+/// and establishes background tasks for event collection, heartbeats, and command handling.
+/// The agent blocks until Ctrl-C is received.
+///
+/// If `--check` is set, validates the configuration and environment, then exits.
+///
+/// # Errors
+///
+/// Returns an error if configuration is invalid, the agent cannot connect to the fleet,
+/// or enrollment fails.
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     dotenvy::dotenv().ok();

agent/crates/agent-core/src/command_handler.rs

@@ -12,6 +12,27 @@ pub struct CommandHandler {
 }
 
 impl CommandHandler {
+    /// Processes a server command and returns a JSON status response.
+    ///
+    /// Handles different command types with specific actions:
+    /// - `Isolate`: Isolates or de-isolates the process based on the flag.
+    /// - `ConfigUpdate`: Acknowledges the configuration update.
+    /// - `Ack`: Acknowledges the message.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the command is missing from the message or if an isolation operation fails.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # async fn example() {
+    /// let handler = CommandHandler { /* ... */ };
+    /// let cmd = ServerCommand { command: Some(Command::Ack(())) };
+    /// let result = handler.handle(cmd).await;
+    /// assert!(result.is_ok());
+    /// # }
+    /// ```
     pub async fn handle(&self, msg: ServerCommand) -> Result<Value, String> {
         let command = msg.command.ok_or("missing command")?;
 

agent/crates/agent-core/src/lib.rs

@@ -41,12 +41,17 @@ pub struct AgentCore {
 }
 
 impl AgentCore {
-    /// Start all background tasks and block until the shutdown token fires.
+    /// Executes the agent's background tasks until shutdown is triggered.
     ///
-    /// # Parameters
-    /// - `agent_uuid`: The node UUID assigned during enrollment.  Passed into
-    ///   `OsqueryCollector::start` so that every `OsqueryResult` carries the
-    ///   correct identity before it is serialised into an `AgentEvent`.
+    /// Spawns an OSQuery polling task and a command listener task. Both tasks
+    /// are given up to 5 seconds to gracefully exit after shutdown is signaled.
+    ///
+    /// # Examples
+    ///
+    /// ```ignore
+    /// let core = AgentCore { /* ... */ };
+    /// core.run("my-agent-uuid").await?;
+    /// ```
     pub async fn run(&self, agent_uuid: &str) -> Result<()> {
         let shutdown = self.shutdown.clone();
 

agent/crates/agent-core/src/orchestrator.rs

@@ -5,6 +5,19 @@ use event_buffer::EventBuffer;
 use osquery_client::types::OsqueryResult;
 use tokio::sync::mpsc;
 
+/// Runs the agent orchestrator, loading configuration and managing osquery event collection.
+///
+/// Reads configuration from the `EDR_AGENT_CONFIG` environment variable
+/// (defaulting to `"agent.toml"`), initializes logging, starts collecting events
+/// from osquery, and attempts to enroll with the fleet server. Processes collected
+/// events until interrupted by Ctrl-C. Fleet enrollment failures do not prevent
+/// the agent from continuing to operate offline.
+///
+/// # Examples
+///
+/// ```no_run
+/// run().await?;
+/// ```
 pub async fn run() -> Result<()> {
     let config_path =
         std::env::var("EDR_AGENT_CONFIG").unwrap_or_else(|_| "agent.toml".to_string());
@@ -157,13 +170,29 @@ fn encode_result(result: &OsqueryResult) -> String {
     serde_json::to_string(&event).unwrap_or_default()
 }
 
+/// Retrieves the system hostname, or a fallback value if unavailable.
+///
+/// # Examples
+///
+/// ```
+/// let hostname = hostname_or_default();
+/// assert!(!hostname.is_empty());
+/// ```
 fn hostname_or_default() -> String {
     hostname::get()
         .ok()
         .and_then(|h| h.into_string().ok())
         .unwrap_or_else(|| "unknown-host".to_string())
 }
 
+/// Retrieves the system machine ID or a fallback value if unavailable.
+///
+/// # Examples
+///
+/// ```
+/// let machine_id = read_machine_id();
+/// assert!(!machine_id.is_empty());
+/// ```
 pub fn read_machine_id() -> String {
     if let Ok(id) = std::fs::read_to_string("/etc/machine-id") {
         let trimmed = id.trim();
@@ -180,6 +209,16 @@ pub fn read_machine_id() -> String {
     "unknown-machine-id".to_string()
 }
 
+/// Retrieves the operating system version string.
+///
+/// Reads the system's `/etc/os-release` file to determine the OS version. If the file is missing or cannot be read, returns a descriptive fallback message.
+///
+/// # Examples
+///
+/// ```
+/// let version = get_os_version();
+/// assert!(!version.is_empty());
+/// ```
 pub fn get_os_version() -> String {
     use std::fs::File;
     use std::io::{BufRead, BufReader};

agent/crates/agent-core/src/preflight.rs

@@ -13,6 +13,26 @@ pub struct PreflightReport {
 }
 
 impl PreflightReport {
+    /// Checks whether the system meets all required preflight conditions.
+    ///
+    /// Returns `true` if the configuration and log directories are writable, both osqueryd and nft
+    /// are installed, and the process is running as root.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let report = PreflightReport {
+    ///     config_dir_writable: Ok(()),
+    ///     data_dir_writable: Ok(()),
+    ///     log_dir_writable: Ok(()),
+    ///     osqueryd_installed: Ok("Found".to_string()),
+    ///     nft_installed: Ok("Found".to_string()),
+    ///     bpf_jit_enabled: Ok(true),
+    ///     inotify_watches: Ok(524288),
+    ///     is_root: true,
+    /// };
+    /// assert!(report.is_ok());
+    /// ```
     pub fn is_ok(&self) -> bool {
         self.config_dir_writable.is_ok()
             && self.data_dir_writable.is_ok()
@@ -22,6 +42,17 @@ impl PreflightReport {
             && self.is_root
     }
 
+    /// Prints a human-readable report of the preflight environment checks.
+    ///
+    /// Displays the status of all checks (root, directory accessibility, BPF JIT, inotify limits,
+    /// and required dependencies) using `[OK]`, `[WARN]`, or `[FAIL]` indicators.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let report = run_preflight(&config);
+    /// report.print();
+    /// ```
     pub fn print(&self) {
         println!("Aigis-Zero Agent Pre-flight Environment Check");
 
@@ -77,6 +108,22 @@ impl PreflightReport {
     }
 }
 
+/// Validates the system environment.
+///
+/// Checks root privilege, directory writability, BPF JIT status, inotify limits,
+/// and the presence of required executables (`osqueryd` and `nft`).
+///
+/// # Examples
+///
+/// ```no_run
+/// let config = /* ... */;
+/// let report = run_preflight(&config);
+/// if report.is_ok() {
+///     println!("Environment is ready");
+/// } else {
+///     report.print();
+/// }
+/// ```
 pub fn run_preflight(config: &crate::config::AgentConfig) -> PreflightReport {
     let is_root = unsafe { libc::getuid() } == 0;
 
@@ -151,6 +198,14 @@ pub fn run_preflight(config: &crate::config::AgentConfig) -> PreflightReport {
     }
 }
 
+/// Checks if a command is available in the system PATH.
+///
+/// # Examples
+///
+/// ```ignore
+/// assert!(which("sh"));
+/// assert!(!which("nonexistent_xyz_command"));
+/// ```
 fn which(cmd: &str) -> bool {
     Command::new("which")
         .arg(cmd)

agent/crates/fleet-client/src/enrollment.rs

@@ -6,6 +6,20 @@ use edr_sdk::codec::JsonCodec;
 pub struct AgentEnrollment;
 
 impl AgentEnrollment {
+    /// Enrolls an agent with a remote service.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the enrollment request fails.
+    ///
+    /// # Examples
+    ///
+    /// ```ignore
+    /// let channel = Channel::from_static("http://localhost:50051").connect().await?;
+    /// let request = RegisterRequest { hostname: "my-host".into(), ..Default::default() };
+    /// let result = AgentEnrollment::enroll(channel, request).await?;
+    /// println!("Node ID: {}", result.node_id);
+    /// ```
     pub async fn enroll(channel: Channel, request: RegisterRequest) -> Result<EnrollmentResult> {
         tracing::info!("Enrolling agent: {:?}", request.hostname);