diff --git a/implants/lib/eldritch/src/assets/list_impl.rs b/implants/lib/eldritch/src/assets/list_impl.rs index 9354fae83..fca856f80 100644 --- a/implants/lib/eldritch/src/assets/list_impl.rs +++ b/implants/lib/eldritch/src/assets/list_impl.rs @@ -1,9 +1,29 @@ -use anyhow::Result; +use anyhow::{Context, Result}; +use starlark::{eval::Evaluator, values::list::ListRef}; -pub fn list() -> Result> { +pub fn list(starlark_eval: &Evaluator<'_, '_>) -> Result> { let mut res: Vec = Vec::new(); - for file_path in super::Asset::iter() { - res.push(file_path.to_string()); + 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")?; + for asset_path in tmp_list.iter() { + let mut asset_path_string = asset_path.to_str(); + if let Some(local_asset_path_string) = asset_path_string.strip_prefix('"') { + asset_path_string = local_asset_path_string.to_string(); + } + if let Some(local_asset_path_string) = asset_path_string.strip_suffix('"') { + asset_path_string = local_asset_path_string.to_string(); + } + res.push(asset_path_string) + } + if res.len() > 0 { + return Ok(res); + } + } + + for asset_path in super::Asset::iter() { + res.push(asset_path.to_string()); } Ok(res) @@ -11,24 +31,61 @@ pub fn list() -> Result> { #[cfg(test)] mod tests { - use super::*; - - #[test] - fn test_assets_list() -> anyhow::Result<()> { - let res_all_embedded_files = list()?; - - assert_eq!( - res_all_embedded_files, - [ - "exec_script/hello_world.bat", - "exec_script/hello_world.sh", - "exec_script/main.eldritch", - "exec_script/metadata.yml", - "print/main.eldritch", - "print/metadata.yml", - ] - ); - - Ok(()) + use std::collections::HashMap; + + use crate::runtime::Message; + use pb::eldritch::Tome; + + macro_rules! test_cases { + ($($name:ident: $value:expr,)*) => { + $( + #[tokio::test] + async fn $name() { + let tc: TestCase = $value; + + // Run Eldritch (until finished) + let mut runtime = crate::start(tc.id, tc.tome).await; + runtime.finish().await; + + // Read Messages + let mut found = false; + for msg in runtime.messages() { + if let Message::ReportText(m) = msg { + assert_eq!(tc.id, m.id); + assert_eq!(tc.want_text, m.text); + found = true; + } + } + assert!(found); + } + )* + } + } + + struct TestCase { + pub id: i64, + pub tome: Tome, + pub want_text: String, + } + + test_cases! { + test_asset_list_remote: TestCase{ + id: 123, + tome: Tome{ + eldritch: String::from(r#"print(assets.list())"#), + parameters: HashMap::new(), + file_names: Vec::new(), + }, + want_text: String::from("[\"exec_script/hello_world.bat\", \"exec_script/hello_world.sh\", \"exec_script/main.eldritch\", \"exec_script/metadata.yml\", \"print/main.eldritch\", \"print/metadata.yml\"]\n"), + }, + test_asset_list_local: TestCase{ + id: 123, + tome: Tome{ + eldritch: String::from(r#"print(assets.list())"#), + parameters: HashMap::new(), + file_names: Vec::from(["remote_asset/just_a_remote_asset.txt".to_string()]), + }, + want_text: String::from("[\"remote_asset/just_a_remote_asset.txt\"]\n"), + }, } } diff --git a/implants/lib/eldritch/src/assets/mod.rs b/implants/lib/eldritch/src/assets/mod.rs index 985bca52e..c733b8b75 100644 --- a/implants/lib/eldritch/src/assets/mod.rs +++ b/implants/lib/eldritch/src/assets/mod.rs @@ -44,17 +44,17 @@ fn methods(builder: &mut MethodsBuilder) { } #[allow(unused_variables)] - fn list(this: &AssetsLibrary) -> anyhow::Result> { - list_impl::list() + fn list(this: &AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>) -> anyhow::Result> { + list_impl::list(starlark_eval) } #[allow(unused_variables)] - fn read_binary(this: &AssetsLibrary, src: String) -> anyhow::Result> { - read_binary_impl::read_binary(src) + fn read_binary(this: &AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>, src: String) -> anyhow::Result> { + read_binary_impl::read_binary(starlark_eval, src) } #[allow(unused_variables)] - fn read(this: &AssetsLibrary, src: String) -> anyhow::Result { - read_impl::read(src) + fn read(this: &AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>, src: String) -> anyhow::Result { + read_impl::read(starlark_eval, src) } } diff --git a/implants/lib/eldritch/src/assets/read_binary_impl.rs b/implants/lib/eldritch/src/assets/read_binary_impl.rs index 137b3a5fb..0ed439d5e 100644 --- a/implants/lib/eldritch/src/assets/read_binary_impl.rs +++ b/implants/lib/eldritch/src/assets/read_binary_impl.rs @@ -1,6 +1,24 @@ -use anyhow::Result; +use std::sync::mpsc::{channel, Receiver}; -pub fn read_binary(src: String) -> Result> { +use anyhow::{Context, Result}; +use pb::c2::FetchAssetResponse; +use starlark::{eval::Evaluator, values::list::ListRef}; + +use crate::runtime::{messages::FetchAssetMessage, Environment}; + +fn read_binary_remote(rx: Receiver) -> Result> { + let mut res: Vec = vec![]; + + // Listen for more chunks and write them + for resp in rx { + let mut new_chunk = resp.chunk.iter().map(|x| *x as u32).collect::>(); + res.append(&mut new_chunk); + } + + Ok(res) +} + +fn read_binary_embedded(src: String) -> Result> { let src_file_bytes = 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.")), @@ -9,20 +27,40 @@ pub fn read_binary(src: String) -> Result> { .iter() .map(|x| *x as u32) .collect::>(); - // let mut result = Vec::new(); - // for byt: Vece in src_file_bytes.iter() { - // result.push(*byte as u32); - // } Ok(result) } +pub fn read_binary(starlark_eval: &Evaluator<'_, '_>, src: 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 env = Environment::from_extra(starlark_eval.extra)?; + let (tx, rx) = channel(); + env.send(FetchAssetMessage { name: src, tx })?; + + return read_binary_remote(rx); + } + } + read_binary_embedded(src) +} + #[cfg(test)] mod tests { + use crate::runtime::Message; + use std::collections::HashMap; + + use crate::runtime::messages::FetchAssetMessage; + use pb::{c2::FetchAssetResponse, eldritch::Tome}; + use super::*; #[test] fn test_assets_read_binary() -> anyhow::Result<()> { - let res = read_binary("print/main.eldritch".to_string())?; + let res = read_binary_embedded("print/main.eldritch".to_string())?; #[cfg(not(windows))] assert_eq!( res, @@ -41,4 +79,91 @@ mod tests { ); Ok(()) } + + pub fn init_logging() { + pretty_env_logger::formatted_timed_builder() + .filter_level(log::LevelFilter::Info) + .parse_env("IMIX_LOG") + .init(); + } + + #[tokio::test] + async fn test_asset_read_binary_remote() -> anyhow::Result<()> { + init_logging(); + // Create files + let tc = Tome { + eldritch: r#"print(assets.read_binary("remote_asset/just_a_remote_asset.txt"))"# + .to_owned(), + parameters: HashMap::new(), + file_names: Vec::from(["remote_asset/just_a_remote_asset.txt".to_string()]), + }; + + // Run Eldritch (in it's own thread) + let mut runtime = crate::start(123, tc.clone()).await; + + // We now mock the agent, looping until eldritch requests a file + // We omit the sleep performed by the agent, just to save test time + loop { + // The runtime only returns the data that is currently available + // So this may return an empty vec if our eldritch tokio task has not yet been scheduled + let messages = runtime.collect(); + let mut fetch_asset_msgs: Vec<&FetchAssetMessage> = messages + .iter() + .filter_map(|m| match m { + Message::FetchAsset(msg) => Some(msg), + _ => None, + }) + .collect(); + + // If no asset request is yet available, just continue looping + if fetch_asset_msgs.is_empty() { + continue; + } + + // Ensure the right asset was requested + assert!(fetch_asset_msgs.len() == 1); + let msg = fetch_asset_msgs.pop().expect("no asset request received!"); + assert!(msg.name == "remote_asset/just_a_remote_asset.txt"); + + // Now, we provide the file to eldritch (as a series of chunks) + msg.tx + .send(FetchAssetResponse { + chunk: "chunk1\n".as_bytes().to_vec(), + }) + .expect("failed to send file chunk to eldritch"); + msg.tx + .send(FetchAssetResponse { + chunk: "chunk2\n".as_bytes().to_vec(), + }) + .expect("failed to send file chunk to eldritch"); + + // We've finished providing the file, so we stop looping + // This will drop `req`, which consequently drops the underlying `Sender` for the file channel + // This will cause the next `recv()` to error with "channel is empty and sending half is closed" + // which is what tells eldritch that there are no more file chunks to wait for + break; + } + + // Now that we've finished writing data, we wait for eldritch to finish + runtime.finish().await; + + let mut found = false; + for msg in runtime.messages() { + if let Message::ReportText(m) = msg { + log::debug!("{}", m.text); + assert_eq!(123, m.id); + assert_eq!( + "[99, 104, 117, 110, 107, 49, 10, 99, 104, 117, 110, 107, 50, 10]\n" + .to_string(), + m.text + ); + found = true; + } + } + assert!(found); + + // Lastly, assert the file was written correctly + + Ok(()) + } } diff --git a/implants/lib/eldritch/src/assets/read_impl.rs b/implants/lib/eldritch/src/assets/read_impl.rs index 8823cf57e..e9f6e897e 100644 --- a/implants/lib/eldritch/src/assets/read_impl.rs +++ b/implants/lib/eldritch/src/assets/read_impl.rs @@ -1,6 +1,12 @@ -use anyhow::Result; +use std::sync::mpsc::{channel, Receiver}; -pub fn read(src: String) -> Result { +use anyhow::{Context, Result}; +use pb::c2::FetchAssetResponse; +use starlark::{eval::Evaluator, values::list::ListRef}; + +use crate::runtime::{messages::FetchAssetMessage, Environment}; + +fn read_local(src: String) -> Result { let src_file_bytes = 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.")), @@ -12,14 +18,165 @@ pub fn read(src: String) -> Result { Ok(result) } +fn read_remote(rx: Receiver) -> Result { + // Wait for our first chunk + let resp = rx.recv()?; + + let mut res = String::new(); + + res.push_str(&String::from_utf8_lossy(resp.chunk.as_slice())); + + // Listen for more chunks and write them + for resp in rx { + res.push_str(&String::from_utf8_lossy(resp.chunk.as_slice())); + } + + Ok(res) +} + +pub fn read(starlark_eval: &Evaluator<'_, '_>, src: 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 env = Environment::from_extra(starlark_eval.extra)?; + let (tx, rx) = channel(); + env.send(FetchAssetMessage { name: src, tx })?; + + return read_remote(rx); + } + } + read_local(src) +} + #[cfg(test)] mod tests { - use super::*; + use std::collections::HashMap; + + use crate::runtime::{messages::FetchAssetMessage, Message}; + use pb::{c2::FetchAssetResponse, eldritch::Tome}; + + macro_rules! test_cases { + ($($name:ident: $value:expr,)*) => { + $( + #[tokio::test] + async fn $name() { + let tc: TestCase = $value; + + // Run Eldritch (until finished) + let mut runtime = crate::start(tc.id, tc.tome).await; + + runtime.finish().await; + + // Read Messages + let mut found = false; + for msg in runtime.messages() { + if let Message::ReportText(m) = msg { + println!("{}", m.text); + assert_eq!(tc.id, m.id); + assert_eq!(tc.want_text, m.text); + found = true; + } + } + assert!(found); + } + )* + } + } + + struct TestCase { + pub id: i64, + pub tome: Tome, + pub want_text: String, + } + + test_cases! { + test_asset_read_local: TestCase{ + id: 123, + tome: Tome{ + eldritch: String::from(r#"print(assets.read("print/main.eldritch").strip())"#), + parameters: HashMap::new(), + file_names: Vec::new(), + }, + want_text: String::from("print(\"This script just prints\")\n"), + }, + + } + + #[tokio::test] + async fn test_asset_read_remote() -> anyhow::Result<()> { + // Create files + let tc = Tome { + eldritch: r#"print(assets.read("remote_asset/just_a_remote_asset.txt").strip())"# + .to_owned(), + parameters: HashMap::new(), + file_names: Vec::from(["remote_asset/just_a_remote_asset.txt".to_string()]), + }; + + // Run Eldritch (in it's own thread) + let mut runtime = crate::start(123, tc.clone()).await; + + // We now mock the agent, looping until eldritch requests a file + // We omit the sleep performed by the agent, just to save test time + loop { + // The runtime only returns the data that is currently available + // So this may return an empty vec if our eldritch tokio task has not yet been scheduled + let messages = runtime.collect(); + let mut fetch_asset_msgs: Vec<&FetchAssetMessage> = messages + .iter() + .filter_map(|m| match m { + Message::FetchAsset(msg) => Some(msg), + _ => None, + }) + .collect(); + + // If no asset request is yet available, just continue looping + if fetch_asset_msgs.is_empty() { + continue; + } + + // Ensure the right asset was requested + assert!(fetch_asset_msgs.len() == 1); + let msg = fetch_asset_msgs.pop().expect("no asset request received!"); + assert!(msg.name == "remote_asset/just_a_remote_asset.txt"); + + // Now, we provide the file to eldritch (as a series of chunks) + msg.tx + .send(FetchAssetResponse { + chunk: "chunk1\n".as_bytes().to_vec(), + }) + .expect("failed to send file chunk to eldritch"); + msg.tx + .send(FetchAssetResponse { + chunk: "chunk2\n".as_bytes().to_vec(), + }) + .expect("failed to send file chunk to eldritch"); + + // We've finished providing the file, so we stop looping + // This will drop `req`, which consequently drops the underlying `Sender` for the file channel + // This will cause the next `recv()` to error with "channel is empty and sending half is closed" + // which is what tells eldritch that there are no more file chunks to wait for + break; + } + + // Now that we've finished writing data, we wait for eldritch to finish + runtime.finish().await; + + let mut found = false; + for msg in runtime.messages() { + if let Message::ReportText(m) = msg { + assert_eq!(123, m.id); + assert_eq!("chunk1\nchunk2\n", m.text); + found = true; + } + } + assert!(found); + + // Lastly, assert the file was written correctly - #[test] - fn test_assets_read() -> anyhow::Result<()> { - let res = read("print/main.eldritch".to_string())?; - assert_eq!(res.trim(), r#"print("This script just prints")"#); Ok(()) } } diff --git a/implants/lib/eldritch/src/file/replace_all_impl.rs b/implants/lib/eldritch/src/file/replace_all_impl.rs index b156a25f9..f4ade2641 100644 --- a/implants/lib/eldritch/src/file/replace_all_impl.rs +++ b/implants/lib/eldritch/src/file/replace_all_impl.rs @@ -1,17 +1,20 @@ use anyhow::Result; -use regex::{NoExpand, Regex}; -use std::fs::{read_to_string, write}; +use regex::bytes::{NoExpand, Regex}; +use std::fs::{self, write}; pub fn replace_all(path: String, pattern: String, value: String) -> Result<()> { - let file_contents = read_to_string(path.clone())?; + let data = fs::read(path.clone())?; + let file_contents_bytes = data.as_slice(); let re = Regex::new(&pattern)?; - let result = re.replace_all(&file_contents, NoExpand(&value)); - write(path, String::from(result))?; + let result = re.replace_all(file_contents_bytes, NoExpand(value.as_bytes())); + write(path, result)?; Ok(()) } #[cfg(test)] mod tests { + use std::fs::read_to_string; + use super::*; use tempfile::NamedTempFile; @@ -55,6 +58,31 @@ mod tests { assert_eq!(res, "Not Match\n".repeat(3)); Ok(()) } + #[test] + fn test_replace_all_regex_binary() -> anyhow::Result<()> { + let tmp_file_new = NamedTempFile::new()?; + let path_new = String::from(tmp_file_new.path().to_str().unwrap()).clone(); + let _ = write( + path_new.clone(), + b"\x90\x90\x90\x90\x90\x90\x90\x90\x90AAAAAAAA\x91\x91\x91\x91\x91\x91\x91\x91\x91AAAAAAAA", + ); + + // Run our code + replace_all( + path_new.clone(), + String::from("AAAAAAAA"), + String::from("BBBBBBBB"), + )?; + + let data = fs::read(path_new.clone())?; + let file_contents_bytes = data.as_slice(); + assert_eq!( + file_contents_bytes, + b"\x90\x90\x90\x90\x90\x90\x90\x90\x90BBBBBBBB\x91\x91\x91\x91\x91\x91\x91\x91\x91BBBBBBBB" + ); + Ok(()) + } + #[test] fn test_replace_all_regex_complex() -> anyhow::Result<()> { let tmp_file_new = NamedTempFile::new()?; diff --git a/implants/lib/eldritch/src/file/replace_impl.rs b/implants/lib/eldritch/src/file/replace_impl.rs index 267a00a23..69b0f5e8f 100644 --- a/implants/lib/eldritch/src/file/replace_impl.rs +++ b/implants/lib/eldritch/src/file/replace_impl.rs @@ -1,17 +1,20 @@ use anyhow::Result; -use regex::{NoExpand, Regex}; -use std::fs::{read_to_string, write}; +use regex::bytes::{NoExpand, Regex}; +use std::fs::{self, write}; pub fn replace(path: String, pattern: String, value: String) -> Result<()> { - let file_contents = read_to_string(path.clone())?; + let data = fs::read(path.clone())?; + let file_contents_bytes = data.as_slice(); let re = Regex::new(&pattern)?; - let result = re.replace(&file_contents, NoExpand(&value)); - write(path, String::from(result))?; + let result = re.replace(file_contents_bytes, NoExpand(value.as_bytes())); + write(path, result)?; Ok(()) } #[cfg(test)] mod tests { + use std::fs::read_to_string; + use super::*; use tempfile::NamedTempFile; @@ -59,6 +62,30 @@ mod tests { Ok(()) } #[test] + fn test_replace_regex_binary() -> anyhow::Result<()> { + let tmp_file_new = NamedTempFile::new()?; + let path_new = String::from(tmp_file_new.path().to_str().unwrap()).clone(); + let _ = write( + path_new.clone(), + b"\x90\x90\x90\x90\x90\x90\x90\x90\x90AAAAAAAA\x91\x91\x91\x91\x91\x91\x91\x91\x91", + ); + + // Run our code + replace( + path_new.clone(), + String::from("AAAAAAAA"), + String::from("BBBBBBBB"), + )?; + + let data = fs::read(path_new.clone())?; + let file_contents_bytes = data.as_slice(); + assert_eq!( + file_contents_bytes, + b"\x90\x90\x90\x90\x90\x90\x90\x90\x90BBBBBBBB\x91\x91\x91\x91\x91\x91\x91\x91\x91" + ); + Ok(()) + } + #[test] fn test_replace_regex_complex() -> anyhow::Result<()> { let tmp_file_new = NamedTempFile::new()?; let path_new = String::from(tmp_file_new.path().to_str().unwrap()).clone(); diff --git a/implants/lib/eldritch/src/file/template_impl.rs b/implants/lib/eldritch/src/file/template_impl.rs index 1b7166c5b..da2f6d49b 100644 --- a/implants/lib/eldritch/src/file/template_impl.rs +++ b/implants/lib/eldritch/src/file/template_impl.rs @@ -21,8 +21,10 @@ pub fn template( autoescape: bool, ) -> Result<()> { let context = build_context(args)?; - let template_content = fs::read_to_string(template_path)?; - let res_content = Tera::one_off(template_content.as_str(), &context, autoescape)?; + let data = fs::read(template_path.clone())?; + let template_content = String::from_utf8_lossy(&data); + + let res_content = Tera::one_off(template_content.to_string().as_str(), &context, autoescape)?; fs::write(dst_path, res_content)?; Ok(()) }