diff --git a/Cargo.lock b/Cargo.lock index 438e4b2..5ca1a84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1178,9 +1178,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.11" +version = "0.103.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20a6af516fea4b20eccceaf166e8aa666ac996208e8a644ce3ef5aa783bc7cd4" +checksum = "8279bb85272c9f10811ae6a6c547ff594d6a7f3c6c6b02ee9726d1d0dcfcdd06" dependencies = [ "ring", "rustls-pki-types", diff --git a/src/api/manage_reporting.rs b/src/api/manage_reporting.rs index 5d00d75..97ecec4 100644 --- a/src/api/manage_reporting.rs +++ b/src/api/manage_reporting.rs @@ -1,5 +1,6 @@ use crate::api::manage_inject::UpdateInput; use crate::api::Client; +use crate::common::constants::{STATUS_ERROR, STATUS_SUCCESS}; use crate::handle::ExecutionOutput; pub fn report_success( @@ -22,7 +23,7 @@ pub fn report_success( agent_id.clone(), UpdateInput { execution_message, - execution_status: String::from("SUCCESS"), + execution_status: String::from(STATUS_SUCCESS), execution_duration: duration, execution_action: String::from(semantic), }, @@ -49,7 +50,7 @@ pub fn report_error( agent_id.clone(), UpdateInput { execution_message, - execution_status: String::from("ERROR"), + execution_status: String::from(STATUS_ERROR), execution_duration: duration, execution_action: String::from(semantic), }, diff --git a/src/common/constants.rs b/src/common/constants.rs new file mode 100644 index 0000000..a6f917b --- /dev/null +++ b/src/common/constants.rs @@ -0,0 +1,24 @@ +// -- EXECUTOR CONSTANTS -- +pub const EXECUTOR_BASH: &str = "bash"; +pub const EXECUTOR_CMD: &str = "cmd"; +pub const EXECUTOR_POWERSHELL: &str = "powershell"; +#[cfg(unix)] +pub const EXECUTOR_PSH: &str = "psh"; +#[cfg(windows)] +pub const EXECUTOR_PWSH: &str = "pwsh"; +pub const EXECUTOR_SH: &str = "sh"; + +// -- EXECUTION STATUS CONSTANTS -- +// Info +pub const STATUS_INFO: &str = "INFO"; +// Success +pub const STATUS_SUCCESS: &str = "SUCCESS"; +pub const STATUS_WARNING: &str = "WARNING"; +pub const STATUS_ACCESS_DENIED: &str = "ACCESS_DENIED"; +// Error +pub const STATUS_ERROR: &str = "ERROR"; +pub const STATUS_COMMAND_NOT_FOUND: &str = "COMMAND_NOT_FOUND"; +pub const STATUS_COMMAND_CANNOT_BE_EXECUTED: &str = "COMMAND_CANNOT_BE_EXECUTED"; +pub const STATUS_INVALID_USAGE: &str = "INVALID_USAGE"; +pub const STATUS_TIMEOUT: &str = "TIMEOUT"; +pub const STATUS_INTERRUPTED: &str = "INTERRUPTED"; diff --git a/src/common/execution_result.rs b/src/common/execution_result.rs new file mode 100644 index 0000000..1ca3416 --- /dev/null +++ b/src/common/execution_result.rs @@ -0,0 +1,108 @@ +use std::process::Output; + +use serde::Deserialize; + +use crate::common::constants::*; +use crate::common::error_model::Error; + +#[derive(Debug, Deserialize)] +pub struct ExecutionResult { + pub stdout: String, + pub stderr: String, + pub status: String, + pub exit_code: i32, +} + +pub fn manage_result( + invoke_output: Output, + pre_check: bool, + executor: &str, +) -> Result { + let exit_code = invoke_output.status.code().unwrap_or(-99); + let stdout = decode_output(&invoke_output.stdout); + let stderr = decode_output(&invoke_output.stderr); + + let exit_status = match exit_code { + 0 if stderr.is_empty() => STATUS_SUCCESS, + 0 if !stderr.is_empty() => STATUS_WARNING, + _ if stderr.contains("CommandNotFoundException") => STATUS_COMMAND_NOT_FOUND, + 1 if pre_check => STATUS_SUCCESS, + _ => map_exit_code(exit_code, executor, &stderr), + }; + + Ok(ExecutionResult { + stdout, + stderr, + exit_code, + status: String::from(exit_status), + }) +} + +pub fn handle_io_error(e: std::io::Error) -> Result { + let status = match e.kind() { + std::io::ErrorKind::PermissionDenied => STATUS_ACCESS_DENIED, + _ => STATUS_ERROR, + }; + Ok(ExecutionResult { + stdout: String::new(), + stderr: format!("{e}"), + exit_code: -1, + status: String::from(status), + }) +} + +// -- PRIVATE -- + +pub(crate) fn decode_output(raw_bytes: &[u8]) -> String { + // Try decoding as UTF-8 + if let Ok(decoded) = String::from_utf8(raw_bytes.to_vec()) { + return decoded; // Return if successful + } + // Fallback to UTF-8 lossy decoding + String::from_utf8_lossy(raw_bytes).to_string() +} + +#[cfg(windows)] +fn map_exit_code(exit_code: i32, executor: &str, stderr: &str) -> &'static str { + use crate::common::constants::{EXECUTOR_CMD, EXECUTOR_POWERSHELL, EXECUTOR_PWSH}; + + match executor { + EXECUTOR_CMD => match exit_code { + 5 => STATUS_ACCESS_DENIED, + 9009 => STATUS_COMMAND_NOT_FOUND, + 1460 => STATUS_TIMEOUT, + _ => STATUS_ERROR, + }, + EXECUTOR_POWERSHELL | EXECUTOR_PWSH => match exit_code { + 1 if stderr.contains("CommandNotFoundException") => STATUS_COMMAND_NOT_FOUND, + 5 => STATUS_ACCESS_DENIED, + 126 => STATUS_COMMAND_CANNOT_BE_EXECUTED, + 127 => STATUS_COMMAND_NOT_FOUND, + -1073741510 => STATUS_INTERRUPTED, + _ => STATUS_ERROR, + }, + // bash/sh on Windows + _ => match exit_code { + 2 => STATUS_INVALID_USAGE, + 5 => STATUS_ACCESS_DENIED, + 124 => STATUS_TIMEOUT, + 126 => STATUS_COMMAND_CANNOT_BE_EXECUTED, + 127 => STATUS_COMMAND_NOT_FOUND, + 130 => STATUS_INTERRUPTED, + _ => STATUS_ERROR, + }, + } +} + +#[cfg(unix)] +fn map_exit_code(exit_code: i32, _executor: &str, _stderr: &str) -> &'static str { + match exit_code { + 2 => STATUS_INVALID_USAGE, + 77 => STATUS_ACCESS_DENIED, + 124 => STATUS_TIMEOUT, + 126 => STATUS_COMMAND_CANNOT_BE_EXECUTED, + 127 => STATUS_COMMAND_NOT_FOUND, + 130 => STATUS_INTERRUPTED, + _ => STATUS_ERROR, + } +} diff --git a/src/common/logger.rs b/src/common/logger.rs new file mode 100644 index 0000000..1602a1c --- /dev/null +++ b/src/common/logger.rs @@ -0,0 +1,18 @@ +use std::path::Path; + +use rolling_file::{BasicRollingFileAppender, RollingConditionBasic}; +use tracing_appender::non_blocking::WorkerGuard; + +const PREFIX_LOG_NAME: &str = "openaev-implant.log"; + +pub fn init_logger(exe_dir: &Path) -> WorkerGuard { + let log_file = exe_dir.join(PREFIX_LOG_NAME); + let condition = RollingConditionBasic::new().daily(); + let file_appender = BasicRollingFileAppender::new(log_file, condition, 3).unwrap(); + let (file_writer, guard) = tracing_appender::non_blocking(file_appender); + tracing_subscriber::fmt() + .json() + .with_writer(file_writer) + .init(); + guard +} diff --git a/src/common/mod.rs b/src/common/mod.rs index 1214539..cd9e081 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1 +1,4 @@ +pub mod constants; pub mod error_model; +pub mod execution_result; +pub mod logger; diff --git a/src/handle/handle_dns_resolution.rs b/src/handle/handle_dns_resolution.rs index 4db3de8..c6c50da 100644 --- a/src/handle/handle_dns_resolution.rs +++ b/src/handle/handle_dns_resolution.rs @@ -4,6 +4,7 @@ use log::info; use crate::api::manage_inject::{InjectorContractPayload, UpdateInput}; use crate::api::Client; +use crate::common::constants::{STATUS_ERROR, STATUS_SUCCESS}; use crate::handle::ExecutionOutput; pub fn handle_dns_resolution( @@ -41,7 +42,7 @@ pub fn handle_dns_resolution( }; UpdateInput { execution_message: serde_json::to_string(&message).unwrap(), - execution_status: String::from("SUCCESS"), + execution_status: String::from(STATUS_SUCCESS), execution_duration: 0, execution_action: String::from("dns_resolution"), } @@ -56,7 +57,7 @@ pub fn handle_dns_resolution( }; UpdateInput { execution_message: serde_json::to_string(&message).unwrap(), - execution_status: String::from("ERROR"), + execution_status: String::from(STATUS_ERROR), execution_duration: 0, execution_action: String::from("dns_resolution"), } diff --git a/src/handle/handle_execution.rs b/src/handle/handle_execution.rs index 135870a..b08fda1 100644 --- a/src/handle/handle_execution.rs +++ b/src/handle/handle_execution.rs @@ -2,9 +2,10 @@ use log::info; use crate::api::manage_inject::UpdateInput; use crate::api::Client; +use crate::common::constants::STATUS_ERROR; use crate::common::error_model::Error; +use crate::common::execution_result::ExecutionResult; use crate::handle::{ExecutionOutput, ExecutionParam}; -use crate::process::command_exec::ExecutionResult; const PREVIEW_LOGS_SIZE: usize = 100000; @@ -57,7 +58,7 @@ pub fn handle_execution_result( params.agent_id.clone(), UpdateInput { execution_message: info_message, - execution_status: String::from("ERROR"), + execution_status: String::from(STATUS_ERROR), execution_duration: elapsed, execution_action: String::from(params.semantic.as_str()), }, @@ -94,7 +95,7 @@ pub fn handle_execution_result( params.agent_id.clone(), UpdateInput { execution_message, - execution_status: String::from("ERROR"), + execution_status: String::from(STATUS_ERROR), execution_duration: elapsed, execution_action: params.semantic.clone(), }, diff --git a/src/main.rs b/src/main.rs index 9a130ab..07f406f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,5 @@ use clap::Parser; use log::{error, info}; -use rolling_file::{BasicRollingFileAppender, RollingConditionBasic}; use std::env; use std::fs::create_dir_all; use std::ops::Deref; @@ -10,6 +9,7 @@ use std::time::Instant; use crate::api::manage_inject::{InjectorContractPayload, UpdateInput}; use crate::api::Client; +use crate::common::constants::{STATUS_ERROR, STATUS_INFO}; use crate::common::error_model::Error; use crate::handle::handle_command::{compute_command, handle_command, handle_execution_command}; use crate::handle::handle_dns_resolution::handle_dns_resolution; @@ -28,7 +28,6 @@ mod tests; pub static THREADS_CONTROL: AtomicBool = AtomicBool::new(true); const ENV_PRODUCTION: &str = "production"; const VERSION: &str = env!("CARGO_PKG_VERSION"); -const PREFIX_LOG_NAME: &str = "openaev-implant.log"; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -84,7 +83,7 @@ pub fn set_error_hook() { args.agent_id, UpdateInput { execution_message: String::from(cause), - execution_status: String::from("ERROR"), + execution_status: String::from(STATUS_ERROR), execution_duration: 0, execution_action: String::from("complete"), }, @@ -105,8 +104,6 @@ pub fn handle_payload( max_size: usize, ) { let mut prerequisites_code = 0; - let mut execution_message = "Payload completed"; - let mut execution_status = "INFO"; // region download files parameters if let Some(slice_arguments) = contract_payload.payload_arguments.as_deref() { // println!("Slice reference exists. Length: {}", slice_arguments.len()); @@ -174,25 +171,29 @@ pub fn handle_payload( if prerequisites_code == 0 { let payload_type = &contract_payload.payload_type; match payload_type.as_str() { - "Command" => handle_command( - inject_id.clone(), - agent_id.clone(), - api, - contract_payload, - max_size, - ), + "Command" => { + handle_command( + inject_id.clone(), + agent_id.clone(), + api, + contract_payload, + max_size, + ); + } "DnsResolution" => { - handle_dns_resolution(inject_id.clone(), agent_id.clone(), api, contract_payload) + handle_dns_resolution(inject_id.clone(), agent_id.clone(), api, contract_payload); + } + "Executable" => { + handle_file_execute( + inject_id.clone(), + agent_id.clone(), + api, + contract_payload, + max_size, + ); } - "Executable" => handle_file_execute( - inject_id.clone(), - agent_id.clone(), - api, - contract_payload, - max_size, - ), "FileDrop" => { - handle_file_drop(inject_id.clone(), agent_id.clone(), api, contract_payload) + handle_file_drop(inject_id.clone(), agent_id.clone(), api, contract_payload); } // "NetworkTraffic" => {}, // Not implemented yet _ => { @@ -201,16 +202,13 @@ pub fn handle_payload( agent_id.clone(), UpdateInput { execution_message: String::from("Payload execution type not supported."), - execution_status: String::from("ERROR"), + execution_status: String::from(STATUS_ERROR), execution_duration: duration.elapsed().as_millis(), - execution_action: String::from("complete"), + execution_action: String::from("command_execution"), }, ); } } - } else { - execution_message = "Payload execution not executed due to dependencies failure."; - execution_status = "ERROR"; } // endregion // region cleanup execution @@ -239,8 +237,8 @@ pub fn handle_payload( inject_id.clone(), agent_id.clone(), UpdateInput { - execution_message: String::from(execution_message), - execution_status: String::from(execution_status), + execution_message: String::from("Payload completed"), + execution_status: String::from(STATUS_INFO), execution_duration: duration.elapsed().as_millis(), execution_action: String::from("complete"), }, @@ -249,11 +247,10 @@ pub fn handle_payload( fn main() -> Result<(), Error> { set_error_hook(); - // region Init logger let duration = Instant::now(); let current_exe_path = env::current_exe().unwrap(); let parent_path = current_exe_path.parent().unwrap(); - let log_file = parent_path.join(PREFIX_LOG_NAME); + let _logger_file = common::logger::init_logger(parent_path); // Resolve the payloads path and create it on the fly let folder_name = parent_path.file_name().unwrap().to_str().unwrap(); @@ -266,14 +263,6 @@ fn main() -> Result<(), Error> { .join(folder_name); create_dir_all(payloads_path).expect("Unable to create payload directory"); - let condition = RollingConditionBasic::new().daily(); - let file_appender = BasicRollingFileAppender::new(log_file, condition, 3).unwrap(); - let (file_writer, _guard) = tracing_appender::non_blocking(file_appender); - tracing_subscriber::fmt() - .json() - .with_writer(file_writer) - .init(); - // endregion // region Process execution let args = Args::parse(); diff --git a/src/process/command_exec.rs b/src/process/command_exec.rs index 626516b..9e1f3de 100644 --- a/src/process/command_exec.rs +++ b/src/process/command_exec.rs @@ -1,39 +1,30 @@ -use std::io::ErrorKind; -use std::process::{Command, ExitStatus, Output, Stdio}; +use std::process::{Command, Stdio}; use base64::{engine::general_purpose::STANDARD, Engine as _}; -use serde::Deserialize; +#[cfg(unix)] +use crate::common::constants::EXECUTOR_PSH; +use crate::common::constants::{EXECUTOR_BASH, EXECUTOR_CMD, EXECUTOR_POWERSHELL, EXECUTOR_SH}; use crate::common::error_model::Error; +use crate::common::execution_result::{handle_io_error, manage_result, ExecutionResult}; use crate::handle::handle_command::compute_command; use crate::process::exec_utils::is_executor_present; -#[cfg(unix)] -use std::os::unix::process::ExitStatusExt; #[cfg(windows)] use std::os::windows::process::CommandExt; -#[cfg(windows)] -use std::os::windows::process::ExitStatusExt; - -#[derive(Debug, Deserialize)] -pub struct ExecutionResult { - pub stdout: String, - pub stderr: String, - pub status: String, - pub exit_code: i32, -} pub fn invoke_command( executor: &str, cmd_expression: &str, args: &[&str], -) -> std::io::Result { + pre_check: bool, +) -> Result { let mut command = Command::new(executor); let result = match executor { // For CMD we use "raw_args" to fix issue #3161; #[cfg(windows)] - "cmd" => command.args(args).raw_arg(cmd_expression), + EXECUTOR_CMD => command.args(args).raw_arg(cmd_expression), // for other executors, we still use "args" as they are working properly. _ => command.args(args).arg(cmd_expression), } @@ -41,21 +32,8 @@ pub fn invoke_command( .output(); match result { - Ok(output) => Ok(output), - Err(e) if e.kind() == ErrorKind::PermissionDenied => { - let exit_status = if cfg!(unix) { - ExitStatus::from_raw(256) - } else { - ExitStatus::from_raw(1) - }; - - Ok(Output { - status: exit_status, - stdout: Vec::new(), - stderr: format!("{e}").into_bytes(), - }) - } - Err(e) => Err(e), + Ok(output) => manage_result(output, pre_check, executor), + Err(e) => handle_io_error(e), } } @@ -77,58 +55,24 @@ pub fn format_windows_command(command: String) -> String { format!("setlocal & {command} & exit /b errorlevel") } -pub fn manage_result(invoke_output: Output, pre_check: bool) -> Result { - let invoke_result = invoke_output.clone(); - let exit_code = invoke_result.status.code().unwrap_or(-99); - - let stdout = decode_output(&invoke_result.stdout); - let stderr = decode_output(&invoke_result.stderr); - - let exit_status = match exit_code { - 0 if stderr.is_empty() => "SUCCESS", - 0 if !stderr.is_empty() => "WARNING", - 1 if pre_check => "SUCCESS", - -99 => "ERROR", - 127 => "COMMAND_NOT_FOUND", - 126 => "COMMAND_CANNOT_BE_EXECUTED", - _ => "MAYBE_PREVENTED", - }; - - Ok(ExecutionResult { - stdout, - stderr, - exit_code, - status: String::from(exit_status), - }) -} - -pub fn decode_output(raw_bytes: &[u8]) -> String { - // Try decoding as UTF-8 - if let Ok(decoded) = String::from_utf8(raw_bytes.to_vec()) { - return decoded; // Return if successful - } - // Fallback to UTF-8 lossy decoding - String::from_utf8_lossy(raw_bytes).to_string() -} - -#[cfg(target_os = "windows")] +#[cfg(windows)] pub fn get_executor(executor: &str) -> &str { match executor { - "cmd" | "bash" | "sh" => executor, - _ => "powershell", + EXECUTOR_CMD | EXECUTOR_BASH | EXECUTOR_SH => executor, + _ => EXECUTOR_POWERSHELL, } } -#[cfg(any(target_os = "linux", target_os = "macos"))] +#[cfg(unix)] pub fn get_executor(executor: &str) -> &str { match executor { - "bash" => executor, - "psh" => "powershell", - _ => "sh", + EXECUTOR_BASH => executor, + EXECUTOR_PSH => EXECUTOR_POWERSHELL, + _ => EXECUTOR_SH, } } -#[cfg(target_os = "windows")] +#[cfg(windows)] pub fn get_psh_arg() -> Vec<&'static str> { Vec::from([ "-ExecutionPolicy", @@ -141,7 +85,7 @@ pub fn get_psh_arg() -> Vec<&'static str> { ]) } -#[cfg(any(target_os = "linux", target_os = "macos"))] +#[cfg(unix)] pub fn get_psh_arg() -> Vec<&'static str> { Vec::from([ "-ExecutionPolicy", @@ -167,14 +111,13 @@ pub fn command_execution( ))); } - if final_executor == "cmd" { + if final_executor == EXECUTOR_CMD { formatted_cmd = format_windows_command(formatted_cmd); args = vec!["/V", "/C"]; - } else if final_executor == "powershell" { + } else if final_executor == EXECUTOR_POWERSHELL { formatted_cmd = format_powershell_command(formatted_cmd); args = get_psh_arg(); } - let invoke_output = invoke_command(final_executor, &formatted_cmd, args.as_slice()); - manage_result(invoke_output?, pre_check) + invoke_command(final_executor, &formatted_cmd, args.as_slice(), pre_check) } diff --git a/src/process/file_exec.rs b/src/process/file_exec.rs index efdddf7..01c6238 100644 --- a/src/process/file_exec.rs +++ b/src/process/file_exec.rs @@ -1,9 +1,13 @@ use std::env; use std::path::PathBuf; -use std::process::{Command, Output, Stdio}; +use std::process::{Command, Stdio}; +#[cfg(unix)] +use crate::common::constants::EXECUTOR_BASH; +#[cfg(windows)] +use crate::common::constants::EXECUTOR_POWERSHELL; use crate::common::error_model::Error; -use crate::process::command_exec::ExecutionResult; +use crate::common::execution_result::{handle_io_error, manage_result, ExecutionResult}; use crate::process::exec_utils::is_executor_present; fn compute_working_file(filename: &str) -> PathBuf { @@ -26,31 +30,7 @@ fn compute_working_file(filename: &str) -> PathBuf { executable_path.join(filename) } -pub fn manage_result(invoke_output: Output) -> Result { - let invoke_result = invoke_output.clone(); - // 0 success | other = maybe prevented - let exit_code = invoke_result.status.code().unwrap_or(-99); - let stdout = String::from_utf8_lossy(&invoke_result.stdout).to_string(); - let stderr = String::from_utf8_lossy(&invoke_result.stderr).to_string(); - - let exit_status = match exit_code { - 0 if stderr.is_empty() => "SUCCESS", - 0 if !stderr.is_empty() => "WARNING", - -99 => "ERROR", - 127 => "COMMAND_NOT_FOUND", - 126 => "COMMAND_CANNOT_BE_EXECUTED", - _ => "MAYBE_PREVENTED", - }; - - Ok(ExecutionResult { - stdout, - stderr, - exit_code, - status: String::from(exit_status), - }) -} - -#[cfg(target_os = "windows")] +#[cfg(windows)] pub fn file_execution(filename: &str) -> Result { let executor = "powershell.exe"; if !is_executor_present(executor) { @@ -79,12 +59,16 @@ pub fn file_execution(filename: &str) -> Result { .stdout(Stdio::piped()) .spawn()? .wait_with_output(); - manage_result(invoke_output?) + + match invoke_output { + Ok(output) => manage_result(output, false, EXECUTOR_POWERSHELL), + Err(e) => handle_io_error(e), + } } -#[cfg(any(target_os = "linux", target_os = "macos"))] +#[cfg(unix)] pub fn file_execution(filename: &str) -> Result { - let executor = "bash"; + let executor = EXECUTOR_BASH; if !is_executor_present(executor) { return Err(Error::Internal(format!( "Executor '{executor}' is not available." @@ -99,5 +83,9 @@ pub fn file_execution(filename: &str) -> Result { .stdout(Stdio::piped()) .spawn()? .wait_with_output(); - manage_result(invoke_output?) + + match invoke_output { + Ok(output) => manage_result(output, false, EXECUTOR_BASH), + Err(e) => handle_io_error(e), + } } diff --git a/src/tests/process/command_exec_tests.rs b/src/tests/process/command_exec_tests.rs index 7c3dcd4..ce2cf44 100644 --- a/src/tests/process/command_exec_tests.rs +++ b/src/tests/process/command_exec_tests.rs @@ -1,7 +1,17 @@ -use crate::process::command_exec::decode_output; +#[cfg(windows)] +use crate::common::constants::EXECUTOR_CMD; +#[cfg(windows)] +use crate::common::constants::EXECUTOR_POWERSHELL; +use crate::common::constants::STATUS_ERROR; +#[cfg(windows)] +use crate::common::constants::{STATUS_COMMAND_NOT_FOUND, STATUS_SUCCESS, STATUS_WARNING}; +use crate::common::execution_result::decode_output; +#[cfg(windows)] use crate::process::command_exec::format_powershell_command; use crate::process::command_exec::invoke_command; +// -- DECODE_OUTPUT -- + #[test] fn test_decode_output_with_hello() { let output = vec![72, 101, 108, 108, 111]; @@ -27,20 +37,20 @@ fn test_decode_output_with_wrong_character() { assert_eq!(decoded_output, "Hello�"); } +// -- INVOKE_COMMAND -- + #[ignore] #[test] +#[cfg(windows)] fn test_invoke_command_powershell_special_character() { let command = "echo Helloé"; let formatted_cmd = format_powershell_command(command.to_string()); let args: Vec<&str> = vec!["-c"]; - let invoke_output = match invoke_command("powershell", &formatted_cmd, args.as_slice()) { - Ok(output) => output, - Err(e) => panic!("Failed to invoke PowerShell command: {}", e), - }; + let result = invoke_command(EXECUTOR_POWERSHELL, &formatted_cmd, args.as_slice(), false) + .expect("Failed to invoke PowerShell command"); - let stdout = decode_output(&invoke_output.stdout); - assert_eq!(stdout, "Helloé\r\n"); + assert_eq!(result.stdout, "Helloé\r\n"); } #[test] @@ -49,11 +59,138 @@ fn test_invoke_command_cmd_with_quote() { let command = r#"echo "Hello""#; let args: Vec<&str> = vec!["/V", "/C"]; - let invoke_output = match invoke_command("cmd", &command, args.as_slice()) { - Ok(output) => output, - Err(e) => panic!("Failed to invoke CMD command: {}", e), - }; + let result = invoke_command(EXECUTOR_CMD, &command, args.as_slice(), false) + .expect("Failed to invoke CMD command"); + + assert_eq!(result.stdout, "\"Hello\"\r\n"); +} + +// ============================================================= +// Integration tests: real payload execution +// ============================================================= + +// -- SUCCESS scenarios -- + +#[test] +#[cfg(windows)] +fn test_real_cmd_echo_success() { + let command = "echo Hello World"; + let args: Vec<&str> = vec!["/V", "/C"]; + + let result = invoke_command(EXECUTOR_CMD, command, args.as_slice(), false) + .expect("Failed to invoke CMD command"); + + assert_eq!(result.status, STATUS_SUCCESS); + assert_eq!(result.exit_code, 0); + assert!(result.stdout.contains("Hello World")); +} + +#[test] +#[cfg(windows)] +fn test_real_powershell_echo_success() { + let command = format_powershell_command("Write-Output 'Hello World'".to_string()); + let args: Vec<&str> = vec!["-NonInteractive", "-NoProfile", "-Command"]; + + let result = invoke_command(EXECUTOR_POWERSHELL, &command, args.as_slice(), false) + .expect("Failed to invoke PowerShell command"); + + assert_eq!(result.status, STATUS_SUCCESS); + assert_eq!(result.exit_code, 0); + assert!(result.stdout.contains("Hello World")); +} + +// -- COMMAND_NOT_FOUND scenarios -- + +#[test] +#[cfg(windows)] +fn test_real_cmd_command_not_found() { + let command = "this_command_does_not_exist_xyz"; + let args: Vec<&str> = vec!["/V", "/C"]; + + let result = invoke_command(EXECUTOR_CMD, command, args.as_slice(), false) + .expect("Should return result, not error"); + + // CMD returns 9009 for command not found, mapped to COMMAND_NOT_FOUND + // but some CMD versions may return 1; we accept either COMMAND_NOT_FOUND or ERROR + assert!(result.exit_code != 0, "Expected non-zero exit code"); + assert!( + result.status == STATUS_COMMAND_NOT_FOUND || result.status == STATUS_ERROR, + "Expected COMMAND_NOT_FOUND or ERROR, got: {}", + result.status + ); +} + +#[test] +#[cfg(windows)] +fn test_real_powershell_command_not_found() { + let command = format_powershell_command("this_command_does_not_exist_xyz".to_string()); + let args: Vec<&str> = vec!["-NonInteractive", "-NoProfile", "-Command"]; + + let result = invoke_command(EXECUTOR_POWERSHELL, &command, args.as_slice(), false) + .expect("Should return result, not error"); + + assert_eq!(result.status, STATUS_COMMAND_NOT_FOUND); + assert!(result.stderr.contains("CommandNotFoundException")); +} + +// -- ERROR scenarios -- + +#[test] +#[cfg(windows)] +fn test_real_cmd_error_exit_code() { + let command = "exit /b 42"; + let args: Vec<&str> = vec!["/V", "/C"]; + + let result = invoke_command(EXECUTOR_CMD, command, args.as_slice(), false) + .expect("Should return result, not error"); + + assert_eq!(result.status, STATUS_ERROR); + assert_eq!(result.exit_code, 42); +} + +#[test] +#[cfg(windows)] +fn test_real_powershell_error_exit_code() { + let command = format_powershell_command("throw 'Something went wrong'".to_string()); + let args: Vec<&str> = vec!["-NonInteractive", "-NoProfile", "-Command"]; + + let result = invoke_command(EXECUTOR_POWERSHELL, &command, args.as_slice(), false) + .expect("Should return result, not error"); + + assert_eq!(result.status, STATUS_ERROR); + assert_ne!(result.exit_code, 0); + assert!(!result.stderr.is_empty()); +} + +// -- Executor not found (io error) -- + +#[test] +fn test_real_nonexistent_executor() { + let command = "echo hello"; + let args: Vec<&str> = vec!["-c"]; + + let result = invoke_command("nonexistent_shell_xyz", command, args.as_slice(), false) + .expect("Should return result with error status"); + + assert_eq!(result.status, STATUS_ERROR); + assert_eq!(result.exit_code, -1); + assert!(!result.stderr.is_empty()); +} + +// -- WARNING scenario (exit 0 but stderr not empty) -- + +#[test] +#[cfg(windows)] +fn test_real_powershell_warning() { + let command = format_powershell_command( + "[Console]::Error.WriteLine('warn'); Write-Output 'ok'".to_string(), + ); + let args: Vec<&str> = vec!["-NonInteractive", "-NoProfile", "-Command"]; + + let result = invoke_command(EXECUTOR_POWERSHELL, &command, args.as_slice(), false) + .expect("Should return result"); - let stdout = decode_output(&invoke_output.stdout); - assert_eq!(stdout, "\"Hello\"\r\n"); //assert that raw_args was used otherwise " should have been escaped twice + assert_eq!(result.exit_code, 0); + assert_eq!(result.status, STATUS_WARNING); + assert!(result.stdout.contains("ok")); } diff --git a/src/tests/process/execution_result_tests.rs b/src/tests/process/execution_result_tests.rs new file mode 100644 index 0000000..b9ee7eb --- /dev/null +++ b/src/tests/process/execution_result_tests.rs @@ -0,0 +1,179 @@ +#[cfg(unix)] +use std::os::unix::process::ExitStatusExt; +#[cfg(windows)] +use std::os::windows::process::ExitStatusExt; + +use std::process::{ExitStatus, Output}; + +use crate::common::constants::*; +use crate::common::execution_result::manage_result; + +// -- HELPERS -- + +fn make_output(exit_code: i32, stdout: &str, stderr: &str) -> Output { + #[cfg(windows)] + let status = ExitStatus::from_raw(exit_code as u32); + #[cfg(unix)] + let status = ExitStatus::from_raw(exit_code << 8); + Output { + status, + stdout: stdout.as_bytes().to_vec(), + stderr: stderr.as_bytes().to_vec(), + } +} + +// -- CROSS-CUTTING RULES (ALL OS) -- + +#[test] +fn test_success() { + let output = make_output(0, "ok", ""); + let result = manage_result(output, false, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_SUCCESS); + assert_eq!(result.exit_code, 0); +} + +#[test] +fn test_warning_when_stderr_non_empty() { + let output = make_output(0, "ok", "some warning"); + let result = manage_result(output, false, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_WARNING); +} + +#[test] +fn test_pre_check_success_on_code_1() { + let output = make_output(1, "", "stderr"); + let result = manage_result(output, true, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_SUCCESS); +} + +// -- CMD (WINDOWS) -- + +#[test] +#[cfg(windows)] +fn test_cmd_exit_codes() { + let output = make_output(1, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_ERROR); + + let output = make_output(5, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_ACCESS_DENIED); + + let output = make_output(9009, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_COMMAND_NOT_FOUND); + + let output = make_output(1460, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_TIMEOUT); + + let output = make_output(42, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_CMD).unwrap(); + assert_eq!(result.status, STATUS_ERROR); +} + +// -- POWERSHELL (WINDOWS) -- + +#[test] +#[cfg(windows)] +fn test_powershell_exit_codes() { + let output = make_output(1, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_POWERSHELL).unwrap(); + assert_eq!(result.status, STATUS_ERROR); + + let output = make_output(5, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_POWERSHELL).unwrap(); + assert_eq!(result.status, STATUS_ACCESS_DENIED); + + let output = make_output(126, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_POWERSHELL).unwrap(); + assert_eq!(result.status, STATUS_COMMAND_CANNOT_BE_EXECUTED); + + let output = make_output(127, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_POWERSHELL).unwrap(); + assert_eq!(result.status, STATUS_COMMAND_NOT_FOUND); + + let output = make_output(-1073741510, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_POWERSHELL).unwrap(); + assert_eq!(result.status, STATUS_INTERRUPTED); + + let output = make_output(99, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_POWERSHELL).unwrap(); + assert_eq!(result.status, STATUS_ERROR); +} + +// -- BASH/SH ON WINDOWS -- + +#[test] +#[cfg(windows)] +fn test_bash_windows_exit_codes() { + let output = make_output(1, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_ERROR); + + let output = make_output(2, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_INVALID_USAGE); + + let output = make_output(5, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_ACCESS_DENIED); + + let output = make_output(124, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_TIMEOUT); + + let output = make_output(126, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_COMMAND_CANNOT_BE_EXECUTED); + + let output = make_output(127, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_COMMAND_NOT_FOUND); + + let output = make_output(130, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_INTERRUPTED); + + let output = make_output(42, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_BASH).unwrap(); + assert_eq!(result.status, STATUS_ERROR); +} + +// -- UNIX (BASH/SH) -- + +#[test] +#[cfg(unix)] +fn test_unix_exit_codes() { + let output = make_output(1, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_ERROR); + + let output = make_output(2, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_INVALID_USAGE); + + let output = make_output(77, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_ACCESS_DENIED); + + let output = make_output(124, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_TIMEOUT); + + let output = make_output(126, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_COMMAND_CANNOT_BE_EXECUTED); + + let output = make_output(127, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_COMMAND_NOT_FOUND); + + let output = make_output(130, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_INTERRUPTED); + + let output = make_output(42, "", "stderr"); + let result = manage_result(output, false, EXECUTOR_SH).unwrap(); + assert_eq!(result.status, STATUS_ERROR); +} diff --git a/src/tests/process/mod.rs b/src/tests/process/mod.rs index 6d67ff2..562b7e2 100644 --- a/src/tests/process/mod.rs +++ b/src/tests/process/mod.rs @@ -1 +1,2 @@ mod command_exec_tests; +mod execution_result_tests;