Split stdlib into seperate impl files. (#1538)

* Split stdlib into seperate impl files.

* fmt

Author: hulto

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/claim_tasks_impl.rs

@@ -0,0 +1,16 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+use alloc::vec::Vec;
+
+use crate::TaskWrapper;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+#[cfg(feature = "stdlib")]
+use pb::c2;
+
+pub fn claim_tasks(agent: Arc<dyn Agent>) -> Result<Vec<TaskWrapper>, String> {
+    let req = c2::ClaimTasksRequest { beacon: None };
+    let resp = agent.claim_tasks(req)?;
+    Ok(resp.tasks.into_iter().map(TaskWrapper).collect())
+}

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/fetch_asset_impl.rs

@@ -0,0 +1,13 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+use alloc::vec::Vec;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+#[cfg(feature = "stdlib")]
+use pb::c2;
+
+pub fn fetch_asset(agent: Arc<dyn Agent>, name: String) -> Result<Vec<u8>, String> {
+    let req = c2::FetchAssetRequest { name };
+    agent.fetch_asset(req)
+}

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/get_callback_interval_impl.rs

@@ -0,0 +1,9 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+
+pub fn get_callback_interval(agent: Arc<dyn Agent>) -> Result<i64, String> {
+    agent.get_callback_interval().map(|i| i as i64)
+}

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/get_config_impl.rs

@@ -0,0 +1,23 @@
+use alloc::collections::BTreeMap;
+use alloc::string::String;
+use alloc::sync::Arc;
+use eldritch_core::Value;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+
+pub fn get_config(agent: Arc<dyn Agent>) -> Result<BTreeMap<String, Value>, String> {
+    let config = agent.get_config()?;
+    let mut result = BTreeMap::new();
+    for (k, v) in config {
+        // Try to parse numbers, otherwise keep as string
+        if let Ok(i) = v.parse::<i64>() {
+            result.insert(k, Value::Int(i));
+        } else if let Ok(b) = v.parse::<bool>() {
+            result.insert(k, Value::Bool(b));
+        } else {
+            result.insert(k, Value::String(v));
+        }
+    }
+    Ok(result)
+}

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/get_transport_impl.rs

@@ -0,0 +1,9 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+
+pub fn get_transport(agent: Arc<dyn Agent>) -> Result<String, String> {
+    agent.get_transport()
+}

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/list_tasks_impl.rs

@@ -0,0 +1,13 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+use alloc::vec::Vec;
+
+use crate::TaskWrapper;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+
+pub fn list_tasks(agent: Arc<dyn Agent>) -> Result<Vec<TaskWrapper>, String> {
+    let tasks = agent.list_tasks()?;
+    Ok(tasks.into_iter().map(TaskWrapper).collect())
+}

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/list_transports_impl.rs

@@ -0,0 +1,10 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+use alloc::vec::Vec;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+
+pub fn list_transports(agent: Arc<dyn Agent>) -> Result<Vec<String>, String> {
+    agent.list_transports()
+}

…chv2/stdlib/eldritch-libagent/src/std.rs …/stdlib/eldritch-libagent/src/std/mod.rsimplants/lib/eldritchv2/stdlib/eldritch-libagent/src/std.rs renamed to implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/mod.rs implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std.rs renamed to implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/mod.rs

@@ -10,8 +10,22 @@ use crate::{CredentialWrapper, FileWrapper, ProcessListWrapper, TaskWrapper};
 
 #[cfg(feature = "stdlib")]
 use crate::agent::Agent;
-#[cfg(feature = "stdlib")]
-use pb::c2;
+
+pub mod claim_tasks_impl;
+pub mod fetch_asset_impl;
+pub mod get_callback_interval_impl;
+pub mod get_config_impl;
+pub mod get_transport_impl;
+pub mod list_tasks_impl;
+pub mod list_transports_impl;
+pub mod report_credential_impl;
+pub mod report_file_impl;
+pub mod report_process_list_impl;
+pub mod report_task_output_impl;
+pub mod set_callback_interval_impl;
+pub mod set_callback_uri_impl;
+pub mod stop_task_impl;
+pub mod terminate_impl;
 
 // We need manual Debug impl, and we need to put the macro on the struct.
 #[eldritch_library_impl(AgentLibrary)]
@@ -36,103 +50,65 @@ impl StdAgentLibrary {
 
 impl AgentLibrary for StdAgentLibrary {
     fn get_config(&self) -> Result<BTreeMap<String, Value>, String> {
-        let config = self.agent.get_config()?;
-        let mut result = BTreeMap::new();
-        for (k, v) in config {
-            // Try to parse numbers, otherwise keep as string
-            if let Ok(i) = v.parse::<i64>() {
-                result.insert(k, Value::Int(i));
-            } else if let Ok(b) = v.parse::<bool>() {
-                result.insert(k, Value::Bool(b));
-            } else {
-                result.insert(k, Value::String(v));
-            }
-        }
-        Ok(result)
+        get_config_impl::get_config(self.agent.clone())
     }
 
     fn _terminate_this_process_clowntown(&self) -> Result<(), String> {
-        ::std::process::exit(0);
+        terminate_impl::terminate()
     }
 
     fn set_callback_interval(&self, interval: i64) -> Result<(), String> {
-        self.agent.set_callback_interval(interval as u64)
+        set_callback_interval_impl::set_callback_interval(self.agent.clone(), interval)
     }
 
     // Interactivity
     fn fetch_asset(&self, name: String) -> Result<Vec<u8>, String> {
-        let req = c2::FetchAssetRequest { name };
-        self.agent.fetch_asset(req)
+        fetch_asset_impl::fetch_asset(self.agent.clone(), name)
     }
 
     fn report_credential(&self, credential: CredentialWrapper) -> Result<(), String> {
-        let req = c2::ReportCredentialRequest {
-            task_id: self.task_id,
-            credential: Some(credential.0),
-        };
-        self.agent.report_credential(req).map(|_| ())
+        report_credential_impl::report_credential(self.agent.clone(), self.task_id, credential)
     }
 
     fn report_file(&self, file: FileWrapper) -> Result<(), String> {
-        let req = c2::ReportFileRequest {
-            task_id: self.task_id,
-            chunk: Some(file.0),
-        };
-        self.agent.report_file(req).map(|_| ())
+        report_file_impl::report_file(self.agent.clone(), self.task_id, file)
     }
 
     fn report_process_list(&self, list: ProcessListWrapper) -> Result<(), String> {
-        let req = c2::ReportProcessListRequest {
-            task_id: self.task_id,
-            list: Some(list.0),
-        };
-        self.agent.report_process_list(req).map(|_| ())
+        report_process_list_impl::report_process_list(self.agent.clone(), self.task_id, list)
     }
 
     fn report_task_output(&self, output: String, error: Option<String>) -> Result<(), String> {
-        let task_error = error.map(|msg| c2::TaskError { msg });
-        let output_msg = c2::TaskOutput {
-            id: self.task_id,
-            output,
-            error: task_error,
-            exec_started_at: None,
-            exec_finished_at: None,
-        };
-        let req = c2::ReportTaskOutputRequest {
-            output: Some(output_msg),
-        };
-        self.agent.report_task_output(req).map(|_| ())
+        report_task_output_impl::report_task_output(self.agent.clone(), self.task_id, output, error)
     }
 
     fn claim_tasks(&self) -> Result<Vec<TaskWrapper>, String> {
-        let req = c2::ClaimTasksRequest { beacon: None };
-        let resp = self.agent.claim_tasks(req)?;
-        Ok(resp.tasks.into_iter().map(TaskWrapper).collect())
+        claim_tasks_impl::claim_tasks(self.agent.clone())
     }
 
     // Agent Configuration
     fn get_transport(&self) -> Result<String, String> {
-        self.agent.get_transport()
+        get_transport_impl::get_transport(self.agent.clone())
     }
 
     fn list_transports(&self) -> Result<Vec<String>, String> {
-        self.agent.list_transports()
+        list_transports_impl::list_transports(self.agent.clone())
     }
+
     fn set_callback_uri(&self, uri: String) -> Result<(), String> {
-        self.agent.set_callback_uri(uri)
+        set_callback_uri_impl::set_callback_uri(self.agent.clone(), uri)
     }
 
     fn get_callback_interval(&self) -> Result<i64, String> {
-        self.agent.get_callback_interval().map(|i| i as i64)
+        get_callback_interval_impl::get_callback_interval(self.agent.clone())
     }
 
     // Task Management
     fn list_tasks(&self) -> Result<Vec<TaskWrapper>, String> {
-        let tasks = self.agent.list_tasks()?;
-        Ok(tasks.into_iter().map(TaskWrapper).collect())
+        list_tasks_impl::list_tasks(self.agent.clone())
     }
 
     fn stop_task(&self, task_id: i64) -> Result<(), String> {
-        self.agent.stop_task(task_id)
+        stop_task_impl::stop_task(self.agent.clone(), task_id)
     }
 }

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/report_credential_impl.rs

@@ -0,0 +1,21 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+
+use crate::CredentialWrapper;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+#[cfg(feature = "stdlib")]
+use pb::c2;
+
+pub fn report_credential(
+    agent: Arc<dyn Agent>,
+    task_id: i64,
+    credential: CredentialWrapper,
+) -> Result<(), String> {
+    let req = c2::ReportCredentialRequest {
+        task_id,
+        credential: Some(credential.0),
+    };
+    agent.report_credential(req).map(|_| ())
+}

implants/lib/eldritchv2/stdlib/eldritch-libagent/src/std/report_file_impl.rs

@@ -0,0 +1,17 @@
+use alloc::string::String;
+use alloc::sync::Arc;
+
+use crate::FileWrapper;
+
+#[cfg(feature = "stdlib")]
+use crate::agent::Agent;
+#[cfg(feature = "stdlib")]
+use pb::c2;
+
+pub fn report_file(agent: Arc<dyn Agent>, task_id: i64, file: FileWrapper) -> Result<(), String> {
+    let req = c2::ReportFileRequest {
+        task_id,
+        chunk: Some(file.0),
+    };
+    agent.report_file(req).map(|_| ())
+}