diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index efc3221e5..d5b6a74c1 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -34,7 +34,7 @@ jobs: uses: codecov/codecov-action@v3 implants: runs-on: ${{ matrix.os }} - timeout-minutes: 30 + timeout-minutes: 60 strategy: matrix: os: diff --git a/implants/imix/src/task.rs b/implants/imix/src/task.rs index 7a79b3a3f..7b3babe91 100644 --- a/implants/imix/src/task.rs +++ b/implants/imix/src/task.rs @@ -158,14 +158,18 @@ impl TaskHandle { let resp = match file_chunk.recv() { Ok(r) => r, Err(_err) => { - #[cfg(debug_assertions)] - log::error!( - "failed to download file chunk: task_id={}, name={}: {}", - task_id, - req.name(), - _err - ); - + match _err.to_string().as_str() { + "receiving on a closed channel" => {} + _ => { + #[cfg(debug_assertions)] + log::error!( + "failed to download file chunk: task_id={}, name={}: {}", + task_id, + req.name(), + _err + ); + } + } return; } }; diff --git a/implants/lib/eldritch/Cargo.toml b/implants/lib/eldritch/Cargo.toml index 3ee30108b..28afd0306 100644 --- a/implants/lib/eldritch/Cargo.toml +++ b/implants/lib/eldritch/Cargo.toml @@ -30,6 +30,7 @@ nix = { workspace = true } notify = { workspace = true } object = { workspace = true } openssl = { workspace = true, features = ["vendored"] } +pretty_env_logger = { workspace = true } prost = { workspace = true} prost-types = { workspace = true } regex = { workspace = true } diff --git a/implants/lib/eldritch/src/assets.rs b/implants/lib/eldritch/src/assets.rs index c04274699..04439be51 100644 --- a/implants/lib/eldritch/src/assets.rs +++ b/implants/lib/eldritch/src/assets.rs @@ -6,16 +6,16 @@ mod read_impl; use allocative::Allocative; use derive_more::Display; +use rust_embed::RustEmbed; +use serde::{Serialize, Serializer}; use starlark::environment::{Methods, MethodsBuilder, MethodsStatic}; +use starlark::eval::Evaluator; use starlark::values::none::NoneType; use starlark::values::{ starlark_value, ProvidesStaticType, StarlarkValue, UnpackValue, Value, ValueLike, }; use starlark::{starlark_module, starlark_simple_value}; -use rust_embed::RustEmbed; -use serde::{Serialize, Serializer}; - #[cfg(debug_assertions)] #[derive(RustEmbed)] #[folder = "../../../bin/embedded_files_test"] @@ -71,9 +71,9 @@ impl<'v> UnpackValue<'v> for AssetsLibrary { #[rustfmt::skip] #[allow(clippy::needless_lifetimes, clippy::type_complexity, clippy::too_many_arguments)] fn methods(builder: &mut MethodsBuilder) { - fn copy(this: AssetsLibrary, src: String, dest: String) -> anyhow::Result { + fn copy<'v>(this: AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>, src: String, dest: String) -> anyhow::Result { if false { println!("Ignore unused this var. _this isn't allowed by starlark. {:?}", this); } - copy_impl::copy(src, dest)?; + copy_impl::copy(starlark_eval, src, dest)?; Ok(NoneType{}) } fn list(this: AssetsLibrary) -> anyhow::Result> { diff --git a/implants/lib/eldritch/src/assets/copy_impl.rs b/implants/lib/eldritch/src/assets/copy_impl.rs index 17325d494..e50beabcc 100644 --- a/implants/lib/eldritch/src/assets/copy_impl.rs +++ b/implants/lib/eldritch/src/assets/copy_impl.rs @@ -1,7 +1,9 @@ -use anyhow::Result; -use std::fs; +use crate::runtime::Client; +use anyhow::{Context, Result}; +use starlark::{eval::Evaluator, values::list::ListRef}; +use std::{fs, sync::mpsc::Receiver}; -pub fn copy(src: String, dst: String) -> Result<()> { +fn copy_local(src: String, dst: String) -> Result<()> { let src_file = match super::Asset::get(src.as_str()) { Some(local_src_file) => local_src_file.data, None => return Err(anyhow::anyhow!("Embedded file {src} not found.")), @@ -13,28 +15,139 @@ pub fn copy(src: String, dst: String) -> Result<()> { } } +fn copy_remote(file_reciever: Receiver>, dst: String) -> Result<()> { + loop { + let val = match file_reciever.recv() { + Ok(v) => v, + Err(err) => { + match err.to_string().as_str() { + "channel is empty and sending half is closed" => { + break; + } + "timed out waiting on channel" => { + continue; + } + _ => { + #[cfg(debug_assertions)] + log::debug!("failed to drain channel: {}", err) + } + } + break; + } + }; + match fs::write(dst.clone(), val) { + Ok(_) => {} + Err(local_err) => return Err(local_err.try_into()?), + }; + } + + Ok(()) +} + +pub fn copy(starlark_eval: &mut Evaluator<'_, '_>, src: String, dst: String) -> Result<()> { + let remote_assets = starlark_eval.module().get("remote_assets"); + + if let Some(assets) = remote_assets { + let tmp_list = ListRef::from_value(assets).context("`remote_assets` is not type list")?; + let src_value = starlark_eval.module().heap().alloc_str(&src); + + if tmp_list.contains(&src_value.to_value()) { + let client = Client::from_extra(starlark_eval.extra)?; + let file_reciever = client.request_file(src)?; + + return copy_remote(file_reciever, dst); + } + } + copy_local(src, dst) +} + #[cfg(test)] mod tests { - use super::*; - use std::io::prelude::*; + use crate::Runtime; + + use std::{collections::HashMap, io::prelude::*}; use tempfile::NamedTempFile; + // fn init_log() { + // pretty_env_logger::formatted_timed_builder() + // .filter_level(log::LevelFilter::Info) + // .parse_env("IMIX_LOG") + // .init(); + // } + + // #[tokio::test] + // async fn test_remote_copy() -> anyhow::Result<()> { + // // Create files + // let mut tmp_file_dst = NamedTempFile::new()?; + // let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); + + // let (sender, reciver) = channel::>(); + // sender.send("Hello from a remote asset".as_bytes().to_vec())?; + + // copy_remote(reciver, path_dst)?; + + // let mut contents = String::new(); + // tmp_file_dst.read_to_string(&mut contents)?; + // assert!(contents.contains("Hello from a remote asset")); + // Ok(()) + // } + + // #[tokio::test] + // async fn test_remote_copy_full() -> anyhow::Result<()> { + // init_log(); + // log::debug!("Testing123"); + + // // Create files + // let mut tmp_file_dst = NamedTempFile::new()?; + // let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); + + // let (runtime, broker) = Runtime::new(); + // let handle = tokio::task::spawn_blocking(move || { + // runtime.run(crate::pb::Tome { + // eldritch: r#"assets.copy("test_tome/test_file.txt", input_params['test_output'])"# + // .to_owned(), + // parameters: HashMap::from([("test_output".to_string(), path_dst)]), + // file_names: Vec::from(["test_tome/test_file.txt".to_string()]), + // }) + // }); + // handle.await?; + // println!("{:?}", broker.collect_file_requests().len()); + // assert!(broker.collect_errors().is_empty()); // No errors even though the remote asset is inaccessible + + // let mut contents = String::new(); + // tmp_file_dst.read_to_string(&mut contents)?; + // // Compare - Should be empty basically just didn't error + // assert!(contents.contains("")); + + // Ok(()) + // } + #[test] fn test_embedded_copy() -> anyhow::Result<()> { // Create files let mut tmp_file_dst = NamedTempFile::new()?; let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); - // Run our code #[cfg(any(target_os = "linux", target_os = "macos"))] - copy("exec_script/hello_world.sh".to_string(), path_dst)?; + let path_src = "exec_script/hello_world.sh".to_string(); #[cfg(target_os = "windows")] - copy("exec_script/hello_world.bat".to_string(), path_dst)?; + let path_src = "exec_script/hello_world.bat".to_string(); + + let (runtime, broker) = Runtime::new(); + runtime.run(crate::pb::Tome { + eldritch: r#"assets.copy(input_params['src_file'], input_params['test_output'])"# + .to_owned(), + parameters: HashMap::from([ + ("src_file".to_string(), path_src), + ("test_output".to_string(), path_dst), + ]), + file_names: Vec::from(["test_tome/test_file.txt".to_string()]), + }); + + assert!(broker.collect_errors().is_empty()); // No errors even though the remote asset is inaccessible - // Read let mut contents = String::new(); tmp_file_dst.read_to_string(&mut contents)?; - // Compare assert!(contents.contains("hello from an embedded shell script")); Ok(()) diff --git a/implants/lib/eldritch/src/runtime/exec.rs b/implants/lib/eldritch/src/runtime/exec.rs index 5fbebedb0..eec0b9113 100644 --- a/implants/lib/eldritch/src/runtime/exec.rs +++ b/implants/lib/eldritch/src/runtime/exec.rs @@ -196,6 +196,10 @@ impl Runtime { input_params.insert_hashed(hashed_key, new_value); } module.set("input_params", input_params.alloc_value(module.heap())); + module.set( + "remote_assets", + tome.file_names.clone().alloc_value(module.heap()), + ); Ok(module) }