From 7845f85714cdc4feedf049d6f3e1c870054b88de Mon Sep 17 00:00:00 2001 From: savacano28 Date: Tue, 5 May 2026 18:17:42 +0200 Subject: [PATCH 1/2] [agent] fix: customize clean up thresholds --- config/default.toml | 8 +- src/config/settings.rs | 35 +++++- src/main.rs | 2 +- src/process/agent_cleanup.rs | 94 +++++++++------- src/tests/process/agent_cleanup_tests.rs | 130 +++++++++++++++++++++++ src/tests/process/mod.rs | 1 + 6 files changed, 228 insertions(+), 42 deletions(-) create mode 100644 src/tests/process/agent_cleanup_tests.rs diff --git a/config/default.toml b/config/default.toml index 2d0773a2..31ebd8c8 100644 --- a/config/default.toml +++ b/config/default.toml @@ -7,4 +7,10 @@ unsecured_certificate = false with_proxy = false installation_mode= "session-user" service_name = "OAEVAgentService" -service_full_name = "OAEVAgentService" \ No newline at end of file +service_full_name = "OAEVAgentService" +tenant_id = "ChangeMe" + +[cleanup] +executing_max_time_minutes = 10 # inject.execution.threshold.minutes default is 10minutes +directory_max_time_minutes = 10 # default +cleanup_interval_seconds = 180 diff --git a/src/config/settings.rs b/src/config/settings.rs index 201a4295..91c52b27 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -5,7 +5,7 @@ use std::env; const ENV_PRODUCTION: &str = "production"; const ENV_PRODUCTION_CONFIG_FILE: &str = "openaev-agent-config"; -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Clone)] #[allow(unused)] pub struct OpenAEV { pub url: String, @@ -16,11 +16,42 @@ pub struct OpenAEV { pub service_name: String, } -#[derive(Debug, Deserialize)] +fn default_executing_max_time_minutes() -> u64 { + 10 +} +fn default_directory_max_time_minutes() -> u64 { + 10 +} +fn default_cleanup_interval_seconds() -> u64 { + 180 +} +#[derive(Debug, Deserialize, Clone)] +pub struct CleanupSettings { + #[serde(default = "default_executing_max_time_minutes")] + pub executing_max_time_minutes: u64, + #[serde(default = "default_directory_max_time_minutes")] + pub directory_max_time_minutes: u64, + #[serde(default = "default_cleanup_interval_seconds")] + pub cleanup_interval_seconds: u64, +} + +impl Default for CleanupSettings { + fn default() -> Self { + Self { + executing_max_time_minutes: default_executing_max_time_minutes(), + directory_max_time_minutes: default_directory_max_time_minutes(), + cleanup_interval_seconds: default_cleanup_interval_seconds(), + } + } +} + +#[derive(Debug, Deserialize, Clone)] #[allow(unused)] pub struct Settings { pub debug: bool, pub openaev: OpenAEV, + #[serde(default)] + pub cleanup: CleanupSettings, } impl Settings { diff --git a/src/main.rs b/src/main.rs index 54eb3689..2d00488f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -95,7 +95,7 @@ fn agent_start(settings_data: Settings, is_service: bool) -> Result Result, Error> { let now = SystemTime::now(); - let current_exe_patch = env::current_exe().unwrap(); - let executable_path = current_exe_patch.parent().unwrap(); - let entries = fs::read_dir(executable_path.join(subfolder)).unwrap(); + let current_exe_patch = env::current_exe()?; + let executable_path = current_exe_patch.parent().ok_or_else(|| { + Error::new( + std::io::ErrorKind::NotFound, + "Cannot resolve parent directory of current executable", + ) + })?; + let entries = fs::read_dir(executable_path.join(subfolder))?; entries .into_iter() .filter(|entry| { @@ -60,71 +64,85 @@ fn create_cleanup_scripts() { } } -pub fn clean() -> Result, Error> { +fn kill_processes_for_directory(dirname: &str) { + let escaped_dirname = format!("\"{dirname}\""); + if cfg!(target_os = "windows") { + Command::new("powershell") + .args([ + "-ExecutionPolicy", + "Bypass", + "openaev_agent_kill.ps1", + escaped_dirname.as_str(), + ]) + .output() + .unwrap(); + } + if cfg!(target_os = "linux") || cfg!(target_os = "macos") { + Command::new("bash") + .args(["openaev_agent_kill.sh", dirname]) + .output() + .unwrap(); + } +} + +pub fn clean(cleanup: CleanupSettings) -> Result, Error> { info!("Starting cleanup thread"); let handle = thread::spawn(move || { // Create the expected script per operating system. create_cleanup_scripts(); + + let executing_max_time = cleanup.executing_max_time_minutes; + let directory_max_time = cleanup.directory_max_time_minutes; + let cleanup_interval = cleanup.cleanup_interval_seconds; + // While no stop signal received while THREADS_CONTROL.load(Ordering::Relaxed) { + // region Handle killing old execution- directories let kill_runtimes_directories = - get_old_execution_directories("runtimes", "execution-", EXECUTING_MAX_TIME) + get_old_execution_directories("runtimes", EXECUTION_PREFIX, executing_max_time) .unwrap(); - // region Handle killing old execution- directories for dir in kill_runtimes_directories { let dir_path = dir.path(); let dirname = dir_path.to_str().unwrap(); - info!("[cleanup thread] Killing process for directory {dirname}"); - let escaped_dirname = format!("\"{dirname}\""); - if cfg!(target_os = "windows") { - Command::new("powershell") - .args([ - "-ExecutionPolicy", - "Bypass", - "openaev_agent_kill.ps1", - escaped_dirname.as_str(), - ]) - .output() - .unwrap(); - } - if cfg!(target_os = "linux") || cfg!(target_os = "macos") { - Command::new("bash") - .args(["openaev_agent_kill.sh", dirname]) - .output() - .unwrap(); - } + info!("[cleanup thread] Killing process for runtime directory {dirname}"); + kill_processes_for_directory(dirname); // After kill, rename from execution to executed + info!("[cleanup thread] Renaming runtime directory {dirname}"); fs::rename(dirname, dirname.replace("execution", "executed")).unwrap(); } let rename_payloads_directories = - get_old_execution_directories("payloads", "execution-", EXECUTING_MAX_TIME) + get_old_execution_directories("payloads", EXECUTION_PREFIX, executing_max_time) .unwrap(); for dir in rename_payloads_directories { let dir_path = dir.path(); let dirname = dir_path.to_str().unwrap(); + info!("[cleanup thread] Renaming payload directory {dirname}"); fs::rename(dirname, dirname.replace("execution", "executed")).unwrap(); } // endregion + // region Handle remove of old executed- directories let remove_runtimes_directories = - get_old_execution_directories("runtimes", "executed-", DIRECTORY_MAX_TIME).unwrap(); + get_old_execution_directories("runtimes", EXECUTED_PREFIX, directory_max_time) + .unwrap(); for dir in remove_runtimes_directories { let dir_path = dir.path(); let dirname = dir_path.to_str().unwrap(); - info!("[cleanup thread] Removing directory {dirname}"); + info!("[cleanup thread] Removing runtime directory {dirname}"); fs::remove_dir_all(dir_path).unwrap() } let remove_payloads_directories = - get_old_execution_directories("payloads", "executed-", DIRECTORY_MAX_TIME).unwrap(); + get_old_execution_directories("payloads", EXECUTED_PREFIX, directory_max_time) + .unwrap(); for dir in remove_payloads_directories { let dir_path = dir.path(); let dirname = dir_path.to_str().unwrap(); - info!("[cleanup thread] Removing directory {dirname}"); + info!("[cleanup thread] Removing payload directory {dirname}"); fs::remove_dir_all(dir_path).unwrap() } // endregion - // Wait for the next cleanup (3 minutes) - sleep(Duration::from_secs(3 * 60)); + // Wait for the next cleanup + sleep(Duration::from_secs(cleanup_interval)); } }); Ok(handle) diff --git a/src/tests/process/agent_cleanup_tests.rs b/src/tests/process/agent_cleanup_tests.rs new file mode 100644 index 00000000..95bbbf75 --- /dev/null +++ b/src/tests/process/agent_cleanup_tests.rs @@ -0,0 +1,130 @@ +#[cfg(test)] +mod tests { + use std::env; + use std::fs; + use std::fs::create_dir_all; + use std::path::PathBuf; + use std::thread::sleep; + use std::time::Duration; + + fn compute_working_dir() -> PathBuf { + let current_exe_path = env::current_exe().unwrap(); + current_exe_path.parent().unwrap().to_path_buf() + } + + fn create_test_directory(subfolder: &str, prefix: &str, id: &str) -> PathBuf { + let working_dir = compute_working_dir(); + let dir = working_dir.join(subfolder).join(format!("{prefix}{id}")); + create_dir_all(&dir).unwrap(); + // Write a dummy file inside to simulate execution output + fs::write(dir.join("test.txt"), "test content").unwrap(); + dir + } + + fn cleanup_test_directory(path: &PathBuf) { + if path.exists() { + let _ = fs::remove_dir_all(path); + } + } + + // -- Tests for get_old_execution_directories -- + + #[test] + fn test_get_old_execution_directories_finds_execution_prefix() { + let working_dir = compute_working_dir(); + create_dir_all(working_dir.join("runtimes")).unwrap(); + + let test_id = "test-exec-find-001"; + let dir = create_test_directory("runtimes", "execution-", test_id); + + // With since_minutes=0, any directory should be returned (it's older than 0 minutes) + // We need to wait at least 1 second so modified time is in the past + sleep(Duration::from_millis(100)); + + // Verify the directory exists + assert!(dir.exists()); + assert!(dir.join("test.txt").exists()); + + // Cleanup + cleanup_test_directory(&dir); + } + + #[test] + fn test_get_old_execution_directories_ignores_unmatched_prefix() { + let working_dir = compute_working_dir(); + create_dir_all(working_dir.join("runtimes")).unwrap(); + + let test_id = "test-unknown-001"; + let dir = create_test_directory("runtimes", "unknown-", test_id); + + // Verify it exists but would not be matched by cleanup + assert!(dir.exists()); + let file_name = dir.file_name().unwrap().to_str().unwrap(); + assert!(!file_name.contains("execution-")); + assert!(!file_name.contains("executed-")); + + // Cleanup + cleanup_test_directory(&dir); + } + + #[test] + fn test_execution_directory_rename_logic() { + let working_dir = compute_working_dir(); + create_dir_all(working_dir.join("runtimes")).unwrap(); + + let test_id = "test-rename-001"; + let dir = create_test_directory("runtimes", "execution-", test_id); + + // Simulate rename logic (same as in agent_cleanup) + let dirname = dir.to_str().unwrap(); + let new_name = dirname.replace("execution", "executed"); + fs::rename(dirname, &new_name).unwrap(); + + let new_path = PathBuf::from(&new_name); + assert!(new_path.exists()); + assert!(!dir.exists()); + // File inside should still be there after rename + assert!(new_path.join("test.txt").exists()); + + // Cleanup + cleanup_test_directory(&new_path); + } + + #[test] + fn test_executed_directory_delete_logic() { + let working_dir = compute_working_dir(); + create_dir_all(working_dir.join("runtimes")).unwrap(); + + let test_id = "test-executed-delete-001"; + let dir = create_test_directory("runtimes", "executed-", test_id); + + assert!(dir.exists()); + assert!(dir.join("test.txt").exists()); + + // Simulate executed cleanup (permanent delete) + fs::remove_dir_all(&dir).unwrap(); + + assert!(!dir.exists()); + } + + #[test] + fn test_payloads_directory_cleanup() { + let working_dir = compute_working_dir(); + create_dir_all(working_dir.join("payloads")).unwrap(); + + let test_id = "test-payload-001"; + let exec_dir = create_test_directory("payloads", "execution-", test_id); + + assert!(exec_dir.exists()); + + // Simulate rename for execution- + let dirname = exec_dir.to_str().unwrap(); + let new_name = dirname.replace("execution", "executed"); + fs::rename(dirname, &new_name).unwrap(); + let renamed_path = PathBuf::from(&new_name); + assert!(renamed_path.exists()); + + // Cleanup + cleanup_test_directory(&renamed_path); + } +} diff --git a/src/tests/process/mod.rs b/src/tests/process/mod.rs index 99cf709a..156a6d02 100644 --- a/src/tests/process/mod.rs +++ b/src/tests/process/mod.rs @@ -1 +1,2 @@ +mod agent_cleanup_tests; mod agent_exec_tests; From 6027649ea0d9dd948def0ca5600d9b03f77dca8c Mon Sep 17 00:00:00 2001 From: savacano28 Date: Wed, 6 May 2026 18:19:19 +0200 Subject: [PATCH 2/2] [agent] fix: customize clean up thresholds --- config/default.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/config/default.toml b/config/default.toml index 31ebd8c8..74771f73 100644 --- a/config/default.toml +++ b/config/default.toml @@ -8,7 +8,6 @@ with_proxy = false installation_mode= "session-user" service_name = "OAEVAgentService" service_full_name = "OAEVAgentService" -tenant_id = "ChangeMe" [cleanup] executing_max_time_minutes = 10 # inject.execution.threshold.minutes default is 10minutes