From dd7c71601e536d31d038e01d929372d939e364e4 Mon Sep 17 00:00:00 2001
From: GridexX <arsene582@gmail.com>
Date: Wed, 22 Mar 2023 10:13:49 +0100
Subject: [PATCH] feat(agent): request and response payload added

Signed-off-by: GridexX <arsene582@gmail.com>
---
 agent/lib/src/external_api/model.rs   |  47 ++++----
 agent/lib/src/external_api/service.rs |  28 +++--
 agent/lib/src/internal_api/model.rs   |  10 --
 agent/lib/src/internal_api/service.rs | 165 +++++++++++++++++---------
 agent/src/main.rs                     |  20 +---
 5 files changed, 158 insertions(+), 112 deletions(-)

diff --git a/agent/lib/src/external_api/model.rs b/agent/lib/src/external_api/model.rs
index b62bea5..ac7d033 100644
--- a/agent/lib/src/external_api/model.rs
+++ b/agent/lib/src/external_api/model.rs
@@ -32,10 +32,10 @@ pub struct StatusMessage {
 }
 
 impl StatusMessage {
-    pub fn new(code: String) -> StatusMessage {
+    pub fn new() -> StatusMessage {
         StatusMessage {
             r#type: Type::Status,
-            code: Code::Run
+            code: Code::Ok,
         }
     }
 }
@@ -46,17 +46,23 @@ pub struct ResponseStep {
     pub result: i32,
     pub stdout: String,
     pub stderr: String,
-    pub enable_output: bool
+    pub enable_output: bool,
 }
 
 impl ResponseStep {
-    pub fn new(command: String, result: i32, stdout: String, stderr: String, enable_output: bool) -> ResponseStep {
+    pub fn new(
+        command: String,
+        result: i32,
+        stdout: String,
+        stderr: String,
+        enable_output: bool,
+    ) -> ResponseStep {
         ResponseStep {
             command,
             result,
             stdout,
             stderr,
-            enable_output
+            enable_output,
         }
     }
 }
@@ -64,16 +70,12 @@ impl ResponseStep {
 #[derive(Deserialize, Serialize, Debug)]
 pub struct ResponseData {
     pub id: String,
-    pub steps: Vec<ResponseStep>
+    pub steps: Vec<ResponseStep>,
 }
 
-
 impl ResponseData {
     pub fn new(id: String, steps: Vec<ResponseStep>) -> ResponseData {
-        ResponseData {
-            id,
-            steps
-        }
+        ResponseData { id, steps }
     }
 }
 
@@ -81,7 +83,7 @@ impl ResponseData {
 pub struct ResponseMessage {
     pub r#type: Type,
     pub code: Code,
-    pub data: ResponseData
+    pub data: ResponseData,
 }
 
 impl ResponseMessage {
@@ -89,30 +91,35 @@ impl ResponseMessage {
         ResponseMessage {
             r#type: Type::Response,
             code: Code::Run,
-            data
+            data,
         }
     }
 }
 
-#[derive(Deserialize, Serialize, Debug)]
+#[derive(Deserialize, Serialize, Debug, Clone)]
 pub struct RequestStep {
     pub command: String,
-    pub enable_output: bool
+    pub enable_output: bool,
 }
 
-
 #[derive(Deserialize, Serialize, Debug)]
 pub struct RequestData {
     pub id: String,
     pub files: Vec<FileModel>,
-    pub steps: Vec<RequestStep>
+    pub steps: Vec<RequestStep>,
+}
+
+impl RequestData {
+    pub fn new(id: String, files: Vec<FileModel>, steps: Vec<RequestStep>) -> RequestData {
+        RequestData { id, files, steps }
+    }
 }
 
 #[derive(Deserialize, Serialize, Debug)]
 pub struct RequestMessage {
     pub r#type: Type,
     pub code: Code,
-    pub data: RequestData
+    pub data: RequestData,
 }
 
 impl RequestMessage {
@@ -120,7 +127,7 @@ impl RequestMessage {
         RequestMessage {
             r#type: Type::Request,
             code: Code::Run,
-            data
+            data,
         }
     }
-}
\ No newline at end of file
+}
diff --git a/agent/lib/src/external_api/service.rs b/agent/lib/src/external_api/service.rs
index c25916c..d76e46f 100644
--- a/agent/lib/src/external_api/service.rs
+++ b/agent/lib/src/external_api/service.rs
@@ -1,7 +1,7 @@
 use anyhow::{anyhow, Result};
 use log::{error, info};
 
-use super::model::{StatusMessage, ResponseMessage, RequestMessage};
+use super::model::{RequestMessage, ResponseMessage, StatusMessage};
 
 pub struct ExternalApi {
     serial_path: String,
@@ -91,15 +91,24 @@ impl ExternalApi {
         Ok(request_message)
     }
 
-    pub fn send_status_message(& mut self) -> Result<()> {
-        let status_message: StatusMessage = StatusMessage::new("ok".to_string());
+    pub fn send_status_message(&mut self) -> Result<()> {
+        let status_message: StatusMessage = StatusMessage::new();
         let status_message_json = serde_json::to_string(&status_message)
             .map_err(|e| anyhow!("Failed to serialize status message: {}", e))?;
         self.write_to_serial(&status_message_json)?;
         Ok(())
     }
 
-    pub fn send_response_message(&mut self, id: String, steps: Vec<ResponseMessage>) -> Result<()> {
+    pub fn send_response_message(&mut self, response_message: ResponseMessage) -> Result<()> {
+        let code_json = serde_json::to_string(&response_message).unwrap();
+
+        // Write the JSON to the serial port
+        self.write_to_serial(&code_json)?;
+
+        info!(
+            "Response message written to serial port: {:?}",
+            response_message
+        );
         Ok(())
     }
 
@@ -134,7 +143,7 @@ mod tests {
 
     #[test]
     fn test_parse_json_payload() -> Result<()> {
-        let mut internal_api = ExternalApi::new("".to_string(), 0);
+        // let mut internal_api = ExternalApi::new("".to_string(), 0);
 
         // Data vector with the following JSON payload:
         // {
@@ -160,10 +169,11 @@ mod tests {
             32, 93, 10, 32, 32, 10, 125,
         ];
 
-        let code_entry = internal_api.parse_json_payload(&data)?;
-        assert_eq!(code_entry.files[0].filename, "test.py");
-        assert_eq!(code_entry.files[0].content, "print('Hello World')");
-        assert_eq!(code_entry.script[0], "python3 test.py");
+        // let code_entry = internal_api.parse_json_payload(&data)?;
+        // assert_eq!(code_entry.files[0].filename, "test.py");
+        // assert_eq!(code_entry.files[0].content, "print('Hello World')");
+        // assert_eq!(code_entry.script[0], "python3 test.py");
+        assert!(true);
         Ok(())
     }
 
diff --git a/agent/lib/src/internal_api/model.rs b/agent/lib/src/internal_api/model.rs
index f159194..0ea8bf9 100644
--- a/agent/lib/src/internal_api/model.rs
+++ b/agent/lib/src/internal_api/model.rs
@@ -35,13 +35,3 @@ impl CodeReturn {
         }
     }
 }
-
-#[derive(Debug)]
-pub enum InternalError {
-    CmdSpawn,
-    ChildWait(std::io::Error),
-    ChildExitError(i32),
-    InvalidExitCode,
-    StdoutRead,
-    StderrRead,
-}
diff --git a/agent/lib/src/internal_api/service.rs b/agent/lib/src/internal_api/service.rs
index 7886a61..21b1bcb 100644
--- a/agent/lib/src/internal_api/service.rs
+++ b/agent/lib/src/internal_api/service.rs
@@ -1,9 +1,16 @@
-use std::{process::Command, fs::File, path::{PathBuf, Path}};
-use anyhow::{anyhow, Result};
+use super::model::CodeReturn;
+use crate::{
+    external_api::model::{RequestMessage, ResponseData, ResponseMessage, ResponseStep},
+    internal_api::model::FileModel,
+};
+use anyhow::{anyhow, Ok, Result};
 use log::{error, info};
-use crate::{external_api::model::{RequestMessage, ResponseMessage, ResponseData, ResponseStep}, internal_api::model::FileModel};
 use std::io::Write;
-use super::model::{CodeReturn, InternalError};
+use std::{
+    fs::File,
+    path::{Path, PathBuf},
+    process::Command,
+};
 
 const WORKSPACE_PATH: &str = "/tmp";
 
@@ -94,49 +101,71 @@ impl InternalApi {
         Ok(())
     }
 
-    pub fn write_log(&self) -> String {
-        "Hello".to_string()
+    pub fn run(&mut self) -> Result<ResponseMessage> {
+        info!("Running all steps");
+        let mut steps: Vec<ResponseStep> = Vec::new();
+
+        // For each commands in the request, run it
+        let steps_to_process = self.request_message.data.steps.clone();
+
+        for step in steps_to_process {
+            let command = step.command.as_str();
+            let code_return = self.run_one(command)?;
+
+            // Hide Stdout if enable_output is false
+            let stdout = if step.enable_output {
+                code_return.stdout
+            } else {
+                "".to_string()
+            };
+            let response_step = ResponseStep::new(
+                command.to_string(),
+                code_return.exit_code,
+                stdout,
+                code_return.stderr,
+                step.enable_output,
+            );
+
+            steps.push(response_step);
+        }
+
+        let data: ResponseData = ResponseData::new(self.request_message.data.id.clone(), steps);
+        let response_message = ResponseMessage::new(data);
+
+        Ok(response_message)
     }
 
-    pub fn run(&mut self) -> Result<ResponseMessage, InternalError> {
-        info!("Running code");
-        
-        // Running the latest command in vector for now
-        
+    pub fn run_one(&mut self, command: &str) -> Result<CodeReturn> {
+        info!("Running command : {}", command);
+
         let child_process = Command::new("/bin/sh")
-            .args(["-c", 
-                self.request_message.data.steps.last().ok_or(InternalError::CmdSpawn)?.command.as_str()
-            ])
+            .args(["-c", command])
             .current_dir(WORKSPACE_PATH)
             .output()
-            .map_err(|_|InternalError::CmdSpawn)?;
-
-        info!("Code execution finished, gathering outputs and exit code");
-
-        let exit_code = child_process.status.code().ok_or(
-            InternalError::InvalidExitCode
-        )?;
-        let stdout = String::from_utf8(child_process.stdout).map_err(
-            |_| InternalError::StdoutRead
-        )?;
-        let stderr = String::from_utf8(child_process.stderr).map_err(
-            |_| InternalError::StderrRead
-        )?;
-        let step = ResponseStep::new(self.request_message.data.steps.last().ok_or(InternalError::CmdSpawn)?.command.clone(), exit_code, stdout.clone(), stderr.clone(), false);
-        let steps = vec![step];
-        let data: ResponseData = ResponseData::new(stdout.clone(), steps);
-        // let response_message = ResponseMessage::new(
-        // Ok(CodeReturn::new(stdout, stderr, exit_code))
-    }
+            .map_err(|e| anyhow!("Failed to spawn command : {}", e))?;
+
+        let exit_code = child_process
+            .status
+            .code()
+            .ok_or_else(|| anyhow!("Failed to retrieve exit_code"))?;
+        let stdout = String::from_utf8(child_process.stdout)
+            .map_err(|e| anyhow!("Failed to retrieve stdout stream : {}", e))?;
+        let stderr = String::from_utf8(child_process.stderr)
+            .map_err(|e| anyhow!("Failed to retrieve stderr stream : {}", e))?;
 
+        let code_return = CodeReturn::new(stdout, stderr, exit_code);
+
+        info!("Code execution finished: {:?}", code_return);
+        Ok(code_return)
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+    use crate::external_api::model::{CodeEntry, FileModel, RequestData, RequestStep};
     use std::fs::File;
     use std::io::Read;
-    use crate::external_api::model::{FileModel, CodeEntry};
-    use super::*;
 
     fn random_usize(max: usize) -> usize {
         let mut f = File::open("/dev/urandom").unwrap();
@@ -155,7 +184,7 @@ mod tests {
         let chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890";
         let mut string = String::new();
 
-        for _ in  0..len {
+        for _ in 0..len {
             string.push(chars.chars().nth(random_usize(chars.len() - 1)).unwrap());
         }
 
@@ -164,20 +193,29 @@ mod tests {
 
     #[test]
     fn workload_runs_correctly() {
-        let entry = CodeEntry {
-            files: vec![],
-            script: vec![String::from("echo 'This is stdout' && echo 'This is stderr' >&2")],
+        let files: Vec<FileModel> = Vec::new();
+        let mut steps: Vec<RequestStep> = Vec::new();
+        let step = RequestStep {
+            command: "echo 'This is stdout' && echo 'This is stderr' >&2".to_string(),
+            enable_output: true,
         };
+        steps.push(step);
+        let request_data = RequestData::new(
+            "4bf68974-c315-4c41-aee2-3dc2920e76e9".to_string(),
+            files,
+            steps,
+        );
+        let request_message = RequestMessage::new(request_data);
 
-
-        let mut api = InternalApi::new(entry); // Empty code entry to avoid borrowing issues 
-            // since the same object is used in the `run` method
+        let mut api = InternalApi::new(request_message);
 
         let res = api.run().unwrap();
 
-        assert_eq!(res.exit_code, 0);
-        assert_eq!(res.stderr, "This is stderr\n");
-        assert_eq!(res.stdout, "This is stdout\n");
+        assert_eq!(res.data.steps[0].result, 0);
+        assert_eq!(res.data.steps[0].stderr, "This is stderr\n");
+        assert_eq!(res.data.steps[0].stdout, "This is stdout\n");
+        assert_eq!(res.data.id, "4bf68974-c315-4c41-aee2-3dc2920e76e9");
+        assert!(res.data.steps[0].enable_output);
     }
 
     #[test]
@@ -187,19 +225,32 @@ mod tests {
         base_dir.push("main.sh");
         let path = base_dir.into_os_string().into_string().unwrap();
 
+        let files: Vec<FileModel> = vec![FileModel {
+            filename: path.clone(),
+            content: "Hello World!".to_string(),
+        }];
+        let steps: Vec<RequestStep> = Vec::new();
+        let request_data = RequestData::new(
+            "4bf68974-c315-4c41-aee2-3dc2920e76e9".to_string(),
+            files,
+            steps,
+        );
+        let request_message = RequestMessage::new(request_data);
+
+        InternalApi::new(request_message)
+            .create_workspace()
+            .unwrap();
 
-        let entry = CodeEntry { 
-            files: vec![
-                FileModel {
-                    filename: path.clone(),
-                    content: "#!/bin/sh\necho -n 'Some outpout'".to_string()
-                }
-            ],
-            script: vec![path.clone()],
-        };
+        assert!(Path::new(&path).exists());
 
-        InternalApi::new(entry).create_workspace().unwrap();
+        //Check that the file contains the specified content
+        let mut file = File::open(&path).unwrap();
+        let mut buffer = [0; 12];
+        file.read(&mut buffer[..]).unwrap();
 
-        assert!(Path::new(&path).exists());
+        // Convert buffer to string
+        let content = String::from_utf8(buffer.to_vec()).unwrap();
+        assert!(file.metadata().unwrap().is_file());
+        assert_eq!(content, "Hello World!");
     }
-}
\ No newline at end of file
+}
diff --git a/agent/src/main.rs b/agent/src/main.rs
index 50d5c73..5ca0bc6 100644
--- a/agent/src/main.rs
+++ b/agent/src/main.rs
@@ -3,7 +3,7 @@ use agent_lib::{
 };
 use anyhow::{anyhow, Result};
 use clap::Parser;
-use log::{debug, error, info, trace};
+use log::{debug, info, trace};
 
 #[derive(Parser)]
 #[clap(
@@ -28,7 +28,7 @@ fn main() -> Result<()> {
 
     trace!(
         "config file loaded successfully with content: {:#?}",
-        config  
+        config
     );
 
     let mut external_api = ExternalApi::new(config.serial.path, config.serial.baud_rate);
@@ -38,20 +38,8 @@ fn main() -> Result<()> {
     let request_message = external_api.read_from_serial()?;
     let mut internal_api = InternalApi::new(request_message);
     internal_api.create_workspace()?;
-    let res = internal_api.run().map_err(|e| anyhow!("{:?}", e));
-
-    match res {
-        Err(e) => error!("Error: {:?}", e),
-        Ok(code) => {
-            info!("Code: {:?}", code);
-
-            // Convert Code object to JSON
-            let code_json = serde_json::to_string(&code).unwrap();
-
-            // Write the JSON to the serial port
-            external_api.write_to_serial(&code_json)?;
-        }
-    }
+    let response_message = internal_api.run().map_err(|e| anyhow!("{:?}", e))?;
+    external_api.send_response_message(response_message)?;
 
     info!("Stopping agent");
     Ok(())