[implant] feat(status): enhance Inject Result Clarity & Actionability (#115)

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),
         },

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";

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<ExecutionResult, Error> {
+    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<ExecutionResult, Error> {
+    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,
+    }
+}

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
+}

src/common/mod.rs

@@ -1 +1,4 @@
+pub mod constants;
 pub mod error_model;
+pub mod execution_result;
+pub mod logger;

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"),
                 }

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(),
                 },

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();