From 239c66591cca2f623f270392b472aacefed4f365 Mon Sep 17 00:00:00 2001
From: KCarretto <Kcarretto@gmail.com>
Date: Wed, 21 Feb 2024 22:12:47 +0000
Subject: [PATCH] remove file for assets.copy() instead of truncating

---
 implants/imix/src/install.rs                  |  4 +-
 implants/imix/src/task.rs                     |  2 +-
 implants/lib/eldritch/src/assets/copy_impl.rs | 45 +++++++++----------
 3 files changed, 24 insertions(+), 27 deletions(-)

diff --git a/implants/imix/src/install.rs b/implants/imix/src/install.rs
index f0d614d8c..ef8367726 100644
--- a/implants/imix/src/install.rs
+++ b/implants/imix/src/install.rs
@@ -1,6 +1,8 @@
 use anyhow::{anyhow, Result};
 use pb::eldritch::Tome;
 use std::collections::HashMap;
+#[cfg(debug_assertions)]
+use std::fmt::Write;
 
 pub async fn install() {
     #[cfg(debug_assertions)]
@@ -47,7 +49,7 @@ pub async fn install() {
 
             #[cfg(debug_assertions)]
             for msg in runtime.collect() {
-                if let Message::ReportText(m) = msg {
+                if let eldritch::runtime::Message::ReportText(m) = msg {
                     if let Err(err) = output.write_str(m.text().as_str()) {
                         #[cfg(debug_assertions)]
                         log::error!("failed to write text: {}", err);
diff --git a/implants/imix/src/task.rs b/implants/imix/src/task.rs
index d5dc30f55..5d2b3ed36 100644
--- a/implants/imix/src/task.rs
+++ b/implants/imix/src/task.rs
@@ -65,7 +65,7 @@ impl TaskHandle {
                         match t
                             .report_task_output(ReportTaskOutputRequest {
                                 output: Some(TaskOutput {
-                                    id: id,
+                                    id,
                                     output: String::new(),
                                     error: Some(TaskError {
                                         msg: format!("dispatch error ({}): {:#?}", msg_str, err),
diff --git a/implants/lib/eldritch/src/assets/copy_impl.rs b/implants/lib/eldritch/src/assets/copy_impl.rs
index 19754ab8a..19decf67a 100644
--- a/implants/lib/eldritch/src/assets/copy_impl.rs
+++ b/implants/lib/eldritch/src/assets/copy_impl.rs
@@ -2,10 +2,13 @@ use crate::runtime::{messages::FetchAssetMessage, Environment};
 use anyhow::{Context, Result};
 use pb::c2::FetchAssetResponse;
 use starlark::{eval::Evaluator, values::list::ListRef};
-use std::fs::OpenOptions;
-use std::io::Write;
-use std::sync::mpsc::channel;
-use std::{fs, sync::mpsc::Receiver};
+use std::{
+    fs,
+    fs::OpenOptions,
+    io::Write,
+    path::Path,
+    sync::mpsc::{channel, Receiver},
+};
 
 fn copy_local(src: String, dst: String) -> Result<()> {
     let src_file = match super::Asset::get(src.as_str()) {
@@ -23,20 +26,12 @@ fn copy_remote(rx: Receiver<FetchAssetResponse>, dst_path: String) -> Result<()>
     // Wait for our first chunk
     let resp = rx.recv()?;
 
-    // Truncate file
-    let mut dst = OpenOptions::new()
-        .create(true)
-        .truncate(true)
-        .write(true)
-        .open(&dst_path)
-        .context(format!(
-            "failed to truncate destination file: {}",
-            &dst_path
-        ))?;
-    dst.flush()
-        .context(format!("failed to flush file truncation: {}", &dst_path))?;
+    // Delete file if it exists
+    if Path::new(&dst_path).exists() {
+        fs::remove_file(&dst_path).context("failed to delete existing file")?;
+    }
 
-    // Reopen file for writing
+    // Open file for writing
     let mut dst = OpenOptions::new()
         .create(true)
         .append(true)
@@ -47,7 +42,7 @@ fn copy_remote(rx: Receiver<FetchAssetResponse>, dst_path: String) -> Result<()>
     dst.write_all(&resp.chunk)
         .context(format!("failed to write file chunk: {}", &dst_path))?;
 
-    // Listen for downloaded chunks and write them
+    // Listen for more chunks and write them
     for resp in rx {
         dst.write_all(&resp.chunk)
             .context(format!("failed to write file chunk: {}", &dst_path))?;
@@ -87,15 +82,16 @@ mod tests {
     };
     use pb::c2::FetchAssetResponse;
     use pb::eldritch::Tome;
-    use std::sync::mpsc::channel;
     use std::{collections::HashMap, io::prelude::*};
+    use std::{fs, sync::mpsc::channel};
     use tempfile::NamedTempFile;
 
     #[tokio::test]
     async fn test_remote_copy() -> anyhow::Result<()> {
         // Create files
-        let mut tmp_file_dst = NamedTempFile::new()?;
+        let tmp_file_dst = NamedTempFile::new()?;
         let path_dst = String::from(tmp_file_dst.path().to_str().unwrap());
+        let check_dst = path_dst.clone();
 
         let (tx, rx) = channel();
         let handle =
@@ -113,8 +109,7 @@ mod tests {
 
         handle.await?;
 
-        let mut contents = String::new();
-        tmp_file_dst.read_to_string(&mut contents)?;
+        let contents = fs::read_to_string(check_dst)?;
         assert!(contents.contains("Hello from a remote asset"));
         assert!(contents.contains("Goodbye from a remote asset"));
         Ok(())
@@ -123,8 +118,9 @@ mod tests {
     #[tokio::test]
     async fn test_remote_copy_full() -> anyhow::Result<()> {
         // Create files
-        let mut tmp_file_dst = NamedTempFile::new()?;
+        let tmp_file_dst = NamedTempFile::new()?;
         let path_dst = String::from(tmp_file_dst.path().to_str().unwrap());
+        let check_dst = path_dst.clone();
 
         // Run Eldritch (in it's own thread)
         let mut runtime = crate::start(
@@ -185,8 +181,7 @@ mod tests {
         runtime.finish().await;
 
         // Lastly, assert the file was written correctly
-        let mut contents = String::new();
-        tmp_file_dst.read_to_string(&mut contents)?;
+        let contents = fs::read_to_string(check_dst)?;
         assert_eq!("chunk1\nchunk2\n", contents.as_str());
 
         Ok(())