Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion implants/imix/src/install.rs
Original file line number Diff line number Diff line change
@@ -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)]
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion implants/imix/src/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
45 changes: 20 additions & 25 deletions implants/lib/eldritch/src/assets/copy_impl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()) {
Expand All @@ -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)
Expand All @@ -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))?;
Expand Down Expand Up @@ -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 =
Expand All @@ -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(())
Expand All @@ -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(
Expand Down Expand Up @@ -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(())
Expand Down