Binary read support

implants/lib/eldritch/src/assets/list_impl.rs

@@ -1,34 +1,91 @@
-use anyhow::Result;
+use anyhow::{Context, Result};
+use starlark::{eval::Evaluator, values::list::ListRef};
 
-pub fn list() -> Result<Vec<String>> {
+pub fn list(starlark_eval: &Evaluator<'_, '_>) -> Result<Vec<String>> {
     let mut res: Vec<String> = 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)
 }
 
 #[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"),
+            },
     }
 }

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<Vec<String>> {
-        list_impl::list()
+    fn list(this: &AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>) -> anyhow::Result<Vec<String>> {
+        list_impl::list(starlark_eval)
     }
 
     #[allow(unused_variables)]
-    fn read_binary(this: &AssetsLibrary, src: String) -> anyhow::Result<Vec<u32>> {
-        read_binary_impl::read_binary(src)
+    fn read_binary(this: &AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>, src: String) -> anyhow::Result<Vec<u32>> {
+        read_binary_impl::read_binary(starlark_eval, src)
     }
 
     #[allow(unused_variables)]
-    fn read(this: &AssetsLibrary, src: String) -> anyhow::Result<String> {
-        read_impl::read(src)
+    fn read(this: &AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>, src: String) -> anyhow::Result<String> {
+        read_impl::read(starlark_eval, src)
     }
 }

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<Vec<u32>> {
+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<FetchAssetResponse>) -> Result<Vec<u32>> {
+    let mut res: Vec<u32> = 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::<Vec<u32>>();
+        res.append(&mut new_chunk);
+    }
+
+    Ok(res)
+}
+
+fn read_binary_embedded(src: String) -> Result<Vec<u32>> {
     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<Vec<u32>> {
         .iter()
         .map(|x| *x as u32)
         .collect::<Vec<u32>>();
-    // let mut result = Vec::new();
-    // for byt: Vec<u32>e in src_file_bytes.iter() {
-    //     result.push(*byte as u32);
-    // }
     Ok(result)
 }
 
+pub fn read_binary(starlark_eval: &Evaluator<'_, '_>, src: String) -> Result<Vec<u32>> {
+    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(())
+    }
 }